goclaw/internal/store/pg/memory_search.go

package pg

import (
	"context"
	"fmt"
	"strings"

	"github.com/nextlevelbuilder/goclaw/internal/store"
)

// Search performs hybrid search (FTS + vector) over memory_chunks.
// Merges global (user_id IS NULL) + per-user chunks, with user boost.
func (s *PGMemoryStore) Search(ctx context.Context, query string, agentID, userID string, opts store.MemorySearchOptions) ([]store.MemorySearchResult, error) {
	maxResults := opts.MaxResults
	if maxResults <= 0 {
		maxResults = s.cfg.MaxResults
	}

	aid, err := parseUUID(agentID)
	if err != nil {
		return nil, fmt.Errorf("memory search: %w", err)
	}

	// FTS search using tsvector
	ftsResults, err := s.ftsSearch(ctx, query, aid, userID, maxResults*2)
	if err != nil {
		return nil, err
	}

	// Vector search if provider available
	var vecResults []scoredChunk
	if s.provider != nil {
		embeddings, err := s.provider.Embed(ctx, []string{query})
		if err == nil && len(embeddings) > 0 {
			vecResults, err = s.vectorSearch(ctx, embeddings[0], aid, userID, maxResults*2)
			if err != nil {
				vecResults = nil
			}
		}
	}

	// Merge results — use per-query overrides if set, else store defaults
	textW, vecW := s.cfg.TextWeight, s.cfg.VectorWeight
	if opts.TextWeight > 0 {
		textW = opts.TextWeight
	}
	if opts.VectorWeight > 0 {
		vecW = opts.VectorWeight
	}
	if len(ftsResults) == 0 && len(vecResults) > 0 {
		textW, vecW = 0, 1.0
	} else if len(vecResults) == 0 && len(ftsResults) > 0 {
		textW, vecW = 1.0, 0
	}
	merged := hybridMerge(ftsResults, vecResults, textW, vecW, userID)

	// Apply min score filter
	var filtered []store.MemorySearchResult
	for _, m := range merged {
		if opts.MinScore > 0 && m.Score < opts.MinScore {
			continue
		}
		if opts.PathPrefix != "" && !strings.HasPrefix(m.Path, opts.PathPrefix) {
			continue
		}
		filtered = append(filtered, m)
		if len(filtered) >= maxResults {
			break
		}
	}

	return filtered, nil
}

type scoredChunk struct {
	Path      string
	StartLine int
	EndLine   int
	Text      string
	Score     float64
	UserID    *string
}

func (s *PGMemoryStore) ftsSearch(ctx context.Context, query string, agentID any, userID string, limit int) ([]scoredChunk, error) {
	var q string
	var args []any

	if store.IsSharedMemory(ctx) {
		// Shared: no user_id filter — search ALL chunks for agent
		tc, tcArgs, _, err := scopeClause(ctx, 4)
		if err != nil {
			return nil, err
		}
		limitN := 4 + len(tcArgs)
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				ts_rank(tsv, plainto_tsquery('russian', $1)) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND tsv @@ plainto_tsquery('russian', $3)%s
			ORDER BY score DESC LIMIT $%d`, tc, limitN)
		args = append([]any{query, agentID, query}, tcArgs...)
		args = append(args, limit)
	} else if userID != "" {
		// fixed params: $1=query, $2=agentID, $3=query, $4=userID
		tc, tcArgs, _, err := scopeClause(ctx, 5)
		if err != nil {
			return nil, err
		}
		limitN := 5 + len(tcArgs)
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				ts_rank(tsv, plainto_tsquery('russian', $1)) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND tsv @@ plainto_tsquery('russian', $3)
			AND (user_id IS NULL OR user_id = $4)%s
			ORDER BY score DESC LIMIT $%d`, tc, limitN)
		args = append([]any{query, agentID, query, userID}, tcArgs...)
		args = append(args, limit)
	} else {
		// fixed params: $1=query, $2=agentID, $3=query
		tc, tcArgs, _, err := scopeClause(ctx, 4)
		if err != nil {
			return nil, err
		}
		limitN := 4 + len(tcArgs)
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				ts_rank(tsv, plainto_tsquery('russian', $1)) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND tsv @@ plainto_tsquery('russian', $3)
			AND user_id IS NULL%s
			ORDER BY score DESC LIMIT $%d`, tc, limitN)
		args = append([]any{query, agentID, query}, tcArgs...)
		args = append(args, limit)
	}

	var rows []scoredChunkRow
	if err := pkgSqlxDB.SelectContext(ctx, &rows, q, args...); err != nil {
		return nil, err
	}
	results := make([]scoredChunk, len(rows))
	for i := range rows {
		results[i] = rows[i].toScoredChunk()
	}
	return results, nil
}

func (s *PGMemoryStore) vectorSearch(ctx context.Context, embedding []float32, agentID any, userID string, limit int) ([]scoredChunk, error) {
	vecStr := vectorToString(embedding)

	var q string
	var args []any

	if store.IsSharedMemory(ctx) {
		// Shared: no user_id filter — search ALL chunks for agent
		tc, tcArgs, _, err := scopeClause(ctx, 3)
		if err != nil {
			return nil, err
		}
		orderN := 3 + len(tcArgs)
		limitN := orderN + 1
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				1 - (embedding <=> $1::vector) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND embedding IS NOT NULL%s
			ORDER BY embedding <=> $%d::vector LIMIT $%d`, tc, orderN, limitN)
		args = append([]any{vecStr, agentID}, tcArgs...)
		args = append(args, vecStr, limit)
	} else if userID != "" {
		// fixed params: $1=vec, $2=agentID, $3=userID
		tc, tcArgs, _, err := scopeClause(ctx, 4)
		if err != nil {
			return nil, err
		}
		orderN := 4 + len(tcArgs)
		limitN := orderN + 1
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				1 - (embedding <=> $1::vector) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND embedding IS NOT NULL
			AND (user_id IS NULL OR user_id = $3)%s
			ORDER BY embedding <=> $%d::vector LIMIT $%d`, tc, orderN, limitN)
		args = append([]any{vecStr, agentID, userID}, tcArgs...)
		args = append(args, vecStr, limit)
	} else {
		// fixed params: $1=vec, $2=agentID
		tc, tcArgs, _, err := scopeClause(ctx, 3)
		if err != nil {
			return nil, err
		}
		orderN := 3 + len(tcArgs)
		limitN := orderN + 1
		q = fmt.Sprintf(`SELECT path, start_line, end_line, text, user_id,
				1 - (embedding <=> $1::vector) AS score
			FROM memory_chunks
			WHERE agent_id = $2 AND embedding IS NOT NULL
			AND user_id IS NULL%s
			ORDER BY embedding <=> $%d::vector LIMIT $%d`, tc, orderN, limitN)
		args = append([]any{vecStr, agentID}, tcArgs...)
		args = append(args, vecStr, limit)
	}

	var rows []scoredChunkRow
	if err := pkgSqlxDB.SelectContext(ctx, &rows, q, args...); err != nil {
		return nil, err
	}
	results := make([]scoredChunk, len(rows))
	for i := range rows {
		results[i] = rows[i].toScoredChunk()
	}
	return results, nil
}

// hybridMerge combines FTS and vector results with weighted scoring.
// Per-user results get a 1.2x boost. Deduplication: user copy wins over global.
// NOTE: when shared memory is active, the 1.2x personal boost still applies —
// consider removing it in shared mode if all docs should be treated equally.
func hybridMerge(fts, vec []scoredChunk, textWeight, vectorWeight float64, currentUserID string) []store.MemorySearchResult {
	type key struct {
		Path      string
		StartLine int
	}
	seen := make(map[key]*store.MemorySearchResult)

	addResult := func(r scoredChunk, weight float64) {
		k := key{r.Path, r.StartLine}
		scope := "global"
		boost := 1.0
		if r.UserID != nil && *r.UserID != "" {
			scope = "personal"
			boost = 1.2
		}
		score := r.Score * weight * boost

		if existing, ok := seen[k]; ok {
			existing.Score += score
			// User copy wins
			if scope == "personal" {
				existing.Scope = "personal"
				existing.Snippet = r.Text
			}
		} else {
			seen[k] = &store.MemorySearchResult{
				Path:      r.Path,
				StartLine: r.StartLine,
				EndLine:   r.EndLine,
				Score:     score,
				Snippet:   r.Text,
				Source:    "memory",
				Scope:     scope,
			}
		}
	}

	for _, r := range fts {
		addResult(r, textWeight)
	}
	for _, r := range vec {
		addResult(r, vectorWeight)
	}

	// Collect and sort by score
	results := make([]store.MemorySearchResult, 0, len(seen))
	for _, r := range seen {
		results = append(results, *r)
	}

	// Simple sort (descending score)
	for i := 0; i < len(results); i++ {
		for j := i + 1; j < len(results); j++ {
			if results[j].Score > results[i].Score {
				results[i], results[j] = results[j], results[i]
			}
		}
	}

	return results
}