mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-12 20:15:42 +00:00
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.
904 lines
29 KiB
Go
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
|
|
}
|
|
|