MuyueWorkspace/internal/memory/recall.go

package memory

import (
	"database/sql"
	"strings"
	"time"
)

type SearchResult struct {
	Memory
	Score float64 `json:"score"`
}

func (s *Store) Search(query string, limit int) ([]SearchResult, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if limit <= 0 {
		limit = 10
	}
	if limit > 50 {
		limit = 50
	}

	normalizedQuery := normalizeQuery(query)

	rows, err := s.db.Query(`
		SELECT m.id, m.type, m.key, m.content, m.tags, m.source, m.confidence,
		       m.access_count, m.created_at, m.updated_at,
		       bm25(memories_fts) as score
		FROM memories_fts f
		JOIN memories m ON m.rowid = f.rowid
		WHERE memories_fts MATCH ?
		ORDER BY score
		LIMIT ?
	`, normalizedQuery, limit)
	if err != nil {
		return fallbackSearch(s.db, query, limit)
	}
	defer rows.Close()

	return scanSearchResults(rows)
}

func (s *Store) Recall(query string, limit int) ([]Memory, error) {
	results, err := s.Search(query, limit)
	if err != nil {
		return nil, err
	}

	memories := make([]Memory, len(results))
	for i, r := range results {
		memories[i] = r.Memory
	}
	return memories, nil
}

func (s *Store) RecallByType(memType MemoryType, limit int) ([]Memory, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if limit <= 0 {
		limit = 20
	}

	rows, err := s.db.Query(`
		SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
		FROM memories WHERE type = ?
		ORDER BY access_count DESC, updated_at DESC
		LIMIT ?
	`, string(memType), limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanMemories(rows)
}

func (s *Store) RecallRecent(since time.Time, limit int) ([]Memory, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if limit <= 0 {
		limit = 20
	}

	rows, err := s.db.Query(`
		SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
		FROM memories WHERE updated_at >= ?
		ORDER BY updated_at DESC
		LIMIT ?
	`, since, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanMemories(rows)
}

func (s *Store) RecallPreferences() ([]Memory, error) {
	return s.RecallByType(TypePreference, 50)
}

func (s *Store) RecallFacts() ([]Memory, error) {
	return s.RecallByType(TypeFact, 50)
}

func (s *Store) StorePreference(key, content string) error {
	return s.Store(&Memory{
		Type:       TypePreference,
		Key:        key,
		Content:    content,
		Source:     "user",
		Confidence: 0.9,
	})
}

func (s *Store) StoreContext(key, content string) error {
	return s.Store(&Memory{
		Type:       TypeContext,
		Key:        key,
		Content:    content,
		Source:     "conversation",
		Confidence: 0.7,
	})
}

func (s *Store) StoreSummary(sessionID, summary string) error {
	return s.Store(&Memory{
		Type:       TypeSummary,
		Key:        "session:" + sessionID,
		Content:    summary,
		Source:     "auto",
		Confidence: 0.8,
	})
}

func (s *Store) StoreFact(key, content string) error {
	return s.Store(&Memory{
		Type:       TypeFact,
		Key:        key,
		Content:    content,
		Source:     "auto",
		Confidence: 0.85,
	})
}

func normalizeQuery(query string) string {
	words := strings.Fields(strings.ToLower(query))
	var escaped []string
	for _, w := range words {
		if len(w) > 0 {
			escaped = append(escaped, w+"*")
		}
	}
	return strings.Join(escaped, " OR ")
}

func fallbackSearch(db *sql.DB, query string, limit int) ([]SearchResult, error) {
	likePattern := "%" + strings.ToLower(query) + "%"

	rows, err := db.Query(`
		SELECT id, type, key, content, tags, source, confidence, access_count, created_at, updated_at
		FROM memories
		WHERE LOWER(key) LIKE ? OR LOWER(content) LIKE ? OR LOWER(tags) LIKE ?
		ORDER BY updated_at DESC
		LIMIT ?
	`, likePattern, likePattern, likePattern, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var results []SearchResult
	for rows.Next() {
		var m Memory
		err := rows.Scan(&m.ID, &m.Type, &m.Key, &m.Content, &m.Tags, &m.Source, &m.Confidence, &m.AccessCount, &m.CreatedAt, &m.UpdatedAt)
		if err != nil {
			return results, err
		}
		score := computeFallbackScore(m, query)
		results = append(results, SearchResult{Memory: m, Score: score})
	}
	return results, nil
}

func computeFallbackScore(m Memory, query string) float64 {
	score := m.Confidence * 0.5
	lower := strings.ToLower(query)
	if strings.Contains(strings.ToLower(m.Key), lower) {
		score += 0.3
	}
	if strings.Contains(strings.ToLower(m.Content), lower) {
		score += 0.2
	}
	score += float64(m.AccessCount) * 0.01
	return score
}

func scanSearchResults(rows *sql.Rows) ([]SearchResult, error) {
	var results []SearchResult
	for rows.Next() {
		var m Memory
		var score float64
		err := rows.Scan(&m.ID, &m.Type, &m.Key, &m.Content, &m.Tags, &m.Source,
			&m.Confidence, &m.AccessCount, &m.CreatedAt, &m.UpdatedAt, &score)
		if err != nil {
			return results, err
		}
		results = append(results, SearchResult{Memory: m, Score: score})
	}
	return results, nil
}