chore: tests

Author: Daniele Briggi

src/sqlite_rag/cli.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 import shlex
 import sys
-from pathlib import Path
 from typing import Optional
 
 import typer
@@ -75,7 +74,10 @@ def list_documents():
 
 
 @app.command()
-def remove(identifier: str):
+def remove(
+    identifier: str,
+    yes: bool = typer.Option(False, "-y", "--yes", help="Skip confirmation prompt")
+):
     """Remove document by path or UUID"""
     rag = SQLiteRag()
 
@@ -93,11 +95,12 @@ def remove(identifier: str):
     typer.echo(f"Content preview: {document.content[:200]}{'...' if len(document.content) > 200 else ''}")
     typer.echo()
 
-    # Ask for confirmation
-    confirm = typer.confirm("Are you sure you want to delete this document?")
-    if not confirm:
-        typer.echo("Cancelled.")
-        return
+    # 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)

src/sqlite_rag/engine.py

@@ -46,7 +46,7 @@ def generate_embedding(self, chunks: list[Chunk]) -> list[Chunk]:
         cursor = self._conn.cursor()
 
         for chunk in chunks:
-            cursor.execute("SELECT llm_embed_generate(?, 'json_output=0')", (chunk.content,))
+            cursor.execute("SELECT llm_embed_generate(?)", (chunk.content,))
             result = cursor.fetchone()
 
             if result is None:

tests/test_engine.py

@@ -49,7 +49,7 @@ def test_search_with_semantic_and_fts(self, db_conn):
             uri="document2.txt",
         )
         doc3 = Document(
-            content="This document contains the phrase 'tongue-twister' and discusses woodcutters and wood.",
+            content="This document discusses about woodcutters and wood.",
             uri="document3.txt",
         )
 
@@ -59,16 +59,16 @@ def test_search_with_semantic_and_fts(self, db_conn):
 
         repository = Repository(conn, settings)
         repository.add_document(doc1)
-        doc2_id = repository.add_document(doc2)
-        repository.add_document(doc3)
+        repository.add_document(doc2)
+        doc3_id = repository.add_document(doc3)
 
         engine.quantize()
 
         # Act
-        results = engine.search("wood woodchuck", limit=5)
+        results = engine.search("wood lumberjack", limit=5)
 
         assert len(results) > 0
-        assert doc2_id == results[0].id
+        assert doc3_id == results[0].id
 
     def test_search_semantic_result(self, db_conn):
         # Arrange
@@ -86,7 +86,7 @@ def test_search_semantic_result(self, db_conn):
             uri="document2.txt",
         )
         doc3 = Document(
-            content="This document contains the phrase 'tongue-twister' and discusses woodcutters and wood.",
+            content="This document discusses about woodcutters and wood.",
             uri="document3.txt",
         )
 
@@ -96,16 +96,16 @@ def test_search_semantic_result(self, db_conn):
 
         repository = Repository(conn, settings)
         repository.add_document(doc1)
-        doc2_id = repository.add_document(doc2)
-        repository.add_document(doc3)
+        repository.add_document(doc2)
+        doc3_id = repository.add_document(doc3)
 
         engine.quantize()
 
         # Act
-        results = engine.search("tongue-twister", limit=5)
+        results = engine.search("lumberjack", limit=5)
 
         assert len(results) > 0
-        assert doc2_id == results[0].id
+        assert doc3_id == results[0].id
 
     def test_search_fts_results(self, db_conn):
         # Arrange
@@ -123,7 +123,7 @@ def test_search_fts_results(self, db_conn):
             uri="document2.txt",
         )
         doc3 = Document(
-            content="This document contains the phrase 'tongue-twister' and discusses woodcutters and wood.",
+            content="This document discusses about woodcutters and wood.",
             uri="document3.txt",
         )
 
@@ -132,14 +132,14 @@ def test_search_fts_results(self, db_conn):
         engine.process(doc3)
 
         repository = Repository(conn, settings)
-        repository.add_document(doc1)
+        doc1_id = repository.add_document(doc1)
         repository.add_document(doc2)
-        doc3_id = repository.add_document(doc3)
+        repository.add_document(doc3)
 
         engine.quantize()
 
         # Act
-        results = engine.search("woodcutters", limit=5)
+        results = engine.search("brown fox animal", limit=5)
 
         assert len(results) > 0
-        assert doc3_id == results[0].id
+        assert doc1_id == results[0].id

tests/test_sqlite_rag.py

@@ -13,9 +13,6 @@ def test_add_simple_text_file(self, db_settings):
             )
             temp_file_path = f.name
 
-        db_settings.chunk_size = 2
-        db_settings.chunk_overlap = 0
-
         rag = SQLiteRag(db_settings)
 
         rag.add(temp_file_path)
@@ -102,15 +99,17 @@ def test_list_documents(self, db_settings):
 
     def test_find_document_by_id(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
-        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+
+        rag.add_text(
+            "Test document content.", uri="test.txt", metadata={"author": "test"}
+        )
         documents = rag.list_documents()
         doc_id = documents[0].id
-        
+
         # Find by ID
         assert doc_id is not None
         found_doc = rag.find_document(doc_id)
-        
+
         assert found_doc is not None
         assert found_doc.id == doc_id
         assert found_doc.content == "Test document content."
@@ -119,95 +118,104 @@ def test_find_document_by_id(self, db_settings):
 
     def test_find_document_by_uri(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
-        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
-        
+
+        rag.add_text(
+            "Test document content.", uri="test.txt", metadata={"author": "test"}
+        )
+
         # Find by URI
         found_doc = rag.find_document("test.txt")
-        
+
         assert found_doc is not None
         assert found_doc.content == "Test document content."
         assert found_doc.uri == "test.txt"
         assert found_doc.metadata == {"author": "test"}
 
     def test_find_document_not_found(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
+
         found_doc = rag.find_document("nonexistent")
-        
+
         assert found_doc is None
 
     def test_remove_document_by_id(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
-        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
+
+        rag.add_text(
+            "Test document content.", uri="test.txt", metadata={"author": "test"}
+        )
         documents = rag.list_documents()
         doc_id = documents[0].id
-        
+
         # Verify document exists
         assert len(documents) == 1
-        
+
         # Remove by ID
         assert doc_id is not None
         success = rag.remove_document(doc_id)
-        
+
         assert success is True
-        
+
         # Verify document is removed
         documents = rag.list_documents()
         assert len(documents) == 0
 
     def test_remove_document_by_uri(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
-        rag.add_text("Test document content.", uri="test.txt", metadata={"author": "test"})
-        
+
+        rag.add_text(
+            "Test document content.", uri="test.txt", metadata={"author": "test"}
+        )
+
         # Verify document exists
         documents = rag.list_documents()
         assert len(documents) == 1
-        
+
         # Remove by URI
         success = rag.remove_document("test.txt")
-        
+
         assert success is True
-        
+
         # Verify document is removed
         documents = rag.list_documents()
         assert len(documents) == 0
 
     def test_remove_document_not_found(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
+
         success = rag.remove_document("nonexistent")
-        
+
         assert success is False
 
     def test_remove_document_with_chunks(self, db_settings):
         rag = SQLiteRag(db_settings)
-        
+
         # Add document that will create chunks
-        rag.add_text("This is a longer document that should create multiple chunks when processed by the chunker.", uri="test.txt")
-        
+        rag.add_text(
+            "This is a longer document that should create multiple chunks when processed by the chunker.",
+            uri="test.txt",
+        )
+
         # Verify document and chunks exist
         documents = rag.list_documents()
         assert len(documents) == 1
         doc_id = documents[0].id
-        
+
         cursor = rag._conn.cursor()
         cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
         chunk_count = cursor.fetchone()[0]
         assert chunk_count > 0
-        
+
         # Remove document
         assert doc_id is not None
         success = rag.remove_document(doc_id)
-        
+
         assert success is True
-        
+
         # Verify document and chunks are removed
         documents = rag.list_documents()
         assert len(documents) == 0
-        
+
         cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (doc_id,))
         chunk_count = cursor.fetchone()[0]
         assert chunk_count == 0