Merge pull request #953 from krissetto/cleanup-rag

Cleanup RAG code a bit

Author: dgageot

examples/rag.yaml

@@ -25,6 +25,6 @@ rag:
       - ./rag/blork_field_guide.txt
     strategies:
       - type: chunked-embeddings
-        model: openai/text-embedding-3-small
+        embedding_model: openai/text-embedding-3-small
         database: ./rag/chunked_embeddings.db
         vector_dimensions: 1536

examples/rag/hybrid.yaml

@@ -25,7 +25,7 @@ rag:
     strategies:
       # Chunked embeddings strategy for semantic search
       - type: chunked-embeddings
-        model: openai/text-embedding-3-small
+        embedding_model: openai/text-embedding-3-small
         docs:
           - ./docs                            # Additional docs for chunked-embeddings strategy only
         database: ./chunked_embeddings.db     # chunked-embeddings database to use; path to local sqlite database

examples/rag/reranking.yaml

@@ -31,7 +31,7 @@ rag:
 
     strategies:
       - type: chunked-embeddings
-        model: openai-embedder
+        embedding_model: openai-embedder
         database: ./chunked_embeddings.db
         similarity_metric: cosine_similarity
         chunking:

pkg/config/latest/types.go

@@ -326,7 +326,7 @@ type RAGStrategyConfig struct {
 
 	// Strategy-specific parameters (arbitrary key-value pairs)
 	// Examples:
-	// - chunked-embeddings: model, similarity_metric, threshold, vector_dimensions
+	// - chunked-embeddings: embedding_model, similarity_metric, threshold, vector_dimensions
 	// - bm25: k1, b, threshold
 	Params map[string]any // Flattened into parent JSON
 }

pkg/rag/database/database.go

@@ -1,54 +1,37 @@
 package database
 
 import (
-	"context"
 	"math"
 )
 
-// Document represents a chunk of text with its embedding
+// Document represents a chunk of text - the base type returned by all RAG strategies.
+// Strategy-specific fields (embeddings, semantic summaries) are handled internally
+// by each strategy and don't need to be exposed here.
 type Document struct {
-	ID         string    `json:"id"`
-	SourcePath string    `json:"source_path"`
-	ChunkIndex int       `json:"chunk_index"`
-	Content    string    `json:"content"`
-	Embedding  []float64 `json:"-"`
-	FileHash   string    `json:"file_hash"`
-	CreatedAt  string    `json:"created_at"`
+	ID         string `json:"id"`
+	SourcePath string `json:"source_path"`
+	ChunkIndex int    `json:"chunk_index"`
+	Content    string `json:"content"`
+	FileHash   string `json:"file_hash"`
+	CreatedAt  string `json:"created_at"`
 }
 
-// SearchResult represents a document with its similarity score
+// SearchResult represents a document with its relevance score.
+// This is the common return type for all Strategy.Query() implementations.
 type SearchResult struct {
 	Document   Document `json:"document"`
 	Similarity float64  `json:"similarity"`
 }
 
-// FileMetadata represents metadata about an indexed file
+// FileMetadata represents metadata about an indexed file.
+// Used for change detection and incremental indexing.
 type FileMetadata struct {
 	SourcePath  string
 	FileHash    string
 	LastIndexed string
 	ChunkCount  int
 }
 
-// Database interface for RAG operations
-// Implementations: SQLite (sqlite.go), PostgreSQL (future), Pinecone (future), etc.
-type Database interface {
-	// Document operations
-	AddDocument(ctx context.Context, doc Document) error
-	DeleteDocumentsByPath(ctx context.Context, sourcePath string) error
-	SearchSimilar(ctx context.Context, queryEmbedding []float64, limit int) ([]SearchResult, error)
-	GetDocumentsByPath(ctx context.Context, sourcePath string) ([]Document, error)
-
-	// File metadata operations (for change detection and incremental indexing)
-	GetFileMetadata(ctx context.Context, sourcePath string) (*FileMetadata, error)
-	SetFileMetadata(ctx context.Context, metadata FileMetadata) error
-	GetAllFileMetadata(ctx context.Context) ([]FileMetadata, error)
-	DeleteFileMetadata(ctx context.Context, sourcePath string) error
-
-	// Resource management
-	Close() error
-}
-
 // Helper functions
 
 // CosineSimilarity calculates cosine similarity between two vectors

pkg/rag/strategy/bm25.go

@@ -44,8 +44,9 @@ func NewBM25FromConfig(ctx context.Context, cfg latest.RAGStrategyConfig, buildC
 		return nil, fmt.Errorf("invalid database config: %w", err)
 	}
 
-	// Create BM25-specific database (no vectors needed)
-	db, err := NewBM25Database(dbPath)
+	// Create BM25-specific database (no vectors needed).
+	// Pass strategy type as table prefix so multiple strategies can share the same DB file.
+	db, err := newBM25DB(dbPath, cfg.Type)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create database: %w", err)
 	}
@@ -66,14 +67,14 @@ func NewBM25FromConfig(ctx context.Context, cfg latest.RAGStrategyConfig, buildC
 		"chunk_overlap", chunkOverlap,
 		"respect_word_boundaries", respectWordBoundaries)
 	if chunkSize == 0 {
-		chunkSize = 1000
+		chunkSize = 1500 // General text: good paragraph/section size
 	}
 	if chunkOverlap == 0 {
 		chunkOverlap = 75
 	}
 
 	// Create strategy
-	strategy := NewBM25Strategy(
+	strategy := newBM25Strategy(
 		"bm25",
 		db,
 		events,
@@ -97,7 +98,7 @@ func NewBM25FromConfig(ctx context.Context, cfg latest.RAGStrategyConfig, buildC
 // BM25 is a ranking function that uses term frequency and inverse document frequency
 type BM25Strategy struct {
 	name       string
-	db         database.Database
+	db         *bm25DB
 	processor  *chunk.Processor
 	fileHashes map[string]string
 	watcher    *fsnotify.Watcher
@@ -111,8 +112,8 @@ type BM25Strategy struct {
 	docCount     int     // total number of documents
 }
 
-// NewBM25Strategy creates a new BM25-based retrieval strategy
-func NewBM25Strategy(name string, db database.Database, events chan<- types.Event, k1, b float64) *BM25Strategy {
+// newBM25Strategy creates a new BM25-based retrieval strategy
+func newBM25Strategy(name string, db *bm25DB, events chan<- types.Event, k1, b float64) *BM25Strategy {
 	return &BM25Strategy{
 		name:       name,
 		db:         db,
@@ -460,19 +461,10 @@ func (s *BM25Strategy) calculateBM25Score(queryTerms []string, doc database.Docu
 }
 
 func (s *BM25Strategy) getAllDocuments(ctx context.Context) ([]database.Document, error) {
-	// This is a placeholder - you'd need to add a method to the database interface
-	// For now, we'll use SearchSimilar with an empty embedding to get all docs
-	// In production, add a proper GetAllDocuments method to the database interface
-	results, err := s.db.SearchSimilar(ctx, []float64{}, 10000)
+	docs, err := s.db.GetAllDocuments(ctx)
 	if err != nil {
 		return nil, err
 	}
-
-	docs := make([]database.Document, len(results))
-	for i, result := range results {
-		docs[i] = result.Document
-	}
-
 	s.docCount = len(docs)
 	return docs, nil
 }
@@ -550,13 +542,11 @@ func (s *BM25Strategy) indexFile(ctx context.Context, filePath string, chunkSize
 			continue
 		}
 
-		// For BM25, we don't need embeddings, but we still store the document
 		doc := database.Document{
 			ID:         fmt.Sprintf("%s_%d_%d", filePath, chunk.Index, time.Now().UnixNano()),
 			SourcePath: filePath,
 			ChunkIndex: chunk.Index,
 			Content:    chunk.Content,
-			Embedding:  []float64{}, // Empty embedding for BM25
 			FileHash:   fileHash,
 		}