feat(rag): pgvector knowledge base with API and LLM tools
- services/rag: knowledge_docs store — vector search via the shared embedder (HNSW index, graceful ILIKE fallback), list/delete scoped to the user, search intentionally global for the single-user deployment - gateway: GET/POST /knowledge/docs, DELETE /knowledge/docs/:id, POST /knowledge/search - plugins/knowledge: knowledge_search and knowledge_ingest tools so the persona can save and recall information mid-conversation - main: migrations, store wiring, plugin registration
This commit is contained in:
1 parent
ac0c2b2c73
commit
9c8d5a9e8d
6 files changed
+432
-2
No files matched your search
@@ -0,0 +1,181 @@
|
||||
// Package rag provides a minimal pgvector-backed knowledge base.
|
||||
package rag
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/services/memory"
|
||||
)
|
||||
|
||||
// Document is a knowledge base entry.
|
||||
type Document struct {
|
||||
ID int64 `json:"id"`
|
||||
UserID string `json:"user_id"`
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// Embedder is the subset of memory.Embedder used here.
|
||||
type Embedder interface {
|
||||
Embed(ctx context.Context, text string) ([]float64, error)
|
||||
Dims() int
|
||||
}
|
||||
|
||||
// Store persists and searches documents.
|
||||
type Store struct {
|
||||
db *sql.DB
|
||||
embedder Embedder
|
||||
vectorEnabled bool
|
||||
}
|
||||
|
||||
// NewStore creates a knowledge store. Embedder may be nil (ILIKE-only search).
|
||||
func NewStore(db *sql.DB, embedder Embedder) *Store {
|
||||
return &Store{db: db, embedder: embedder}
|
||||
}
|
||||
|
||||
// EnableVector marks the embedding column as available (after migration).
|
||||
func (s *Store) EnableVector() { s.vectorEnabled = true }
|
||||
|
||||
// CreateTableSQL returns the base schema.
|
||||
func CreateTableSQL() string {
|
||||
return `
|
||||
CREATE TABLE IF NOT EXISTS knowledge_docs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
content TEXT NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_docs_user ON knowledge_docs(user_id, id DESC);
|
||||
`
|
||||
}
|
||||
|
||||
// VectorMigrationSQL adds the embedding column (pgvector required).
|
||||
func VectorMigrationSQL(dims int) string {
|
||||
return fmt.Sprintf(`
|
||||
ALTER TABLE knowledge_docs ADD COLUMN IF NOT EXISTS embedding vector(%d);
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_docs_embedding ON knowledge_docs USING hnsw (embedding vector_cosine_ops);
|
||||
`, dims)
|
||||
}
|
||||
|
||||
// Ingest stores a document, embedding it when an embedder is available.
|
||||
func (s *Store) Ingest(ctx context.Context, userID, title, content string) (*Document, error) {
|
||||
doc := &Document{UserID: userID, Title: title, Content: content}
|
||||
if s.embedder != nil && s.vectorEnabled {
|
||||
if vec, err := s.embedder.Embed(ctx, title+"\n"+content); err == nil {
|
||||
_, err := s.db.ExecContext(ctx, `
|
||||
INSERT INTO knowledge_docs (user_id, title, content, embedding) VALUES ($1, $2, $3, $4::vector)
|
||||
`, userID, title, content, vectorString(vec))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("ingest: %w", err)
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
}
|
||||
err := s.db.QueryRowContext(ctx, `
|
||||
INSERT INTO knowledge_docs (user_id, title, content) VALUES ($1, $2, $3) RETURNING id, created_at
|
||||
`, userID, title, content).Scan(&doc.ID, &doc.CreatedAt)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("ingest: %w", err)
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
// Search finds documents by vector similarity when enabled, else ILIKE.
|
||||
// The knowledge base is shared: userID is recorded on ingest but search
|
||||
// is global (single-user self-hosted deployment).
|
||||
func (s *Store) Search(ctx context.Context, userID, query string, topK int) ([]Document, error) {
|
||||
if topK <= 0 || topK > 20 {
|
||||
topK = 5
|
||||
}
|
||||
if s.embedder != nil && s.vectorEnabled {
|
||||
vec, err := s.embedder.Embed(ctx, query)
|
||||
if err == nil {
|
||||
return s.queryDocs(ctx, `
|
||||
SELECT id, user_id, title, content, created_at FROM knowledge_docs
|
||||
WHERE embedding IS NOT NULL
|
||||
ORDER BY embedding <=> $1::vector LIMIT $2
|
||||
`, vectorString(vec), topK)
|
||||
}
|
||||
}
|
||||
return s.queryDocs(ctx, `
|
||||
SELECT id, user_id, title, content, created_at FROM knowledge_docs
|
||||
WHERE title ILIKE '%' || $1 || '%' OR content ILIKE '%' || $1 || '%'
|
||||
ORDER BY id DESC LIMIT $2
|
||||
`, query, topK)
|
||||
}
|
||||
|
||||
// List returns recent documents (without content payload).
|
||||
func (s *Store) List(ctx context.Context, userID string, limit int) ([]Document, error) {
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
rows, err := s.db.QueryContext(ctx, `
|
||||
SELECT id, user_id, title, created_at FROM knowledge_docs
|
||||
WHERE user_id = $1 ORDER BY id DESC LIMIT $2
|
||||
`, userID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
docs := []Document{}
|
||||
for rows.Next() {
|
||||
var d Document
|
||||
if err := rows.Scan(&d.ID, &d.UserID, &d.Title, &d.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
docs = append(docs, d)
|
||||
}
|
||||
return docs, rows.Err()
|
||||
}
|
||||
|
||||
// Delete removes a document owned by the user.
|
||||
func (s *Store) Delete(ctx context.Context, userID string, id int64) error {
|
||||
res, err := s.db.ExecContext(ctx, `DELETE FROM knowledge_docs WHERE id = $1 AND user_id = $2`, id, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("document not found")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) queryDocs(ctx context.Context, query string, args ...any) ([]Document, error) {
|
||||
rows, err := s.db.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
docs := []Document{}
|
||||
for rows.Next() {
|
||||
var d Document
|
||||
if err := rows.Scan(&d.ID, &d.UserID, &d.Title, &d.Content, &d.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
docs = append(docs, d)
|
||||
}
|
||||
return docs, rows.Err()
|
||||
}
|
||||
|
||||
// vectorString formats an embedding for pgvector input.
|
||||
func vectorString(vec []float64) string {
|
||||
var sb strings.Builder
|
||||
sb.WriteByte('[')
|
||||
for i, v := range vec {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
fmt.Fprintf(&sb, "%g", v)
|
||||
}
|
||||
sb.WriteByte(']')
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// Ensure the memory embedder satisfies the interface at compile time.
|
||||
var _ Embedder = (*memory.Embedder)(nil)
|
||||
Reference in new issue
Block a user