From 8495a4b0559513eff2b20a68d4fd3ad788f378c5 Mon Sep 17 00:00:00 2001 From: Aska Date: Thu, 20 Aug 2026 20:57:13 +0800 Subject: [PATCH] feat(memory): hybrid vector search via embedder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- backend/cmd/cyrene/main.go | 10 ++++ backend/services/memory/model.go | 4 ++ backend/services/memory/repository.go | 67 ++++++++++++++++++++++++++- backend/services/memory/service.go | 46 +++++++++++++++--- 4 files changed, 119 insertions(+), 8 deletions(-) diff --git a/backend/cmd/cyrene/main.go b/backend/cmd/cyrene/main.go index fe8c87d..58ec3a5 100644 --- a/backend/cmd/cyrene/main.go +++ b/backend/cmd/cyrene/main.go @@ -133,6 +133,16 @@ func main() { if database != nil { memRepo := memory.NewRepository(database) memSvc = memory.NewService(memRepo) + + // Embedding: real API when configured, n-gram FNV fallback otherwise + emb := memory.NewEmbedder(cfg.EmbeddingAPIURL, cfg.EmbeddingAPIKey, cfg.EmbeddingModel) + if err := db.RunMigration(context.Background(), database, "memories-vector", db.SplitSQL(memory.VectorMigrationSQL(emb.Dims()))); err != nil { + logger.Infof("[memory] pgvector unavailable, keyword-only search: %v", err) + } else { + memRepo.EnableVector() + memSvc.SetEmbedder(emb) + logger.Infof("[memory] hybrid search enabled (vector dims=%d, api=%v)", emb.Dims(), emb.IsAvailable()) + } logger.Infof("[memory] connected") } else { logger.Infof("[memory] disabled (no database)") diff --git a/backend/services/memory/model.go b/backend/services/memory/model.go index 4ffd978..0ea65c5 100644 --- a/backend/services/memory/model.go +++ b/backend/services/memory/model.go @@ -19,6 +19,10 @@ type Memory struct { LastAcc time.Time `json:"last_accessed_at"` CreatedAt time.Time `json:"created_at"` UpdatedAt time.Time `json:"updated_at"` + + // Embedding is the vector representation used for semantic search. + // Not serialized or scanned by default queries; pgvector-only. + Embedding []float64 `json:"-"` } // SearchQuery represents a memory search request. diff --git a/backend/services/memory/repository.go b/backend/services/memory/repository.go index c9960fa..7a41d32 100644 --- a/backend/services/memory/repository.go +++ b/backend/services/memory/repository.go @@ -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() diff --git a/backend/services/memory/service.go b/backend/services/memory/service.go index 149e78b..f704d62 100644 --- a/backend/services/memory/service.go +++ b/backend/services/memory/service.go @@ -10,7 +10,8 @@ import ( // Service handles business logic for the memory system. type Service struct { - repo *Repository + repo *Repository + embedder *Embedder } // NewService creates a new memory service. @@ -18,6 +19,9 @@ func NewService(repo *Repository) *Service { return &Service{repo: repo} } +// SetEmbedder enables embedding-based semantic search. +func (s *Service) SetEmbedder(e *Embedder) { s.embedder = e } + // Add stores a new memory with deduplication based on keyword overlap. func (s *Service) Add(ctx context.Context, userID, sourceID, content, category string, priority float64) (*Memory, error) { if priority <= 0 { priority = 0.5 } @@ -58,6 +62,11 @@ func (s *Service) Add(ctx context.Context, userID, sourceID, content, category s Priority: priority, Decay: 1.0, } + if s.embedder != nil { + if vec, err := s.embedder.Embed(ctx, content); err == nil { + m.Embedding = vec + } + } if _, err := s.repo.Create(ctx, m); err != nil { return nil, fmt.Errorf("add memory: %w", err) @@ -65,23 +74,46 @@ func (s *Service) Add(ctx context.Context, userID, sourceID, content, category s return m, nil } -// Search performs combined keyword search. +// Search performs hybrid search: vector similarity (when an embedder is +// configured) merged with n-gram keyword matches, vector hits first. func (s *Service) Search(ctx context.Context, q SearchQuery) ([]Memory, error) { if q.Limit <= 0 { q.Limit = 10 } - // Use n-gram based search for better Chinese text matching - nGrams := extractNGrams(q.Query, 2) + seen := make(map[string]bool) var memories []Memory - var err error + if s.embedder != nil { + if vec, err := s.embedder.Embed(ctx, q.Query); err == nil { + if vecHits, err := s.repo.SearchByVector(ctx, q.UserID, vec, q.Limit); err == nil { + for _, m := range vecHits { + memories = append(memories, m) + seen[m.ID] = true + } + } + } + } + + // N-gram keyword search for better Chinese text matching + nGrams := extractNGrams(q.Query, 2) + var kwHits []Memory + var err error if len(nGrams) > 0 { - memories, err = s.repo.SearchByNGrams(ctx, q.UserID, nGrams[:min(3, len(nGrams))], q.Limit) + kwHits, err = s.repo.SearchByNGrams(ctx, q.UserID, nGrams[:min(3, len(nGrams))], q.Limit) } else { - memories, err = s.repo.ListByUser(ctx, q.UserID, q.Limit) + kwHits, err = s.repo.ListByUser(ctx, q.UserID, q.Limit) } if err != nil { return nil, fmt.Errorf("search memory: %w", err) } + for _, m := range kwHits { + if !seen[m.ID] { + memories = append(memories, m) + seen[m.ID] = true + } + } + if len(memories) > q.Limit { + memories = memories[:q.Limit] + } for _, m := range memories { go s.repo.RecordAccess(context.Background(), m.ID)