144 lines
6.6 KiB
Go
144 lines
6.6 KiB
Go
package kb
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"ai-operator/internal/embedding"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type PostgresRepository struct{ pool *pgxpool.Pool }
|
|
|
|
func NewPostgresRepository(pool *pgxpool.Pool) *PostgresRepository {
|
|
return &PostgresRepository{pool: pool}
|
|
}
|
|
|
|
func (r *PostgresRepository) Health(ctx context.Context) (Health, error) {
|
|
h := Health{}
|
|
if err := r.pool.Ping(ctx); err != nil {
|
|
return h, err
|
|
}
|
|
h.DBReachable = true
|
|
_ = r.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname='vector')`).Scan(&h.VectorExtension)
|
|
var mig int
|
|
_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM information_schema.tables WHERE table_name IN ('schema_migrations','knowledge_documents','knowledge_chunks','knowledge_ingest_runs','knowledge_search_logs')`).Scan(&mig)
|
|
h.MigrationsApplied = mig == 5
|
|
_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM knowledge_documents`).Scan(&h.Documents)
|
|
_ = r.pool.QueryRow(ctx, `SELECT count(*) FROM knowledge_chunks`).Scan(&h.Chunks)
|
|
return h, nil
|
|
}
|
|
|
|
func (r *PostgresRepository) UpsertDocument(ctx context.Context, doc Document, chunks []Chunk) error {
|
|
tx, err := r.pool.Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback(ctx)
|
|
docID := stableUUID(doc.ExternalID)
|
|
meta, _ := json.Marshal(doc.Metadata)
|
|
_, err = tx.Exec(ctx, `INSERT INTO knowledge_documents(id,external_id,title,language,region_code,category,source_type,source_uri,source_hash,status,content,question,short_answer,full_answer,keywords,source,metadata,priority,updated_at) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,now()) ON CONFLICT(external_id) DO UPDATE SET title=EXCLUDED.title, language=EXCLUDED.language, region_code=EXCLUDED.region_code, category=EXCLUDED.category, source_type=EXCLUDED.source_type, source_uri=EXCLUDED.source_uri, source_hash=EXCLUDED.source_hash, status=EXCLUDED.status, content=EXCLUDED.content, question=EXCLUDED.question, short_answer=EXCLUDED.short_answer, full_answer=EXCLUDED.full_answer, keywords=EXCLUDED.keywords, source=EXCLUDED.source, metadata=EXCLUDED.metadata, priority=EXCLUDED.priority, updated_at=now()`, docID, doc.ExternalID, doc.Title, doc.Language, doc.RegionCode, doc.Category, doc.SourceType, doc.SourceURI, doc.SourceHash, doc.Status, doc.Content, nullable(doc.Question), nullable(doc.ShortAnswer), nullable(doc.FullAnswer), nonNilKeywords(doc.Keywords), nullable(doc.Source), meta, doc.Priority)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec(ctx, `DELETE FROM knowledge_chunks WHERE document_id=$1`, docID); err != nil {
|
|
return err
|
|
}
|
|
for _, ch := range chunks {
|
|
chID := stableUUID(fmt.Sprintf("%s:%d:%s", doc.ExternalID, ch.ChunkIndex, ch.ContentHash))
|
|
cm, _ := json.Marshal(ch.Metadata)
|
|
_, err = tx.Exec(ctx, `INSERT INTO knowledge_chunks(id,document_id,chunk_index,language,region_code,title,content,content_hash,token_estimate,char_count,embedding,embedding_model,embedding_provider,metadata,tsv) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::vector,$12,$13,$14,to_tsvector('simple',$7))`, chID, docID, ch.ChunkIndex, ch.Language, ch.RegionCode, ch.Title, ch.Content, ch.ContentHash, ch.TokenEstimate, ch.CharCount, embedding.VectorLiteral(ch.Embedding), ch.EmbeddingModel, ch.EmbeddingProvider, cm)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return tx.Commit(ctx)
|
|
}
|
|
|
|
func (r *PostgresRepository) Search(ctx context.Context, req SearchRequest, queryEmbedding []float32) ([]SearchResult, error) {
|
|
limit := req.Limit
|
|
if limit <= 0 {
|
|
limit = 5
|
|
}
|
|
if limit > 10 {
|
|
limit = 10
|
|
}
|
|
langs := []string{req.Language}
|
|
regions := []string{req.RegionCode}
|
|
if req.IncludeGlobal {
|
|
regions = append(regions, "global")
|
|
}
|
|
vec := embedding.VectorLiteral(queryEmbedding)
|
|
rows, err := r.pool.Query(ctx, `WITH q AS (SELECT plainto_tsquery('simple', $1) query) SELECT d.id::text, c.id::text, c.title, c.content, c.language, c.region_code, d.category, coalesce(d.source_uri,''), coalesce(d.source,''), (1 - ((c.embedding <=> $2::vector)/2.0))::float8 AS vector_score, ts_rank_cd(c.tsv, q.query)::float8 AS keyword_score, CASE WHEN c.region_code=$3 THEN 1.0 ELSE 0.6 END AS region_boost FROM knowledge_chunks c JOIN knowledge_documents d ON d.id=c.document_id, q WHERE d.status='published' AND c.language=ANY($4) AND c.region_code=ANY($5) AND (d.valid_from IS NULL OR d.valid_from<=now()) AND (d.valid_to IS NULL OR d.valid_to>=now()) AND (c.tsv @@ q.query OR c.content ILIKE '%' || $1 || '%' OR c.embedding IS NOT NULL) ORDER BY ((0.65*(1-((c.embedding <=> $2::vector)/2.0))) + (0.25*ts_rank_cd(c.tsv,q.query)) + (0.10*CASE WHEN c.region_code=$3 THEN 1.0 ELSE 0.6 END)) DESC LIMIT $6`, sanitizeQuery(req.Query), vec, req.RegionCode, langs, regions, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []SearchResult
|
|
for rows.Next() {
|
|
var r0 SearchResult
|
|
var source string
|
|
var v, k, boost float64
|
|
if err := rows.Scan(&r0.DocumentID, &r0.ChunkID, &r0.Title, &r0.Content, &r0.Language, &r0.RegionCode, &r0.Category, &r0.SourceURI, &source, &v, &k, &boost); err != nil {
|
|
return nil, err
|
|
}
|
|
r0.SourceLanguage = r0.Language
|
|
r0.VectorScore = v
|
|
r0.KeywordScore = k
|
|
r0.Score = 0.65*v + 0.25*k + 0.10*boost
|
|
r0.Citation = Citation{DocumentTitle: r0.Title, SourceURI: r0.SourceURI, ChunkID: r0.ChunkID, RegionCode: r0.RegionCode, Language: r0.Language, Source: source}
|
|
if r0.Score >= req.MinScore {
|
|
out = append(out, r0)
|
|
}
|
|
}
|
|
_ = r.logSearch(ctx, req, out, false)
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (r *PostgresRepository) logSearch(ctx context.Context, req SearchRequest, results []SearchResult, fallback bool) error {
|
|
top := 0.0
|
|
if len(results) > 0 {
|
|
top = results[0].Score
|
|
}
|
|
_, err := r.pool.Exec(ctx, `INSERT INTO knowledge_search_logs(id,call_id,query,language,region_code,result_count,top_score,cross_language_fallback_used) VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, stableUUID(fmt.Sprintf("%s:%s:%d", req.CallID, req.Query, len(results))), nullable(req.CallID), truncate(req.Query, 1000), req.Language, req.RegionCode, len(results), top, fallback)
|
|
return err
|
|
}
|
|
|
|
func nullable(s string) any {
|
|
if strings.TrimSpace(s) == "" {
|
|
return nil
|
|
}
|
|
return s
|
|
}
|
|
func sanitizeQuery(s string) string {
|
|
return strings.Map(func(r rune) rune {
|
|
if r < 32 {
|
|
return -1
|
|
}
|
|
return r
|
|
}, truncate(strings.TrimSpace(s), 1000))
|
|
}
|
|
func truncate(s string, n int) string {
|
|
r := []rune(s)
|
|
if len(r) > n {
|
|
return string(r[:n])
|
|
}
|
|
return s
|
|
}
|
|
func stableUUID(s string) string {
|
|
h := ContentHash(s)
|
|
return fmt.Sprintf("%s-%s-%s-%s-%s", h[0:8], h[8:12], h[12:16], h[16:20], h[20:32])
|
|
}
|
|
|
|
var _ pgx.Tx
|
|
|
|
func nonNilKeywords(v []string) []string {
|
|
if v == nil {
|
|
return []string{}
|
|
}
|
|
return v
|
|
}
|