415 lines
13 KiB
Go
415 lines
13 KiB
Go
|
package vectordb
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"database/sql"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"os"
|
||
|
|
"path/filepath"
|
||
|
|
"regexp"
|
||
|
|
"strings"
|
||
|
|
|
||
|
|
"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
|
||
|
|
|
||
|
|
turso "turso.tech/database/tursogo"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
defaultLibSQLPath = "data/agent/vectors.db"
|
||
|
|
defaultSearchTopK = 10
|
||
|
|
busyTimeoutMillis = 5000
|
||
|
|
)
|
||
|
|
|
||
|
|
var collectionNamePattern = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_]*$`)
|
||
|
|
|
||
|
|
type LibSQLProvider struct {
|
||
|
|
db *sql.DB
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewLibSQLProvider(cfg *config.VectorDBConfig) (*LibSQLProvider, error) {
|
||
|
|
if cfg == nil {
|
||
|
|
return nil, fmt.Errorf("libsql vector database config is required")
|
||
|
|
}
|
||
|
|
path := strings.TrimSpace(cfg.Path)
|
||
|
|
if path == "" {
|
||
|
|
path = defaultLibSQLPath
|
||
|
|
}
|
||
|
|
absPath, err := filepath.Abs(path)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("resolve libsql vector database path: %w", err)
|
||
|
|
}
|
||
|
|
if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
|
||
|
|
return nil, fmt.Errorf("create libsql vector database directory: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
connector, err := turso.NewConnector(absPath, turso.WithBusyTimeout(busyTimeoutMillis))
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("create libsql vector database connector: %w", err)
|
||
|
|
}
|
||
|
|
db := sql.OpenDB(connector)
|
||
|
|
db.SetMaxOpenConns(1)
|
||
|
|
db.SetMaxIdleConns(1)
|
||
|
|
provider := &LibSQLProvider{db: db}
|
||
|
|
if err := provider.initialize(context.Background()); err != nil {
|
||
|
|
_ = db.Close()
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return provider, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) initialize(ctx context.Context) error {
|
||
|
|
if p == nil || p.db == nil {
|
||
|
|
return fmt.Errorf("libsql vector database is closed")
|
||
|
|
}
|
||
|
|
if err := p.db.PingContext(ctx); err != nil {
|
||
|
|
return fmt.Errorf("connect to libsql vector database: %w", err)
|
||
|
|
}
|
||
|
|
_, err := p.db.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS "_agent_vector_collections" (
|
||
|
|
name TEXT PRIMARY KEY NOT NULL,
|
||
|
|
dimension INTEGER NOT NULL,
|
||
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||
|
|
)`)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("initialize libsql collection registry: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) Close() error {
|
||
|
|
if p == nil || p.db == nil {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
err := p.db.Close()
|
||
|
|
p.db = nil
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) CreateCollection(ctx context.Context, name string, dimension int) error {
|
||
|
|
tableName, err := collectionIdentifier(name)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if dimension <= 0 || dimension > 65536 {
|
||
|
|
return fmt.Errorf("invalid libsql vector dimension: %d", dimension)
|
||
|
|
}
|
||
|
|
if info, getErr := p.GetCollection(ctx, name); getErr == nil {
|
||
|
|
if info.Dimension != dimension {
|
||
|
|
return fmt.Errorf("collection %s already uses dimension %d, requested %d", name, info.Dimension, dimension)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
} else if !errors.Is(getErr, sql.ErrNoRows) {
|
||
|
|
return getErr
|
||
|
|
}
|
||
|
|
|
||
|
|
tx, err := p.db.BeginTx(ctx, nil)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("begin libsql collection transaction: %w", err)
|
||
|
|
}
|
||
|
|
defer func() { _ = tx.Rollback() }()
|
||
|
|
|
||
|
|
createTable := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
|
||
|
|
id TEXT PRIMARY KEY NOT NULL,
|
||
|
|
embedding BLOB NOT NULL,
|
||
|
|
knowledge_base_id INTEGER NOT NULL DEFAULT 0,
|
||
|
|
document_id INTEGER NOT NULL DEFAULT 0,
|
||
|
|
document_title TEXT NOT NULL DEFAULT '',
|
||
|
|
faq_id INTEGER NOT NULL DEFAULT 0,
|
||
|
|
faq_question TEXT NOT NULL DEFAULT '',
|
||
|
|
chunk_no INTEGER NOT NULL DEFAULT 0,
|
||
|
|
chunk_type TEXT NOT NULL DEFAULT '',
|
||
|
|
section_path TEXT NOT NULL DEFAULT '',
|
||
|
|
title TEXT NOT NULL DEFAULT '',
|
||
|
|
content TEXT NOT NULL DEFAULT '',
|
||
|
|
provider TEXT NOT NULL DEFAULT ''
|
||
|
|
)`, tableName)
|
||
|
|
if _, err := tx.ExecContext(ctx, createTable); err != nil {
|
||
|
|
return fmt.Errorf("create libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
if _, err := tx.ExecContext(ctx, fmt.Sprintf(
|
||
|
|
`CREATE INDEX IF NOT EXISTS %s ON %s (knowledge_base_id, document_id)`,
|
||
|
|
quoteIdentifier(name+"_payload_idx"), tableName,
|
||
|
|
)); err != nil {
|
||
|
|
return fmt.Errorf("create libsql payload index for %s: %w", name, err)
|
||
|
|
}
|
||
|
|
if _, err := tx.ExecContext(ctx,
|
||
|
|
`INSERT INTO "_agent_vector_collections" (name, dimension) VALUES (?, ?)`, name, dimension,
|
||
|
|
); err != nil {
|
||
|
|
return fmt.Errorf("register libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
if err := tx.Commit(); err != nil {
|
||
|
|
return fmt.Errorf("commit libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) DeleteCollection(ctx context.Context, name string) error {
|
||
|
|
tableName, err := collectionIdentifier(name)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
tx, err := p.db.BeginTx(ctx, nil)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("begin libsql collection transaction: %w", err)
|
||
|
|
}
|
||
|
|
defer func() { _ = tx.Rollback() }()
|
||
|
|
if _, err := tx.ExecContext(ctx, "DROP TABLE IF EXISTS "+tableName); err != nil {
|
||
|
|
return fmt.Errorf("drop libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
if _, err := tx.ExecContext(ctx, `DELETE FROM "_agent_vector_collections" WHERE name = ?`, name); err != nil {
|
||
|
|
return fmt.Errorf("unregister libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
if err := tx.Commit(); err != nil {
|
||
|
|
return fmt.Errorf("commit libsql collection deletion %s: %w", name, err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
|
||
|
|
tableName, err := collectionIdentifier(name)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
var dimension int
|
||
|
|
if err := p.db.QueryRowContext(ctx,
|
||
|
|
`SELECT dimension FROM "_agent_vector_collections" WHERE name = ?`, name,
|
||
|
|
).Scan(&dimension); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
var count int
|
||
|
|
if err := p.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+tableName).Scan(&count); err != nil {
|
||
|
|
return nil, fmt.Errorf("count libsql collection %s: %w", name, err)
|
||
|
|
}
|
||
|
|
return &CollectionInfo{Name: name, Dimension: dimension, PointCount: count, Status: "ready"}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) ListCollections(ctx context.Context) ([]string, error) {
|
||
|
|
rows, err := p.db.QueryContext(ctx, `SELECT name FROM "_agent_vector_collections" ORDER BY name`)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("list libsql collections: %w", err)
|
||
|
|
}
|
||
|
|
defer rows.Close()
|
||
|
|
collections := make([]string, 0)
|
||
|
|
for rows.Next() {
|
||
|
|
var name string
|
||
|
|
if err := rows.Scan(&name); err != nil {
|
||
|
|
return nil, fmt.Errorf("scan libsql collection: %w", err)
|
||
|
|
}
|
||
|
|
collections = append(collections, name)
|
||
|
|
}
|
||
|
|
if err := rows.Err(); err != nil {
|
||
|
|
return nil, fmt.Errorf("iterate libsql collections: %w", err)
|
||
|
|
}
|
||
|
|
return collections, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
|
||
|
|
if len(vectors) == 0 {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
tableName, err := collectionIdentifier(collectionName)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
info, err := p.GetCollection(ctx, collectionName)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("get libsql collection %s: %w", collectionName, err)
|
||
|
|
}
|
||
|
|
|
||
|
|
tx, err := p.db.BeginTx(ctx, nil)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("begin libsql vector upsert: %w", err)
|
||
|
|
}
|
||
|
|
defer func() { _ = tx.Rollback() }()
|
||
|
|
statement := fmt.Sprintf(`INSERT INTO %s (
|
||
|
|
id, embedding, knowledge_base_id, document_id, document_title,
|
||
|
|
faq_id, faq_question, chunk_no, chunk_type, section_path, title, content, provider
|
||
|
|
) VALUES (?, vector32(?), ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
|
|
ON CONFLICT(id) DO UPDATE SET
|
||
|
|
embedding=excluded.embedding,
|
||
|
|
knowledge_base_id=excluded.knowledge_base_id,
|
||
|
|
document_id=excluded.document_id,
|
||
|
|
document_title=excluded.document_title,
|
||
|
|
faq_id=excluded.faq_id,
|
||
|
|
faq_question=excluded.faq_question,
|
||
|
|
chunk_no=excluded.chunk_no,
|
||
|
|
chunk_type=excluded.chunk_type,
|
||
|
|
section_path=excluded.section_path,
|
||
|
|
title=excluded.title,
|
||
|
|
content=excluded.content,
|
||
|
|
provider=excluded.provider`, tableName)
|
||
|
|
stmt, err := tx.PrepareContext(ctx, statement)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("prepare libsql vector upsert: %w", err)
|
||
|
|
}
|
||
|
|
defer stmt.Close()
|
||
|
|
for _, item := range vectors {
|
||
|
|
if strings.TrimSpace(item.ID) == "" {
|
||
|
|
return fmt.Errorf("libsql vector id is required")
|
||
|
|
}
|
||
|
|
if len(item.Vector) != info.Dimension {
|
||
|
|
return fmt.Errorf("invalid vector dimension for %s: got %d, want %d", item.ID, len(item.Vector), info.Dimension)
|
||
|
|
}
|
||
|
|
encoded, err := json.Marshal(item.Vector)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("encode vector %s: %w", item.ID, err)
|
||
|
|
}
|
||
|
|
payload := item.Payload
|
||
|
|
if _, err := stmt.ExecContext(ctx,
|
||
|
|
item.ID, string(encoded), payload.KnowledgeBaseID, payload.DocumentID, payload.DocumentTitle,
|
||
|
|
payload.FaqID, payload.FaqQuestion, payload.ChunkNo, payload.ChunkType,
|
||
|
|
payload.SectionPath, payload.Title, payload.Content, payload.Provider,
|
||
|
|
); err != nil {
|
||
|
|
return fmt.Errorf("upsert libsql vector %s: %w", item.ID, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := tx.Commit(); err != nil {
|
||
|
|
return fmt.Errorf("commit libsql vector upsert: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
|
||
|
|
if len(ids) == 0 {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
tableName, err := collectionIdentifier(collectionName)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
tx, err := p.db.BeginTx(ctx, nil)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("begin libsql vector deletion: %w", err)
|
||
|
|
}
|
||
|
|
defer func() { _ = tx.Rollback() }()
|
||
|
|
stmt, err := tx.PrepareContext(ctx, "DELETE FROM "+tableName+" WHERE id = ?")
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("prepare libsql vector deletion: %w", err)
|
||
|
|
}
|
||
|
|
defer stmt.Close()
|
||
|
|
for _, id := range ids {
|
||
|
|
if _, err := stmt.ExecContext(ctx, id); err != nil {
|
||
|
|
return fmt.Errorf("delete libsql vector %s: %w", id, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := tx.Commit(); err != nil {
|
||
|
|
return fmt.Errorf("commit libsql vector deletion: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *LibSQLProvider) Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
|
||
|
|
if req == nil {
|
||
|
|
return nil, fmt.Errorf("libsql search request is required")
|
||
|
|
}
|
||
|
|
tableName, err := collectionIdentifier(req.CollectionName)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
info, err := p.GetCollection(ctx, req.CollectionName)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("get libsql collection %s: %w", req.CollectionName, err)
|
||
|
|
}
|
||
|
|
if len(req.Vector) != info.Dimension {
|
||
|
|
return nil, fmt.Errorf("invalid search vector dimension: got %d, want %d", len(req.Vector), info.Dimension)
|
||
|
|
}
|
||
|
|
topK := req.TopK
|
||
|
|
if topK <= 0 {
|
||
|
|
topK = defaultSearchTopK
|
||
|
|
}
|
||
|
|
encoded, err := json.Marshal(req.Vector)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("encode search vector: %w", err)
|
||
|
|
}
|
||
|
|
vectorJSON := string(encoded)
|
||
|
|
|
||
|
|
filterSQL, filterArgs := buildSearchFilter(req.Filter)
|
||
|
|
innerWhere := filterSQL
|
||
|
|
innerArgs := []any{vectorJSON}
|
||
|
|
if filterSQL != "" {
|
||
|
|
innerArgs = append(innerArgs, filterArgs...)
|
||
|
|
}
|
||
|
|
|
||
|
|
query := fmt.Sprintf(`SELECT id, score, knowledge_base_id, document_id, document_title,
|
||
|
|
faq_id, faq_question, chunk_no, chunk_type, section_path, title, content, provider
|
||
|
|
FROM (
|
||
|
|
SELECT id, 1.0 - vector_distance_cos(embedding, vector32(?)) AS score,
|
||
|
|
knowledge_base_id, document_id, document_title, faq_id, faq_question,
|
||
|
|
chunk_no, chunk_type, section_path, title, content, provider
|
||
|
|
FROM %s%s
|
||
|
|
) ranked
|
||
|
|
WHERE score >= ?
|
||
|
|
ORDER BY score DESC
|
||
|
|
LIMIT ?`, tableName, innerWhere)
|
||
|
|
innerArgs = append(innerArgs, req.ScoreThreshold, topK)
|
||
|
|
rows, err := p.db.QueryContext(ctx, query, innerArgs...)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("search libsql collection %s: %w", req.CollectionName, err)
|
||
|
|
}
|
||
|
|
defer rows.Close()
|
||
|
|
|
||
|
|
results := make([]SearchResult, 0, topK)
|
||
|
|
for rows.Next() {
|
||
|
|
var result SearchResult
|
||
|
|
if err := rows.Scan(
|
||
|
|
&result.ID, &result.Score,
|
||
|
|
&result.Payload.KnowledgeBaseID, &result.Payload.DocumentID, &result.Payload.DocumentTitle,
|
||
|
|
&result.Payload.FaqID, &result.Payload.FaqQuestion, &result.Payload.ChunkNo,
|
||
|
|
&result.Payload.ChunkType, &result.Payload.SectionPath, &result.Payload.Title,
|
||
|
|
&result.Payload.Content, &result.Payload.Provider,
|
||
|
|
); err != nil {
|
||
|
|
return nil, fmt.Errorf("scan libsql search result: %w", err)
|
||
|
|
}
|
||
|
|
results = append(results, result)
|
||
|
|
}
|
||
|
|
if err := rows.Err(); err != nil {
|
||
|
|
return nil, fmt.Errorf("iterate libsql search results: %w", err)
|
||
|
|
}
|
||
|
|
return results, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func collectionIdentifier(name string) (string, error) {
|
||
|
|
name = strings.TrimSpace(name)
|
||
|
|
if !collectionNamePattern.MatchString(name) {
|
||
|
|
return "", fmt.Errorf("invalid libsql collection name %q", name)
|
||
|
|
}
|
||
|
|
return quoteIdentifier(name), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func quoteIdentifier(value string) string {
|
||
|
|
return `"` + value + `"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func buildSearchFilter(filter *SearchFilter) (string, []any) {
|
||
|
|
if filter == nil {
|
||
|
|
return "", nil
|
||
|
|
}
|
||
|
|
clauses := make([]string, 0, 2)
|
||
|
|
args := make([]any, 0, len(filter.KnowledgeBaseIDs)+len(filter.DocumentIDs))
|
||
|
|
if len(filter.KnowledgeBaseIDs) > 0 {
|
||
|
|
clauses = append(clauses, "knowledge_base_id IN ("+placeholders(len(filter.KnowledgeBaseIDs))+")")
|
||
|
|
for _, id := range filter.KnowledgeBaseIDs {
|
||
|
|
args = append(args, id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if len(filter.DocumentIDs) > 0 {
|
||
|
|
clauses = append(clauses, "document_id IN ("+placeholders(len(filter.DocumentIDs))+")")
|
||
|
|
for _, id := range filter.DocumentIDs {
|
||
|
|
args = append(args, id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if len(clauses) == 0 {
|
||
|
|
return "", nil
|
||
|
|
}
|
||
|
|
return " WHERE " + strings.Join(clauses, " AND "), args
|
||
|
|
}
|
||
|
|
|
||
|
|
func placeholders(count int) string {
|
||
|
|
values := make([]string, count)
|
||
|
|
for i := range values {
|
||
|
|
values[i] = "?"
|
||
|
|
}
|
||
|
|
return strings.Join(values, ",")
|
||
|
|
}
|