Skip to content

Commit 89bf319

Browse files
committed
feat(postgres): add findSimilar method for cosine distance search in PostgresStore
1 parent 9c3682d commit 89bf319

1 file changed

Lines changed: 49 additions & 0 deletions

File tree

pkg/memory/postgres.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,3 +191,52 @@ func (ps *PostgresStore) Store(ctx context.Context, req StoreRequest) (*StoreRes
191191

192192
return result, nil
193193
}
194+
195+
type pgSimilarEntry struct {
196+
id string
197+
text string
198+
distance float64
199+
isDup bool
200+
}
201+
202+
// findSimilar performs a full-scan cosine distance search.
203+
// The comment in the SQLite implementation applies equally here
204+
// for < 10K rows; at larger scale consider pgvector or a separate ANN index.
205+
func (s *PostgresStore) findSimilar(ctx context.Context, embedding []float32) ([]pgSimilarEntry, error) {
206+
rows, err := s.pool.Query(ctx,
207+
`SELECT id, text, embedding FROM memories WHERE embedding IS NOT NULL AND expired = FALSE`,
208+
)
209+
if err != nil {
210+
return nil, err
211+
}
212+
defer rows.Close()
213+
214+
conflictThreshold := s.cfg.ConflictThreshold
215+
if conflictThreshold <= 0 {
216+
conflictThreshold = 0.35
217+
}
218+
219+
var results []pgSimilarEntry
220+
for rows.Next() {
221+
var id, text string
222+
var embBlob []byte
223+
if err := rows.Scan(&id, &text, &embBlob); err != nil {
224+
return nil, err
225+
}
226+
227+
existing := decodeEmbedding(embBlob)
228+
if len(existing) == 0 {
229+
continue
230+
}
231+
232+
dist := distillmath.CosineDistance(embedding, existing)
233+
if dist < s.cfg.DedupThreshold {
234+
return []pgSimilarEntry{{id: id, text: text, distance: dist, isDup: true}}, nil
235+
}
236+
if dist < conflictThreshold {
237+
results = append(results, pgSimilarEntry{id: id, text: text, distance: dist})
238+
}
239+
}
240+
241+
return results, rows.Err()
242+
}

0 commit comments

Comments
 (0)