Files
Cyrene_v2/backend/services/memory/service.go
T
Aska 8495a4b055 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
2026-08-20 20:57:13 +08:00

311 lines
8.2 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package memory
import (
"context"
"fmt"
"math"
"strings"
"time"
)
// Service handles business logic for the memory system.
type Service struct {
repo *Repository
embedder *Embedder
}
// NewService creates a new memory service.
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 }
keywords := extractKeywords(content)
// Check for duplicates via n-gram keyword overlap
nGrams := extractNGrams(content, 2)
if len(nGrams) > 0 {
// Search for memories with overlapping keywords
existing, _ := s.repo.SearchByNGrams(ctx, userID, nGrams[:min(8, len(nGrams))], 5)
for _, e := range existing {
overlap := keywordOverlap(keywords, e.Keywords)
sim := jaccardSimilarity(content, e.Content)
// Merge if keyword overlap > 28% OR Jaccard > 0.45
// (raised from 25%/0.4 to reduce false positives)
if overlap > 0.28 || sim > 0.45 {
e.Priority = math.Max(e.Priority, priority)
e.Content = mergeContent(e.Content, content)
e.Summary = truncate(e.Content, 200)
e.Keywords = mergeKeywords(e.Keywords, keywords)
e.AccessCnt++
e.LastAcc = time.Now()
if err := s.repo.Update(ctx, &e); err != nil {
return nil, fmt.Errorf("merge memory: %w", err)
}
return &e, nil
}
}
}
m := &Memory{
UserID: userID,
SourceID: sourceID,
Content: content,
Summary: truncate(content, 200),
Keywords: keywords,
Category: category,
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)
}
return m, nil
}
// 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 }
seen := make(map[string]bool)
var memories []Memory
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 {
kwHits, err = s.repo.SearchByNGrams(ctx, q.UserID, nGrams[:min(3, len(nGrams))], q.Limit)
} else {
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)
}
if q.SourceID != "" {
filtered := make([]Memory, 0)
for _, m := range memories {
if m.SourceID == q.SourceID { filtered = append(filtered, m) }
}
return filtered, nil
}
return memories, nil
}
// List returns recent memories.
func (s *Service) List(ctx context.Context, userID string, limit int) ([]Memory, error) {
return s.repo.ListByUser(ctx, userID, limit)
}
// GetByID returns a single memory.
func (s *Service) GetByID(ctx context.Context, id string) (*Memory, error) {
return s.repo.GetByID(ctx, id)
}
// Delete removes a memory.
func (s *Service) Delete(ctx context.Context, id string) error {
return s.repo.Delete(ctx, id)
}
// ApplyDecay reduces decay factor based on time elapsed.
func (s *Service) ApplyDecay(ctx context.Context, userID string) (int, error) {
memories, err := s.repo.ListByUser(ctx, userID, 100)
if err != nil { return 0, err }
updated := 0
for _, m := range memories {
elapsed := time.Since(m.LastAcc)
newDecay := math.Exp(-0.01 * elapsed.Hours() / 24)
if newDecay != m.Decay {
m.Decay = newDecay
if err := s.repo.Update(ctx, &m); err == nil { updated++ }
}
}
s.repo.PurgeDecayed(ctx, userID)
return updated, nil
}
// --- Keyword extraction (sliding window n-grams) ---
// extractNGrams extracts all n-gram substrings from text.
func extractNGrams(text string, n int) []string {
runes := []rune(text)
if len(runes) < n { return []string{text} }
seen := make(map[string]bool)
var result []string
for i := 0; i <= len(runes)-n; i++ {
gram := string(runes[i : i+n])
if !seen[gram] {
seen[gram] = true
result = append(result, gram)
}
}
return result
}
// extractKeywords extracts meaningful keywords using 2-3 char sliding window.
func extractKeywords(content string) []string {
seen := make(map[string]bool)
var keywords []string
// Priority 1: Full segments (comma/sentence split) — most meaningful
segments := strings.FieldsFunc(content, func(r rune) bool {
return r == ',' || r == '。' || r == '、' || r == ' ' || r == '\n' || r == '!' || r == '?' || r == ','
})
for _, seg := range segments {
seg = strings.TrimSpace(seg)
if len([]rune(seg)) >= 2 && !seen[seg] {
keywords = append(keywords, seg)
seen[seg] = true
if len(keywords) >= 10 { return keywords }
}
}
// Priority 2: 3-grams (more discriminative than 2-grams)
runes := []rune(content)
for i := 0; i <= len(runes)-3; i++ {
gram := string(runes[i : i+3])
if isAllPunct(gram) { continue }
if !seen[gram] {
seen[gram] = true
keywords = append(keywords, gram)
if len(keywords) >= 20 { return keywords }
}
}
// Priority 3: 2-grams (fallback)
for i := 0; i <= len(runes)-2; i++ {
gram := string(runes[i : i+2])
if isAllPunct(gram) { continue }
if !seen[gram] {
seen[gram] = true
keywords = append(keywords, gram)
if len(keywords) >= 25 { return keywords }
}
}
// Priority 4: Full text (last resort)
if len(keywords) == 0 && len(runes) > 0 {
keywords = append(keywords, string(runes))
}
return keywords
}
// --- Jaccard similarity (n-gram based) ---
func jaccardSimilarity(a, b string) float64 {
setA := ngramSet(a, 2)
setB := ngramSet(b, 2)
if len(setA) == 0 && len(setB) == 0 { return 1.0 }
if len(setA) == 0 || len(setB) == 0 { return 0.0 }
intersection := 0
for token := range setA {
if setB[token] { intersection++ }
}
union := len(setA) + len(setB) - intersection
return float64(intersection) / float64(union)
}
func ngramSet(s string, n int) map[string]bool {
tokens := make(map[string]bool)
runes := []rune(s)
for i := 0; i <= len(runes)-n; i++ {
tokens[string(runes[i:i+n])] = true
}
// Also include single characters for short strings
if len(runes) < n*2 {
for _, r := range runes {
tokens[string(r)] = true
}
}
return tokens
}
// --- Helpers ---
func isAllPunct(s string) bool {
for _, r := range s {
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= 0x4e00 && r <= 0x9fff) || (r >= '0' && r <= '9') {
return false
}
}
return true
}
func mergeContent(old, new string) string {
if len([]rune(new)) > len([]rune(old)) { return new }
return old
}
func mergeKeywords(old, new []string) []string {
seen := make(map[string]bool)
var result []string
for _, k := range old { seen[k] = true; result = append(result, k) }
for _, k := range new {
if !seen[k] {
seen[k] = true
result = append(result, k)
}
}
return result
}
// keywordOverlap calculates what fraction of keywords from `a` are present in `b`.
func keywordOverlap(a, b []string) float64 {
if len(a) == 0 { return 0 }
seen := make(map[string]bool, len(b))
for _, k := range b { seen[k] = true }
matches := 0
for _, k := range a { if seen[k] { matches++ } }
return float64(matches) / float64(len(a))
}
func truncate(s string, maxLen int) string {
runes := []rune(s)
if len(runes) <= maxLen { return s }
return string(runes[:maxLen]) + "..."
}
// ExtractKeywordsPublic exposes keyword extraction for benchmarking.
func ExtractKeywordsPublic(content string) []string {
return extractKeywords(content)
}