feat(memory): hybrid vector search via embedder
- repository: optional pgvector embedding column (migration with CREATE EXTENSION guarded by main), SearchByVector cosine search, vector-enabled insert path - service: SetEmbedder — memories are embedded on Add; Search merges vector hits with n-gram keyword matches (vector first, deduped) - embedder was already implemented (OpenAI-compatible API + n-gram FNV fallback) but never instantiated; main now wires it from EMBEDDING_API_URL/EMBEDDING_MODEL config
This commit is contained in:
1 parent
8190007865
commit
8495a4b055
4 files changed
+119
-8
No files matched your search
@@ -12,7 +12,8 @@ import (
|
||||
|
||||
// Repository handles database operations for memories.
|
||||
type Repository struct {
|
||||
db *sql.DB
|
||||
db *sql.DB
|
||||
vectorEnabled bool // true when the pgvector embedding column exists
|
||||
}
|
||||
|
||||
// NewRepository creates a new memory repository.
|
||||
@@ -20,6 +21,36 @@ func NewRepository(db *sql.DB) *Repository {
|
||||
return &Repository{db: db}
|
||||
}
|
||||
|
||||
// EnableVector marks the embedding column as available (after migration).
|
||||
func (r *Repository) EnableVector() { r.vectorEnabled = true }
|
||||
|
||||
// VectorEnabled reports whether vector search is usable.
|
||||
func (r *Repository) VectorEnabled() bool { return r.vectorEnabled }
|
||||
|
||||
// VectorMigrationSQL returns the statements to enable pgvector and add the
|
||||
// embedding column with the given dimension.
|
||||
func VectorMigrationSQL(dims int) string {
|
||||
return fmt.Sprintf(`
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
ALTER TABLE memories ADD COLUMN IF NOT EXISTS embedding vector(%d);
|
||||
CREATE INDEX IF NOT EXISTS idx_memories_embedding ON memories USING hnsw (embedding vector_cosine_ops);
|
||||
`, dims)
|
||||
}
|
||||
|
||||
// vectorString formats an embedding for pgvector input.
|
||||
func vectorString(vec []float64) string {
|
||||
var sb strings.Builder
|
||||
sb.WriteByte('[')
|
||||
for i, v := range vec {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&sb, "%g", v)
|
||||
}
|
||||
sb.WriteByte(']')
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// Create inserts a new memory and returns its ID.
|
||||
func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
||||
query := `
|
||||
@@ -41,6 +72,24 @@ func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
||||
}
|
||||
|
||||
var id string
|
||||
if r.vectorEnabled && len(m.Embedding) > 0 {
|
||||
query = `
|
||||
INSERT INTO memories (user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at, embedding)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13::vector)
|
||||
RETURNING id
|
||||
`
|
||||
err := r.db.QueryRowContext(ctx, query,
|
||||
m.UserID, m.SourceID, m.Content, m.Summary,
|
||||
pq.Array(m.Keywords), m.Category, m.Priority, m.Decay,
|
||||
m.AccessCnt, m.LastAcc, m.CreatedAt, m.UpdatedAt, vectorString(m.Embedding),
|
||||
).Scan(&id)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("insert memory: %w", err)
|
||||
}
|
||||
m.ID = id
|
||||
return id, nil
|
||||
}
|
||||
|
||||
err := r.db.QueryRowContext(ctx, query,
|
||||
m.UserID, m.SourceID, m.Content, m.Summary,
|
||||
pq.Array(m.Keywords), m.Category, m.Priority, m.Decay,
|
||||
@@ -127,6 +176,22 @@ func (r *Repository) SearchByNGrams(ctx context.Context, userID string, nGrams [
|
||||
return r.queryMemories(ctx, query, userID, DecayThreshold, pq.Array(nGrams), limit)
|
||||
}
|
||||
|
||||
// SearchByVector returns the nearest memories by cosine distance.
|
||||
func (r *Repository) SearchByVector(ctx context.Context, userID string, vec []float64, limit int) ([]Memory, error) {
|
||||
if !r.vectorEnabled || len(vec) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
query := `
|
||||
SELECT id, user_id, source_id, content, summary, keywords, category, priority, decay, access_count, last_accessed_at, created_at, updated_at
|
||||
FROM memories WHERE user_id = $1 AND decay > $2 AND embedding IS NOT NULL
|
||||
ORDER BY embedding <=> $3::vector LIMIT $4
|
||||
`
|
||||
return r.queryMemories(ctx, query, userID, DecayThreshold, vectorString(vec), limit)
|
||||
}
|
||||
|
||||
// Update updates an existing memory.
|
||||
func (r *Repository) Update(ctx context.Context, m *Memory) error {
|
||||
m.UpdatedAt = time.Now()
|
||||
|
||||
Reference in new issue
Block a user