fix(search): add/remove chunks to fts table

Author: Daniele Briggi

src/sqlite_rag/cli.py

@@ -34,7 +34,7 @@ def add(
 ):
     """Add a file path to the database"""
     rag = SQLiteRag()
-    rag.add(path, recursively=recursive)
+    rag.add(path, recursive=recursive)
 
 
 @app.command()
@@ -56,8 +56,6 @@ def list_documents():
 
     # Print stats
     typer.echo(f"Total documents: {len(documents)}")
-    typer.echo("Documents:")
-    typer.echo("-" * 106)
 
     typer.echo(f"{'ID':<36} {'URI/Content':<50} {'Created At':<20}")
     typer.echo("-" * 106)
@@ -76,32 +74,36 @@ def list_documents():
 @app.command()
 def remove(
     identifier: str,
-    yes: bool = typer.Option(False, "-y", "--yes", help="Skip confirmation prompt")
+    yes: bool = typer.Option(False, "-y", "--yes", help="Skip confirmation prompt"),
 ):
     """Remove document by path or UUID"""
     rag = SQLiteRag()
-    
+
     # Find the document first
     document = rag.find_document(identifier)
     if not document:
         typer.echo(f"Document not found: {identifier}")
         raise typer.Exit(1)
-    
+
     # Show document details
     typer.echo(f"Found document:")
     typer.echo(f"ID: {document.id}")
     typer.echo(f"URI: {document.uri or 'N/A'}")
-    typer.echo(f"Created: {document.created_at.strftime('%Y-%m-%d %H:%M:%S') if document.created_at else 'N/A'}")
-    typer.echo(f"Content preview: {document.content[:200]}{'...' if len(document.content) > 200 else ''}")
+    typer.echo(
+        f"Created: {document.created_at.strftime('%Y-%m-%d %H:%M:%S') if document.created_at else 'N/A'}"
+    )
+    typer.echo(
+        f"Content preview: {document.content[:200]}{'...' if len(document.content) > 200 else ''}"
+    )
     typer.echo()
-    
+
     # Ask for confirmation unless -y flag is used
     if not yes:
         confirm = typer.confirm("Are you sure you want to delete this document?")
         if not confirm:
             typer.echo("Cancelled.")
             return
-    
+
     # Remove the document
     success = rag.remove_document(identifier)
     if success:
@@ -128,7 +130,15 @@ def search(
     query: str, limit: int = typer.Option(10, help="Number of results to return")
 ):
     """Search for documents using hybrid vector + full-text search"""
-    pass
+    rag = SQLiteRag()
+    results = rag.search(query, top_k=limit)
+
+    if not results:
+        typer.echo("No documents found matching the query.")
+        return
+
+    typer.echo(f"Found {len(results)} documents:")
+    
 
 
 def repl_mode():

src/sqlite_rag/database.py

@@ -12,7 +12,9 @@ def initialize(conn: sqlite3.Connection, settings: Settings) -> sqlite3.Connecti
         """Initialize the database with extensions and schema"""
         conn.enable_load_extension(True)
         try:
-            conn.load_extension(str(Path(__file__).parent.parent.parent / "extensions" / "ai"))
+            conn.load_extension(
+                str(Path(__file__).parent.parent.parent / "extensions" / "ai")
+            )
             conn.load_extension(
                 str(Path(__file__).parent.parent.parent / "extensions" / "vector")
             )
@@ -61,7 +63,6 @@ def _create_schema(conn: sqlite3.Connection, settings: Settings):
         cursor.execute(
             """
             CREATE TABLE IF NOT EXISTS chunks (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
                 document_id TEXT,
                 content TEXT,
                 embedding BLOB,
@@ -72,7 +73,7 @@ def _create_schema(conn: sqlite3.Connection, settings: Settings):
 
         cursor.execute(
             """
-            CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(content, content='chunks', content_rowid='id');
+            CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(content, content='chunks', content_rowid='rowid');
         """
         )
 

src/sqlite_rag/engine.py

@@ -81,7 +81,7 @@ def search(self, query: str, limit: int = 10) -> list[Document]:
                     v.distance
                 FROM chunks AS c
                     JOIN vector_quantize_scan('chunks', 'embedding', vector_convert_f32(:query_embedding), :k) AS v
-                    ON c.id = v.rowid
+                    ON c.rowid = v.rowid
             ),
             -- Full-text search results
             fts_matches AS (
@@ -90,7 +90,7 @@ def search(self, query: str, limit: int = 10) -> list[Document]:
                     c.document_id,
                     rank AS score
                 FROM chunks_fts AS c_fts
-                    JOIN chunks AS c ON c_fts.rowid = c.id
+                    JOIN chunks AS c ON c_fts.rowid = c.rowid
                 WHERE c_fts.content MATCH :query
                 LIMIT :k
             ),
@@ -127,7 +127,7 @@ def search(self, query: str, limit: int = 10) -> list[Document]:
             FROM matches;
             """,
             {
-                "query": f'"{query}"',
+                "query": query + "*",
                 "query_embedding": query_embedding,
                 "k": limit,
                 # TODO: move to settings or costants

src/sqlite_rag/repository.py

@@ -33,6 +33,10 @@ def add_document(self, document: Document) -> str:
                 "INSERT INTO chunks (document_id, content, embedding) VALUES (?, ?, vector_convert_f32(?))",
                 (document_id, chunk.content, chunk.embedding),
             )
+            cursor.execute(
+                "INSERT INTO chunks_fts (rowid, content) VALUES (last_insert_rowid(), ?)",
+                (chunk.content,),
+            )
 
         self._conn.commit()
 
@@ -62,11 +66,11 @@ def find_document_by_id_or_uri(self, identifier: str) -> Document | None:
         """Find document by ID or URI"""
         cursor = self._conn.cursor()
         cursor.execute(
-            "SELECT id, content, uri, metadata, created_at FROM documents WHERE id = ? OR uri = ?", 
-            (identifier, identifier)
+            "SELECT id, content, uri, metadata, created_at FROM documents WHERE id = ? OR uri = ?",
+            (identifier, identifier),
         )
         row = cursor.fetchone()
-        
+
         if row:
             doc_id, content, uri, metadata, created_at = row
             return Document(
@@ -81,17 +85,21 @@ def find_document_by_id_or_uri(self, identifier: str) -> Document | None:
     def remove_document(self, document_id: str) -> bool:
         """Remove document and its chunks by document ID"""
         cursor = self._conn.cursor()
-        
+
         # Check if document exists
         cursor.execute("SELECT COUNT(*) FROM documents WHERE id = ?", (document_id,))
         if cursor.fetchone()[0] == 0:
             return False
-        
-        # Remove chunks first (foreign key constraint)
+
+        # Remove chunks first
+        cursor.execute(
+            "DELETE FROM chunks_fts WHERE rowid IN (SELECT rowid FROM chunks WHERE document_id = ?)",
+            (document_id,),
+        )
         cursor.execute("DELETE FROM chunks WHERE document_id = ?", (document_id,))
-        
+
         # Remove document
         cursor.execute("DELETE FROM documents WHERE id = ?", (document_id,))
-        
+
         self._conn.commit()
         return True

src/sqlite_rag/sqliterag.py

@@ -19,7 +19,7 @@ def __init__(self, settings: Optional[Settings] = None):
         if settings is None:
             # TODO: load defaults or from the database
             settings = Settings(
-                model_path_or_name="./all-MiniLM-L6-v2.e4ce9877.q8_0.gguf",
+                model_path_or_name="./Qwen3-Embedding-0.6B-Q8_0.gguf",
                 db_path="sqliterag.db",
             )
 
@@ -102,8 +102,10 @@ def remove_document(self, identifier: str) -> bool:
 
         return self._repository.remove_document(document.id or "")
 
-    # def search(
-    #     self, query: str, top_k: int = 10
-    # ) -> list[Document]:
-    #     """Search for documents matching the query"""
-    #     self._ensure_initialized()
+    def search(
+        self, query: str, top_k: int = 10
+    ) -> list[Document]:
+        """Search for documents matching the query"""
+        self._ensure_initialized()
+
+        return self._engine.search(query, limit=top_k)

tests/conftest.py

@@ -11,7 +11,7 @@
 def db_conn():
     with tempfile.NamedTemporaryFile(suffix=".db") as tmp_db:
         settings = Settings(
-            model_path_or_name="./all-MiniLM-L6-v2.e4ce9877.q8_0.gguf",
+            model_path_or_name="./Qwen3-Embedding-0.6B-Q8_0.gguf",
             db_path=tmp_db.name,
         )
 
@@ -27,7 +27,7 @@ def db_conn():
 def db_settings() -> Settings:
     with tempfile.NamedTemporaryFile(suffix=".db") as tmp_db:
         settings = Settings(
-            model_path_or_name="./all-MiniLM-L6-v2.e4ce9877.q8_0.gguf",
+            model_path_or_name="./Qwen3-Embedding-0.6B-Q8_0.gguf",
             db_path=tmp_db.name,
         )
     return settings

tests/test_engine.py

@@ -139,7 +139,7 @@ def test_search_fts_results(self, db_conn):
         engine.quantize()
 
         # Act
-        results = engine.search("brown fox animal", limit=5)
+        results = engine.search("quick brown fox", limit=5)
 
         assert len(results) > 0
         assert doc1_id == results[0].id