Cyrene_v2/backend/services/rag/rag.go

// 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)