125 lines
3.6 KiB
Go
125 lines
3.6 KiB
Go
// Package postgres implements knowledge's DocumentRepository against
|
|
// the pgvector-backed knowledge database. Transplant of
|
|
// internal/rag/store.go's SQL (vectorLiteral and all) plus
|
|
// internal/rag/index.go's pruneMissing, and internal/db.DB.
|
|
// ListIndexedFiles.
|
|
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/lib/pq"
|
|
|
|
"enkly/pkg/pgkit"
|
|
"enkly/services/knowledge/internal/domain"
|
|
)
|
|
|
|
// vectorLiteral formats a float32 embedding as the string pgvector
|
|
// expects for input, e.g. "[0.011,-0.42,...]" — lib/pq has no native
|
|
// vector type, so it's sent as text and cast with ::vector in SQL.
|
|
func vectorLiteral(v []float32) string {
|
|
parts := make([]string, len(v))
|
|
for i, f := range v {
|
|
parts[i] = strconv.FormatFloat(float64(f), 'f', 8, 32)
|
|
}
|
|
return "[" + strings.Join(parts, ",") + "]"
|
|
}
|
|
|
|
type DocumentRepo struct {
|
|
DB *pgkit.DB
|
|
}
|
|
|
|
func NewDocumentRepo(db *pgkit.DB) *DocumentRepo { return &DocumentRepo{DB: db} }
|
|
|
|
func (r *DocumentRepo) ExistingContentHash(ctx context.Context, filePath string, chunkIndex int) (string, error) {
|
|
var hash string
|
|
err := r.DB.QueryRowContext(ctx,
|
|
`SELECT content_hash FROM documents WHERE file_path = $1 AND chunk_index = $2`,
|
|
filePath, chunkIndex).Scan(&hash)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return "", nil
|
|
}
|
|
return hash, err
|
|
}
|
|
|
|
func (r *DocumentRepo) Upsert(ctx context.Context, filePath string, chunkIndex int, content, contentHash string, embedding []float32) error {
|
|
_, err := r.DB.ExecContext(ctx, `
|
|
INSERT INTO documents (file_path, chunk_index, content, content_hash, embedding, updated_at)
|
|
VALUES ($1, $2, $3, $4, $5::vector, now())
|
|
ON CONFLICT (file_path, chunk_index)
|
|
DO UPDATE SET content = EXCLUDED.content,
|
|
content_hash = EXCLUDED.content_hash,
|
|
embedding = EXCLUDED.embedding,
|
|
updated_at = now()`,
|
|
filePath, chunkIndex, content, contentHash, vectorLiteral(embedding))
|
|
return err
|
|
}
|
|
|
|
func (r *DocumentRepo) DeleteChunksFrom(ctx context.Context, filePath string, fromIndex int) error {
|
|
_, err := r.DB.ExecContext(ctx,
|
|
`DELETE FROM documents WHERE file_path = $1 AND chunk_index >= $2`, filePath, fromIndex)
|
|
return err
|
|
}
|
|
|
|
func (r *DocumentRepo) Search(ctx context.Context, embedding []float32, topK int) ([]domain.CodeChunk, error) {
|
|
lit := vectorLiteral(embedding)
|
|
rows, err := r.DB.QueryContext(ctx, `
|
|
SELECT id, file_path, chunk_index, content, embedding <=> $1::vector AS distance
|
|
FROM documents
|
|
ORDER BY embedding <=> $1::vector
|
|
LIMIT $2`, lit, topK)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var out []domain.CodeChunk
|
|
for rows.Next() {
|
|
var c domain.CodeChunk
|
|
if err := rows.Scan(&c.ID, &c.FilePath, &c.ChunkIndex, &c.Content, &c.Distance); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, c)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (r *DocumentRepo) PruneMissing(ctx context.Context, keep []string) (int, error) {
|
|
if keep == nil {
|
|
keep = []string{}
|
|
}
|
|
res, err := r.DB.ExecContext(ctx,
|
|
`DELETE FROM documents WHERE NOT (file_path = ANY($1))`, pq.Array(keep))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
n, err := res.RowsAffected()
|
|
return int(n), err
|
|
}
|
|
|
|
func (r *DocumentRepo) ListIndexedFiles(ctx context.Context, limit int) ([]string, error) {
|
|
if limit <= 0 {
|
|
limit = 200
|
|
}
|
|
rows, err := r.DB.QueryContext(ctx,
|
|
`SELECT DISTINCT file_path FROM documents ORDER BY file_path LIMIT $1`, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var out []string
|
|
for rows.Next() {
|
|
var p string
|
|
if err := rows.Scan(&p); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, p)
|
|
}
|
|
return out, rows.Err()
|
|
}
|