Files
Chandan Kumar 65dbbcfd59 fix(vault): search document bodies and match any query word (#1554)
vault_search only indexed title + path + the auto-summary, and the summary
is written from the first 3000 runes of the file. Anything the summary left
out (room names, prices, codes) could never be found. On top of that the
FTS query used plainto_tsquery, which ANDs every word, so a normal question
like "what is the price of the Deluxe Ocean Suite?" matched nothing even when
the key words were indexed.

- Add vault_document_chunks (migration 98): the file body split into
  chunks with their own tsvector and embedding. body_indexed_hash on
  vault_documents records which content_hash the chunks came from, so
  unchanged files are not re-chunked or re-embedded.
- The enrich worker rebuilds chunks when a file changes. This needs no LLM,
  so it runs even when no provider is configured.
- Rescan backfills chunks for docs indexed before this change, since they
  already have a summary and never go back through the worker.
- FTS matches any query word; ts_rank still ranks docs with more matching
  words first. Both FTS and vector search look at the doc and its chunks
  and score each doc by its best hit.

SQLite is unchanged: its vault search is LIKE on title/path only.
2026-09-27 13:13:40 +05:30

904 lines
29 KiB
Go

package pg
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"strings"
"log/slog"
"sort"
"time"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// appendTeamFilter appends the team_id clause to a vault query.
// TeamIDs (personal + listed teams) takes precedence over TeamID (single team / personal-only).
func appendTeamFilter(q string, args []any, p int, teamID *string, teamIDs []string) (string, []any, int) {
if len(teamIDs) > 0 {
ph := make([]string, len(teamIDs))
for i, id := range teamIDs {
ph[i] = fmt.Sprintf("$%d", p)
args = append(args, parseUUIDOrNil(id))
p++
}
q += " AND (team_id IS NULL OR team_id IN (" + strings.Join(ph, ",") + "))"
} else if teamID != nil {
if *teamID != "" {
q += fmt.Sprintf(" AND (team_id = $%d OR team_id IS NULL)", p)
args = append(args, parseUUIDOrNil(*teamID))
p++
} else {
q += " AND team_id IS NULL"
}
}
return q, args, p
}
// PGVaultStore implements store.VaultStore backed by PostgreSQL.
type PGVaultStore struct {
db *sql.DB
embProvider store.EmbeddingProvider
}
// NewPGVaultStore creates a new PG-backed vault store.
func NewPGVaultStore(db *sql.DB) *PGVaultStore {
return &PGVaultStore{db: db}
}
func (s *PGVaultStore) SetEmbeddingProvider(provider store.EmbeddingProvider) {
s.embProvider = provider
}
func (s *PGVaultStore) Close() error { return nil }
// optAgentUUID converts a nullable *string agent_id to *uuid.UUID for SQL.
// Returns (nil, nil) when the input is nil or empty — a legitimate SQL NULL.
// Returns (nil, error) on a non-empty, non-UUID input — propagating the error
// prevents silent-nil writes that would otherwise corrupt data.
// See docs/agent-identity-conventions.md.
func optAgentUUID(agentID *string) (*uuid.UUID, error) {
if agentID == nil || *agentID == "" {
return nil, nil
}
u, err := parseUUID(*agentID)
if err != nil {
return nil, err
}
return &u, nil
}
// UpsertDocument inserts or updates a vault document.
func (s *PGVaultStore) UpsertDocument(ctx context.Context, doc *store.VaultDocument) error {
tid, err := parseUUID(doc.TenantID)
if err != nil {
return fmt.Errorf("vault upsert: tenant: %w", err)
}
aid, err := optAgentUUID(doc.AgentID)
if err != nil {
return fmt.Errorf("vault upsert: agent: %w", err)
}
now := time.Now().UTC()
meta, err := json.Marshal(doc.Metadata)
if err != nil {
meta = []byte("{}")
}
id := uuid.Must(uuid.NewV7())
var embStr *string
if s.embProvider != nil && doc.Summary != "" {
// Embed title + path + summary for richer vector search.
embedText := doc.Title + " " + doc.Path
if doc.Summary != "" {
embedText += " " + doc.Summary
}
vecs, embErr := s.embProvider.Embed(ctx, []string{embedText})
if embErr == nil && len(vecs) > 0 {
v := vectorToString(vecs[0])
embStr = &v
}
}
var teamID *uuid.UUID
if doc.TeamID != nil && *doc.TeamID != "" {
t, err := parseUUID(*doc.TeamID)
if err != nil {
return fmt.Errorf("vault upsert: team: %w", err)
}
teamID = &t
}
var actualID uuid.UUID
// Normalize chat_id: empty string → NULL.
var chatID *string
if doc.ChatID != nil && *doc.ChatID != "" {
c := *doc.ChatID
chatID = &c
}
err = s.db.QueryRowContext(ctx, `
INSERT INTO vault_documents
(id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, title, doc_type, content_hash, summary, embedding, metadata, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $15)
ON CONFLICT (tenant_id, COALESCE(agent_id, '00000000-0000-0000-0000-000000000000'::uuid), COALESCE(team_id, '00000000-0000-0000-0000-000000000000'::uuid), scope, path) DO UPDATE SET
title = EXCLUDED.title,
doc_type = EXCLUDED.doc_type,
content_hash = EXCLUDED.content_hash,
summary = EXCLUDED.summary,
embedding = COALESCE(EXCLUDED.embedding, vault_documents.embedding),
metadata = EXCLUDED.metadata,
chat_id = COALESCE(EXCLUDED.chat_id, vault_documents.chat_id),
tenant_id = EXCLUDED.tenant_id,
updated_at = EXCLUDED.updated_at
RETURNING id`,
id, tid, aid, teamID, chatID, doc.Scope, doc.CustomScope, doc.Path, doc.Title, doc.DocType,
doc.ContentHash, doc.Summary, embStr, meta, now,
).Scan(&actualID)
if err != nil {
return fmt.Errorf("vault upsert document: %w", err)
}
doc.ID = actualID.String()
return nil
}
// GetDocument retrieves a vault document by tenant, agent, and path.
// Empty agentID means no agent filter (match any agent).
// Team scoping via RunContext: present+TeamID → filter; present+empty → personal; nil → any match.
func (s *PGVaultStore) GetDocument(ctx context.Context, tenantID, agentID, path string) (*store.VaultDocument, error) {
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault get document: tenant: %w", err)
}
q := `SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at
FROM vault_documents WHERE tenant_id = $1 AND path = $2`
args := []any{tid, path}
p := 3
if agentID != "" {
aid, err := parseUUID(agentID)
if err != nil {
return nil, fmt.Errorf("vault get document: agent: %w", err)
}
q += fmt.Sprintf(" AND agent_id = $%d", p)
args = append(args, aid)
p++
}
if rc := store.RunContextFromCtx(ctx); rc != nil {
if rc.TeamID != "" {
tmid, err := parseUUID(rc.TeamID)
if err != nil {
return nil, fmt.Errorf("vault get document: team: %w", err)
}
q += fmt.Sprintf(" AND team_id = $%d", p)
args = append(args, tmid)
} else {
q += " AND team_id IS NULL"
}
}
var row vaultDocRow
// Scan order MUST match SELECT order above: 15 columns including
// path_basename (generated column added in migration 000047).
err = s.db.QueryRowContext(ctx, q, args...).Scan(
&row.ID, &row.TenantID, &row.AgentID, &row.TeamID, &row.ChatID, &row.Scope, &row.CustomScope,
&row.Path, &row.PathBasename, &row.Title, &row.DocType, &row.ContentHash, &row.Summary,
&row.MetaJSON, &row.CreatedAt, &row.UpdatedAt)
if err != nil {
return nil, err
}
doc := row.toVaultDocument()
return &doc, nil
}
// GetDocumentByID retrieves a vault document by ID with tenant isolation.
func (s *PGVaultStore) GetDocumentByID(ctx context.Context, tenantID, id string) (*store.VaultDocument, error) {
uid, err := parseUUID(id)
if err != nil {
return nil, fmt.Errorf("vault get document by id: id: %w", err)
}
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault get document by id: tenant: %w", err)
}
var row vaultDocRow
err = s.db.QueryRowContext(ctx, `
SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at
FROM vault_documents WHERE id = $1 AND tenant_id = $2`, uid, tid,
).Scan(&row.ID, &row.TenantID, &row.AgentID, &row.TeamID, &row.ChatID, &row.Scope, &row.CustomScope,
&row.Path, &row.PathBasename, &row.Title, &row.DocType, &row.ContentHash, &row.Summary,
&row.MetaJSON, &row.CreatedAt, &row.UpdatedAt)
if err != nil {
return nil, err
}
doc := row.toVaultDocument()
return &doc, nil
}
// GetDocumentsByIDs returns documents matching the given IDs with tenant isolation.
// Chunks by 500 to stay within PG param limits.
func (s *PGVaultStore) GetDocumentsByIDs(ctx context.Context, tenantID string, docIDs []string) ([]store.VaultDocument, error) {
if len(docIDs) == 0 {
return nil, nil
}
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault get by ids: tenant: %w", err)
}
const chunkSize = 500
var all []store.VaultDocument
for start := 0; start < len(docIDs); start += chunkSize {
end := min(start+chunkSize, len(docIDs))
var scanned []vaultDocRow
if err := pkgSqlxDB.SelectContext(ctx, &scanned,
`SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at
FROM vault_documents WHERE id = ANY($1) AND tenant_id = $2`,
pqStringArray(docIDs[start:end]), tid); err != nil {
return nil, err
}
for i := range scanned {
all = append(all, scanned[i].toVaultDocument())
}
}
return all, nil
}
// GetDocumentByBasename finds a document by path basename (case-insensitive).
// Uses the stored generated column path_basename + index for fast lookup.
func (s *PGVaultStore) GetDocumentByBasename(ctx context.Context, tenantID, agentID, basename string) (*store.VaultDocument, error) {
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault get by basename: tenant: %w", err)
}
q := `SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at
FROM vault_documents
WHERE tenant_id = $1 AND path_basename = lower($2)`
args := []any{tid, basename}
p := 3
if agentID != "" {
aid, err := parseUUID(agentID)
if err != nil {
return nil, fmt.Errorf("vault get by basename: agent: %w", err)
}
q += fmt.Sprintf(" AND agent_id = $%d", p)
args = append(args, aid)
}
q += " LIMIT 1"
var row vaultDocRow
// Scan order MUST match SELECT order above: 15 columns including
// path_basename (generated column added in migration 000047).
err = s.db.QueryRowContext(ctx, q, args...).Scan(
&row.ID, &row.TenantID, &row.AgentID, &row.TeamID, &row.ChatID, &row.Scope, &row.CustomScope,
&row.Path, &row.PathBasename, &row.Title, &row.DocType, &row.ContentHash, &row.Summary,
&row.MetaJSON, &row.CreatedAt, &row.UpdatedAt)
if err != nil {
return nil, err
}
doc := row.toVaultDocument()
return &doc, nil
}
// DeleteDocument removes a vault document by tenant, agent, and path.
// Empty agentID means no agent filter.
// Team scoping via RunContext (same rules as GetDocument).
func (s *PGVaultStore) DeleteDocument(ctx context.Context, tenantID, agentID, path string) error {
tid, err := parseUUID(tenantID)
if err != nil {
return fmt.Errorf("vault delete document: tenant: %w", err)
}
q := `DELETE FROM vault_documents WHERE tenant_id = $1 AND path = $2`
args := []any{tid, path}
p := 3
if agentID != "" {
aid, err := parseUUID(agentID)
if err != nil {
return fmt.Errorf("vault delete document: agent: %w", err)
}
q += fmt.Sprintf(" AND agent_id = $%d", p)
args = append(args, aid)
p++
}
if rc := store.RunContextFromCtx(ctx); rc != nil {
if rc.TeamID != "" {
tmid, err := parseUUID(rc.TeamID)
if err != nil {
return fmt.Errorf("vault delete document: team: %w", err)
}
q += fmt.Sprintf(" AND team_id = $%d", p)
args = append(args, tmid)
} else {
q += " AND team_id IS NULL"
}
}
_, err = s.db.ExecContext(ctx, q, args...)
return err
}
// ListDocuments returns vault documents for a tenant+agent with optional filters.
func (s *PGVaultStore) ListDocuments(ctx context.Context, tenantID, agentID string, opts store.VaultListOptions) ([]store.VaultDocument, error) {
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault list documents: tenant: %w", err)
}
q := `SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at
FROM vault_documents WHERE tenant_id = $1`
args := []any{tid}
p := 2
// Agent filter is optional — omit for cross-agent listing.
if agentID != "" {
aid, err := parseUUID(agentID)
if err != nil {
return nil, fmt.Errorf("vault list documents: agent: %w", err)
}
q += fmt.Sprintf(" AND (agent_id = $%d OR agent_id IS NULL)", p)
args = append(args, aid)
p++
}
q, args, p = appendTeamFilter(q, args, p, opts.TeamID, opts.TeamIDs)
if opts.Scope != "" {
q += fmt.Sprintf(" AND scope = $%d", p)
args = append(args, opts.Scope)
p++
}
if len(opts.DocTypes) > 0 {
q += fmt.Sprintf(" AND doc_type = ANY($%d)", p)
args = append(args, pqStringArray(opts.DocTypes))
p++
}
q += " ORDER BY updated_at DESC"
limit := opts.Limit
if limit <= 0 {
limit = 100
}
q += fmt.Sprintf(" LIMIT $%d", p)
args = append(args, limit)
p++
if opts.Offset > 0 {
q += fmt.Sprintf(" OFFSET $%d", p)
args = append(args, opts.Offset)
}
var scanned []vaultDocRow
if err := pkgSqlxDB.SelectContext(ctx, &scanned, q, args...); err != nil {
return nil, err
}
docs := make([]store.VaultDocument, 0, len(scanned))
for i := range scanned {
docs = append(docs, scanned[i].toVaultDocument())
}
return docs, nil
}
// CountDocuments returns the total number of vault documents matching the given filters.
func (s *PGVaultStore) CountDocuments(ctx context.Context, tenantID, agentID string, opts store.VaultListOptions) (int, error) {
tid, err := parseUUID(tenantID)
if err != nil {
return 0, fmt.Errorf("vault count documents: tenant: %w", err)
}
q := `SELECT COUNT(*) FROM vault_documents WHERE tenant_id = $1`
args := []any{tid}
p := 2
if agentID != "" {
aid, err := parseUUID(agentID)
if err != nil {
return 0, fmt.Errorf("vault count documents: agent: %w", err)
}
q += fmt.Sprintf(" AND (agent_id = $%d OR agent_id IS NULL)", p)
args = append(args, aid)
p++
}
q, args, p = appendTeamFilter(q, args, p, opts.TeamID, opts.TeamIDs)
if opts.Scope != "" {
q += fmt.Sprintf(" AND scope = $%d", p)
args = append(args, opts.Scope)
p++
}
if len(opts.DocTypes) > 0 {
q += fmt.Sprintf(" AND doc_type = ANY($%d)", p)
args = append(args, pqStringArray(opts.DocTypes))
}
var count int
if err := s.db.QueryRowContext(ctx, q, args...).Scan(&count); err != nil {
return 0, fmt.Errorf("vault count documents: %w", err)
}
return count, nil
}
// UpdateHash updates the content hash for a vault document with tenant isolation.
func (s *PGVaultStore) UpdateHash(ctx context.Context, tenantID, id, newHash string) error {
uid, err := parseUUID(id)
if err != nil {
return fmt.Errorf("vault update hash: id: %w", err)
}
tid, err := parseUUID(tenantID)
if err != nil {
return fmt.Errorf("vault update hash: tenant: %w", err)
}
_, err = s.db.ExecContext(ctx,
`UPDATE vault_documents SET content_hash = $1, updated_at = $2 WHERE id = $3 AND tenant_id = $4`,
newHash, time.Now().UTC(), uid, tid)
return err
}
// UpdateSummaryAndReembed updates summary and re-generates embedding from title+path+summary.
// UpdateSummaryAndReembed and FindSimilarDocs moved to vault_documents_enrichment.go.
// Search performs hybrid FTS + vector search on vault documents and their body chunks.
func (s *PGVaultStore) Search(ctx context.Context, opts store.VaultSearchOptions) ([]store.VaultSearchResult, error) {
tid, err := parseUUID(opts.TenantID)
if err != nil {
return nil, fmt.Errorf("vault search: tenant: %w", err)
}
aid, err := optAgentUUID(&opts.AgentID) // empty string → nil → no agent filter
if err != nil {
return nil, fmt.Errorf("vault search: agent: %w", err)
}
f := vaultSearchScope{
agentID: aid,
team: buildSearchTeamFilter(opts.TeamID, opts.TeamIDs),
// Chat-scope filter (applies only when team is isolated + chat_id non-nil/non-empty).
chat: buildSearchChatFilter(opts.ChatID, opts.TeamIsolated),
scope: opts.Scope,
docTypes: opts.DocTypes,
}
maxResults := opts.MaxResults
if maxResults <= 0 {
maxResults = 10
}
// FTS search
ftsResults, err := s.ftsSearch(ctx, opts.Query, tid, f, maxResults*2)
if err != nil {
return nil, err
}
// Vector search if provider available
var vecResults []store.VaultSearchResult
if s.embProvider != nil {
vecs, embErr := s.embProvider.Embed(ctx, []string{opts.Query})
if embErr == nil && len(vecs) > 0 {
var vecErr error
vecResults, vecErr = s.vectorSearch(ctx, vecs[0], tid, f, maxResults*2)
if vecErr != nil {
slog.Debug("vault.vector_search_fallback", "err", vecErr)
vecResults = nil
}
}
}
// Merge: FTS weight 0.4, vector weight 0.6
merged := s.mergeResults(ftsResults, vecResults, 0.4, 0.6, maxResults)
// Apply min score filter
if opts.MinScore > 0 {
var filtered []store.VaultSearchResult
for _, r := range merged {
if r.Score >= opts.MinScore {
filtered = append(filtered, r)
}
}
return filtered, nil
}
return merged, nil
}
// searchTeamFilter holds pre-parsed team filter info for search sub-queries.
type searchTeamFilter struct {
teamIDs []uuid.UUID // personal + these teams (len > 0 = multi-team mode)
teamID *uuid.UUID // single team filter (nil + active = personal-only)
active bool // whether any team filter is applied
}
func buildSearchTeamFilter(teamID *string, teamIDs []string) searchTeamFilter {
if len(teamIDs) > 0 {
uuids := make([]uuid.UUID, len(teamIDs))
for i, id := range teamIDs {
uuids[i] = parseUUIDOrNil(id)
}
return searchTeamFilter{teamIDs: uuids, active: true}
}
if teamID != nil {
if *teamID != "" {
t := parseUUIDOrNil(*teamID)
return searchTeamFilter{teamID: &t, active: true}
}
return searchTeamFilter{active: true} // personal-only
}
return searchTeamFilter{} // no filter
}
// appendSearchTeamClause appends team filter SQL to a search query.
func (tf searchTeamFilter) append(q string, args []any, p int) (string, []any, int) {
if !tf.active {
return q, args, p
}
if len(tf.teamIDs) > 0 {
ph := make([]string, len(tf.teamIDs))
for i, id := range tf.teamIDs {
ph[i] = fmt.Sprintf("$%d", p)
args = append(args, id)
p++
}
q += " AND (team_id IS NULL OR team_id IN (" + strings.Join(ph, ",") + "))"
} else if tf.teamID != nil {
q += fmt.Sprintf(" AND (team_id = $%d OR team_id IS NULL)", p)
args = append(args, *tf.teamID)
p++
} else {
q += " AND team_id IS NULL"
}
return q, args, p
}
// searchChatFilter isolates vault search by chat_id when the calling team uses isolated workspace.
// Predicate: (chat_id = $N OR chat_id IS NULL). NULL = team-wide doc (legacy or shared-mode write).
type searchChatFilter struct {
chatID string
active bool
}
func buildSearchChatFilter(chatID *string, teamIsolated bool) searchChatFilter {
if !teamIsolated || chatID == nil || *chatID == "" {
return searchChatFilter{}
}
return searchChatFilter{chatID: *chatID, active: true}
}
func (cf searchChatFilter) append(q string, args []any, p int) (string, []any, int) {
if !cf.active {
return q, args, p
}
q += fmt.Sprintf(" AND (chat_id = $%d OR chat_id IS NULL)", p)
args = append(args, cf.chatID)
p++
return q, args, p
}
// vaultAnyTermQuery matches docs containing any word of $1. plainto_tsquery
// ANDs every word, and the 'simple' config keeps filler words like "what" or
// "the", so a plain question never matched anything. Lexemes come out of
// plainto_tsquery already quoted and never contain spaces, so swapping the
// operators is safe. ts_rank still puts docs matching more words first.
const vaultAnyTermQuery = `replace(plainto_tsquery('simple', $1)::text, ' & ', ' | ')::tsquery`
// vaultDocCols is the vault_documents column list scanned into vaultDocRow,
// qualified for queries that join the chunks table.
const vaultDocCols = `d.id, d.tenant_id, d.agent_id, d.team_id, d.chat_id, d.scope, d.custom_scope, d.path, d.path_basename,
d.title, d.doc_type, d.content_hash, d.summary, d.metadata, d.created_at, d.updated_at`
// vaultSearchScope holds the visibility filters shared by every search sub-query.
type vaultSearchScope struct {
agentID *uuid.UUID
team searchTeamFilter
chat searchChatFilter
scope string
docTypes []string
}
// append adds the filters to q. Columns are unqualified; that works for the
// chunk joins too because none of them exist on vault_document_chunks.
func (f vaultSearchScope) append(q string, args []any, p int) (string, []any, int) {
if f.agentID != nil {
q += fmt.Sprintf(" AND (agent_id = $%d OR agent_id IS NULL)", p)
args = append(args, *f.agentID)
p++
}
q, args, p = f.team.append(q, args, p)
q, args, p = f.chat.append(q, args, p)
if f.scope != "" {
q += fmt.Sprintf(" AND scope = $%d", p)
args = append(args, f.scope)
p++
}
if len(f.docTypes) > 0 {
q += fmt.Sprintf(" AND doc_type = ANY($%d)", p)
args = append(args, pqStringArray(f.docTypes))
p++
}
return q, args, p
}
// selectSearchRows applies the scope filters, appends tail (which gets the
// LIMIT placeholder number) and runs the query.
func selectSearchRows(ctx context.Context, q string, args []any, f vaultSearchScope, tail string, limit int) ([]store.VaultSearchResult, error) {
q, args, p := f.append(q, args, len(args)+1)
q += fmt.Sprintf(tail, p)
args = append(args, limit)
var scanned []vaultSearchRow
if err := pkgSqlxDB.SelectContext(ctx, &scanned, q, args...); err != nil {
return nil, err
}
return vaultSearchRowsToResults(scanned, "vault"), nil
}
// ftsSearch matches the query against the document header (title, path,
// summary) and the body chunks, scoring each doc by its best hit.
func (s *PGVaultStore) ftsSearch(ctx context.Context, query string, tenantID uuid.UUID, f vaultSearchScope, limit int) ([]store.VaultSearchResult, error) {
docHits, err := selectSearchRows(ctx, `SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at,
ts_rank(tsv, `+vaultAnyTermQuery+`) AS score
FROM vault_documents
WHERE tenant_id = $2 AND tsv @@ `+vaultAnyTermQuery,
[]any{query, tenantID}, f, " ORDER BY score DESC LIMIT $%d", limit)
if err != nil {
return nil, err
}
// GROUP BY d.id covers the other d columns since id is the primary key.
chunkHits, err := selectSearchRows(ctx, `SELECT `+vaultDocCols+`,
MAX(ts_rank(c.tsv, `+vaultAnyTermQuery+`)) AS score
FROM vault_document_chunks c
JOIN vault_documents d ON d.id = c.document_id
WHERE c.tenant_id = $2 AND c.tsv @@ `+vaultAnyTermQuery,
[]any{query, tenantID}, f, " GROUP BY d.id ORDER BY score DESC LIMIT $%d", limit)
if err != nil {
return nil, err
}
return bestHitPerDoc(docHits, chunkHits), nil
}
// vectorSearch ranks docs by cosine similarity of the header embedding or of
// the closest body chunk, whichever is higher.
func (s *PGVaultStore) vectorSearch(ctx context.Context, embedding []float32, tenantID uuid.UUID, f vaultSearchScope, limit int) ([]store.VaultSearchResult, error) {
vecStr := vectorToString(embedding)
docHits, err := selectSearchRows(ctx, `SELECT id, tenant_id, agent_id, team_id, chat_id, scope, custom_scope, path, path_basename, title, doc_type, content_hash, summary, metadata, created_at, updated_at,
1 - (embedding <=> $1) AS score
FROM vault_documents
WHERE tenant_id = $2 AND embedding IS NOT NULL`,
[]any{vecStr, tenantID}, f, " ORDER BY embedding <=> $1 LIMIT $%d", limit)
if err != nil {
return nil, err
}
// One doc can own several of the nearest chunks, so fetch extra rows to
// still end up with enough distinct docs after bestHitPerDoc.
chunkHits, err := selectSearchRows(ctx, `SELECT `+vaultDocCols+`,
1 - (c.embedding <=> $1) AS score
FROM vault_document_chunks c
JOIN vault_documents d ON d.id = c.document_id
WHERE c.tenant_id = $2 AND c.embedding IS NOT NULL`,
[]any{vecStr, tenantID}, f, " ORDER BY c.embedding <=> $1 LIMIT $%d", limit*3)
if err != nil {
return nil, err
}
return bestHitPerDoc(docHits, chunkHits), nil
}
// bestHitPerDoc keeps the highest score seen for each document, sorted by
// score. mergeResults adds scores per doc, so duplicates must go first.
func bestHitPerDoc(lists ...[]store.VaultSearchResult) []store.VaultSearchResult {
best := make(map[string]int)
var out []store.VaultSearchResult
for _, list := range lists {
for _, r := range list {
if i, ok := best[r.Document.ID]; ok {
if r.Score > out[i].Score {
out[i] = r
}
continue
}
best[r.Document.ID] = len(out)
out = append(out, r)
}
}
sort.SliceStable(out, func(i, j int) bool { return out[i].Score > out[j].Score })
return out
}
// vaultSearchRowsToResults converts a slice of vaultSearchRow to store.VaultSearchResult.
func vaultSearchRowsToResults(rows []vaultSearchRow, source string) []store.VaultSearchResult {
results := make([]store.VaultSearchResult, 0, len(rows))
for i := range rows {
results = append(results, rows[i].toVaultSearchResult(source))
}
return results
}
// mergeResults combines FTS and vector results with weighted scoring.
func (s *PGVaultStore) mergeResults(fts, vec []store.VaultSearchResult, ftsW, vecW float64, maxResults int) []store.VaultSearchResult {
seen := make(map[string]*store.VaultSearchResult)
// Normalize FTS scores
var maxFTS float64
for _, r := range fts {
if r.Score > maxFTS {
maxFTS = r.Score
}
}
for _, r := range fts {
norm := r.Score
if maxFTS > 0 {
norm = r.Score / maxFTS
}
r.Score = norm * ftsW
seen[r.Document.ID] = &r
}
// Normalize vector scores and merge
var maxVec float64
for _, r := range vec {
if r.Score > maxVec {
maxVec = r.Score
}
}
for _, r := range vec {
norm := r.Score
if maxVec > 0 {
norm = r.Score / maxVec
}
if existing, ok := seen[r.Document.ID]; ok {
existing.Score += norm * vecW
} else {
r.Score = norm * vecW
seen[r.Document.ID] = &r
}
}
// Collect and sort
results := make([]store.VaultSearchResult, 0, len(seen))
for _, r := range seen {
results = append(results, *r)
}
// Sort descending by score
sort.Slice(results, func(i, j int) bool {
return results[i].Score > results[j].Score
})
if len(results) > maxResults {
results = results[:maxResults]
}
return results
}
// ListTreeEntries returns immediate children (files + virtual folders) under the given path prefix.
func (s *PGVaultStore) ListTreeEntries(ctx context.Context, tenantID string, opts store.VaultTreeOptions) ([]store.VaultTreeEntry, error) {
tid, err := parseUUID(tenantID)
if err != nil {
return nil, fmt.Errorf("vault tree: tenant: %w", err)
}
prefix := opts.Path
if prefix != "" && !strings.HasSuffix(prefix, "/") {
prefix += "/"
}
fileQ := `SELECT id, path, title, doc_type, scope, updated_at FROM vault_documents WHERE tenant_id = $1`
fileArgs := []any{tid}
fp := 2
if prefix == "" {
fileQ += " AND path NOT LIKE '%/%'"
} else {
fileQ += fmt.Sprintf(" AND path LIKE $%d AND path NOT LIKE $%d", fp, fp+1)
fileArgs = append(fileArgs, prefix+"%", prefix+"%/%")
fp += 2
}
fileQ, fileArgs, fp = appendTreeFilters(fileQ, fileArgs, fp, opts)
fileQ += " ORDER BY path"
deepQ := `SELECT DISTINCT path FROM vault_documents WHERE tenant_id = $1`
deepArgs := []any{tid}
dp := 2
if prefix == "" {
deepQ += " AND path LIKE '%/%'"
} else {
deepQ += fmt.Sprintf(" AND path LIKE $%d", dp)
deepArgs = append(deepArgs, prefix+"%/%")
dp++
}
deepQ, deepArgs, _ = appendTreeFilters(deepQ, deepArgs, dp, opts)
deepQ += " LIMIT 50000"
fileRows, err := s.db.QueryContext(ctx, fileQ, fileArgs...)
if err != nil {
return nil, fmt.Errorf("vault tree files: %w", err)
}
defer fileRows.Close()
var entries []store.VaultTreeEntry
for fileRows.Next() {
var id, path, title, docType, scope string
var updatedAt time.Time
if err := fileRows.Scan(&id, &path, &title, &docType, &scope, &updatedAt); err != nil {
return nil, fmt.Errorf("vault tree scan: %w", err)
}
name := path
if idx := strings.LastIndex(path, "/"); idx >= 0 {
name = path[idx+1:]
}
ua := updatedAt
entries = append(entries, store.VaultTreeEntry{
Name: name, Path: path, DocID: id, DocType: docType, Scope: scope, Title: title, UpdatedAt: &ua,
})
}
if err := fileRows.Err(); err != nil {
return nil, fmt.Errorf("vault tree files: %w", err)
}
deepRows, err := s.db.QueryContext(ctx, deepQ, deepArgs...)
if err != nil {
return nil, fmt.Errorf("vault tree deep: %w", err)
}
defer deepRows.Close()
var deepPaths []string
for deepRows.Next() {
var p string
if err := deepRows.Scan(&p); err != nil {
return nil, fmt.Errorf("vault tree deep scan: %w", err)
}
deepPaths = append(deepPaths, p)
}
if err := deepRows.Err(); err != nil {
return nil, fmt.Errorf("vault tree deep: %w", err)
}
for _, fname := range extractFolderNames(prefix, deepPaths) {
entries = append(entries, store.VaultTreeEntry{
Name: fname, Path: prefix + fname, IsDir: true, HasChildren: true,
})
}
sort.Slice(entries, func(i, j int) bool {
if entries[i].IsDir != entries[j].IsDir {
return entries[i].IsDir
}
return entries[i].Name < entries[j].Name
})
return entries, nil
}
func appendTreeFilters(q string, args []any, p int, opts store.VaultTreeOptions) (string, []any, int) {
if opts.AgentID != "" {
aid, err := parseUUID(opts.AgentID)
if err == nil {
q += fmt.Sprintf(" AND (agent_id = $%d OR agent_id IS NULL)", p)
args = append(args, aid)
p++
}
}
q, args, p = appendTeamFilter(q, args, p, opts.TeamID, opts.TeamIDs)
if opts.Scope != "" {
q += fmt.Sprintf(" AND scope = $%d", p)
args = append(args, opts.Scope)
p++
}
if len(opts.DocTypes) > 0 {
q += fmt.Sprintf(" AND doc_type = ANY($%d)", p)
args = append(args, pqStringArray(opts.DocTypes))
p++
}
return q, args, p
}
func extractFolderNames(prefix string, deepPaths []string) []string {
seen := make(map[string]bool)
var folders []string
for _, p := range deepPaths {
rest := strings.TrimPrefix(p, prefix)
if idx := strings.Index(rest, "/"); idx > 0 {
seg := rest[:idx]
if !seen[seg] {
seen[seg] = true
folders = append(folders, seg)
}
}
}
sort.Strings(folders)
return folders
}