diff options
| -rw-r--r-- | rag/rag.go | 4 | ||||
| -rw-r--r-- | rag/storage.go | 90 |
2 files changed, 55 insertions, 39 deletions
@@ -163,9 +163,9 @@ func (r *RAG) writeVectors(vectorCh chan []models.VectorRow) error { for batch := range vectorCh { for _, vector := range batch { if err := r.storage.WriteVector(&vector); err != nil { - r.logger.Error("failed to write vector", "error", err, "slug", vector.Slug) + r.logger.Error("failed to write vector to DB", "error", err, "slug", vector.Slug) LongJobStatusCh <- ErrRAGStatus - continue // a duplicate is not critical + return err // Stop the entire RAG operation on DB error } } r.logger.Debug("wrote batch to db", "size", len(batch), "vector_chan_len", len(vectorCh)) diff --git a/rag/storage.go b/rag/storage.go index 64d54f7..bb3a24a 100644 --- a/rag/storage.go +++ b/rag/storage.go @@ -30,33 +30,34 @@ func NewVectorStorage(logger *slog.Logger, store storage.FullRepo) *VectorStorag // CreateTables creates the necessary tables for vector storage func (vs *VectorStorage) CreateTables() error { - // Create tables for different embedding dimensions - queries := []string{ - `CREATE TABLE IF NOT EXISTS embeddings_384 ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - embeddings BLOB NOT NULL, - slug TEXT NOT NULL, - raw_text TEXT NOT NULL, - filename TEXT NOT NULL, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP - )`, - `CREATE TABLE IF NOT EXISTS embeddings_5120 ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - embeddings BLOB NOT NULL, - slug TEXT NOT NULL, - raw_text TEXT NOT NULL, - filename TEXT NOT NULL, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP - )`, - // Indexes for better performance - `CREATE INDEX IF NOT EXISTS idx_embeddings_384_filename ON embeddings_384(filename)`, - `CREATE INDEX IF NOT EXISTS idx_embeddings_5120_filename ON embeddings_5120(filename)`, - `CREATE INDEX IF NOT EXISTS idx_embeddings_384_slug ON embeddings_384(slug)`, - `CREATE INDEX IF NOT EXISTS idx_embeddings_5120_slug ON embeddings_5120(slug)`, - - // Additional indexes that may help with searches - `CREATE INDEX IF NOT EXISTS idx_embeddings_384_created_at ON embeddings_384(created_at)`, - `CREATE INDEX IF NOT EXISTS idx_embeddings_5120_created_at ON embeddings_5120(created_at)`, + // Create tables for common embedding dimensions + embeddingSizes := []int{384, 768, 1024, 1536, 2048, 3072, 4096, 5120} + // Pre-allocate queries slice: each embedding size needs 1 table + 3 indexes = 4 queries per size + queries := make([]string, 0, len(embeddingSizes)*4) + + // Generate table creation queries for each embedding size + for _, size := range embeddingSizes { + tableName := fmt.Sprintf("embeddings_%d", size) + queries = append(queries, + fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + embeddings BLOB NOT NULL, + slug TEXT NOT NULL, + raw_text TEXT NOT NULL, + filename TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + )`, tableName), + ) + } + + // Add indexes for all supported sizes + for _, size := range embeddingSizes { + tableName := fmt.Sprintf("embeddings_%d", size) + queries = append(queries, + fmt.Sprintf(`CREATE INDEX IF NOT EXISTS idx_%s_filename ON %s(filename)`, tableName, tableName), + fmt.Sprintf(`CREATE INDEX IF NOT EXISTS idx_%s_slug ON %s(slug)`, tableName, tableName), + fmt.Sprintf(`CREATE INDEX IF NOT EXISTS idx_%s_created_at ON %s(created_at)`, tableName, tableName), + ) } for _, query := range queries { @@ -120,14 +121,25 @@ func (vs *VectorStorage) WriteVector(row *models.VectorRow) error { // getTableName determines which table to use based on embedding size func (vs *VectorStorage) getTableName(emb []float32) (string, error) { - switch len(emb) { - case 384: - return "embeddings_384", nil - case 5120: - return "embeddings_5120", nil - default: - return "", fmt.Errorf("no table for embedding size of %d", len(emb)) + size := len(emb) + + // Check if we support this embedding size + supportedSizes := map[int]bool{ + 384: true, + 768: true, + 1024: true, + 1536: true, + 2048: true, + 3072: true, + 4096: true, + 5120: true, } + + if supportedSizes[size] { + return fmt.Sprintf("embeddings_%d", size), nil + } + + return "", fmt.Errorf("no table for embedding size of %d", size) } // SearchClosest finds vectors closest to the query vector using efficient cosine similarity calculation @@ -211,8 +223,10 @@ func (vs *VectorStorage) SearchClosest(query []float32) ([]models.VectorRow, err func (vs *VectorStorage) ListFiles() ([]string, error) { fileLists := make([][]string, 0) - // Query both tables and combine results - for _, table := range []string{"embeddings_384", "embeddings_5120"} { + // Query all supported tables and combine results + embeddingSizes := []int{384, 768, 1024, 1536, 2048, 3072, 4096, 5120} + for _, size := range embeddingSizes { + table := fmt.Sprintf("embeddings_%d", size) query := "SELECT DISTINCT filename FROM " + table rows, err := vs.sqlxDB.Query(query) if err != nil { @@ -252,7 +266,9 @@ func (vs *VectorStorage) ListFiles() ([]string, error) { func (vs *VectorStorage) RemoveEmbByFileName(filename string) error { var errors []string - for _, table := range []string{"embeddings_384", "embeddings_5120"} { + embeddingSizes := []int{384, 768, 1024, 1536, 2048, 3072, 4096, 5120} + for _, size := range embeddingSizes { + table := fmt.Sprintf("embeddings_%d", size) query := fmt.Sprintf("DELETE FROM %s WHERE filename = ?", table) if _, err := vs.sqlxDB.Exec(query, filename); err != nil { errors = append(errors, err.Error()) |