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
@@ -133,6 +133,16 @@ func main() {
|
|||||||
if database != nil {
|
if database != nil {
|
||||||
memRepo := memory.NewRepository(database)
|
memRepo := memory.NewRepository(database)
|
||||||
memSvc = memory.NewService(memRepo)
|
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")
|
logger.Infof("[memory] connected")
|
||||||
} else {
|
} else {
|
||||||
logger.Infof("[memory] disabled (no database)")
|
logger.Infof("[memory] disabled (no database)")
|
||||||
|
|||||||
@@ -19,6 +19,10 @@ type Memory struct {
|
|||||||
LastAcc time.Time `json:"last_accessed_at"`
|
LastAcc time.Time `json:"last_accessed_at"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
CreatedAt time.Time `json:"created_at"`
|
||||||
UpdatedAt time.Time `json:"updated_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.
|
// SearchQuery represents a memory search request.
|
||||||
|
|||||||
@@ -12,7 +12,8 @@ import (
|
|||||||
|
|
||||||
// Repository handles database operations for memories.
|
// Repository handles database operations for memories.
|
||||||
type Repository struct {
|
type Repository struct {
|
||||||
db *sql.DB
|
db *sql.DB
|
||||||
|
vectorEnabled bool // true when the pgvector embedding column exists
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewRepository creates a new memory repository.
|
// NewRepository creates a new memory repository.
|
||||||
@@ -20,6 +21,36 @@ func NewRepository(db *sql.DB) *Repository {
|
|||||||
return &Repository{db: db}
|
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.
|
// Create inserts a new memory and returns its ID.
|
||||||
func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
||||||
query := `
|
query := `
|
||||||
@@ -41,6 +72,24 @@ func (r *Repository) Create(ctx context.Context, m *Memory) (string, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
var id string
|
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,
|
err := r.db.QueryRowContext(ctx, query,
|
||||||
m.UserID, m.SourceID, m.Content, m.Summary,
|
m.UserID, m.SourceID, m.Content, m.Summary,
|
||||||
pq.Array(m.Keywords), m.Category, m.Priority, m.Decay,
|
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)
|
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.
|
// Update updates an existing memory.
|
||||||
func (r *Repository) Update(ctx context.Context, m *Memory) error {
|
func (r *Repository) Update(ctx context.Context, m *Memory) error {
|
||||||
m.UpdatedAt = time.Now()
|
m.UpdatedAt = time.Now()
|
||||||
|
|||||||
@@ -10,7 +10,8 @@ import (
|
|||||||
|
|
||||||
// Service handles business logic for the memory system.
|
// Service handles business logic for the memory system.
|
||||||
type Service struct {
|
type Service struct {
|
||||||
repo *Repository
|
repo *Repository
|
||||||
|
embedder *Embedder
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewService creates a new memory service.
|
// NewService creates a new memory service.
|
||||||
@@ -18,6 +19,9 @@ func NewService(repo *Repository) *Service {
|
|||||||
return &Service{repo: repo}
|
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.
|
// 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) {
|
func (s *Service) Add(ctx context.Context, userID, sourceID, content, category string, priority float64) (*Memory, error) {
|
||||||
if priority <= 0 { priority = 0.5 }
|
if priority <= 0 { priority = 0.5 }
|
||||||
@@ -58,6 +62,11 @@ func (s *Service) Add(ctx context.Context, userID, sourceID, content, category s
|
|||||||
Priority: priority,
|
Priority: priority,
|
||||||
Decay: 1.0,
|
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 {
|
if _, err := s.repo.Create(ctx, m); err != nil {
|
||||||
return nil, fmt.Errorf("add memory: %w", err)
|
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
|
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) {
|
func (s *Service) Search(ctx context.Context, q SearchQuery) ([]Memory, error) {
|
||||||
if q.Limit <= 0 { q.Limit = 10 }
|
if q.Limit <= 0 { q.Limit = 10 }
|
||||||
|
|
||||||
// Use n-gram based search for better Chinese text matching
|
seen := make(map[string]bool)
|
||||||
nGrams := extractNGrams(q.Query, 2)
|
|
||||||
var memories []Memory
|
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 {
|
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 {
|
} else {
|
||||||
memories, err = s.repo.ListByUser(ctx, q.UserID, q.Limit)
|
kwHits, err = s.repo.ListByUser(ctx, q.UserID, q.Limit)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("search memory: %w", err)
|
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 {
|
for _, m := range memories {
|
||||||
go s.repo.RecordAccess(context.Background(), m.ID)
|
go s.repo.RecordAccess(context.Background(), m.ID)
|
||||||
|
|||||||
Reference in new issue
Block a user