Cyrene_v2/backend/services/memory/service.go

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
}

// ApplyDecayAll applies decay for every user that has memories.
func (s *Service) ApplyDecayAll(ctx context.Context) (int, error) {
	users, err := s.repo.ListUserIDs(ctx)
	if err != nil {
		return 0, err
	}
	total := 0
	for _, uid := range users {
		n, err := s.ApplyDecay(ctx, uid)
		if err != nil {
			return total, err
		}
		total += n
	}
	return total, 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)
}