refact(quantize): manage it manually

- detach from quantize_scan setting

Author: Daniele Briggi

src/sqlite_rag/cli.py

@@ -137,7 +137,7 @@ def add(
     rag.add(
         path,
         recursive=recursive,
-        absolute_paths=absolute_paths,
+        use_absolute_paths=absolute_paths,
         metadata=json.loads(metadata or "{}"),
     )
 
@@ -345,6 +345,29 @@ def search(
             typer.echo(f"{idx:<3} {snippet:<60} {uri:<40}")
 
 
+@app.command()
+def quantize(
+    cleanup: bool = typer.Option(
+        False,
+        "--cleanup",
+        help="Clean up quantization structures instead of creating them",
+    )
+):
+    """Quantize vectors for faster search or clean up quantization structures"""
+    rag = SQLiteRag.create()
+
+    if cleanup:
+        typer.echo("Cleaning up quantization structures...")
+        rag.quantize_cleanup()
+        typer.echo("Quantization cleanup completed.")
+    else:
+        typer.echo("Starting vector quantization...")
+        rag.quantize_vectors()
+        typer.echo(
+            "Vector quantization completed. Now you can search with `--quantize-scan` and `--quantize-preload` enabled."
+        )
+
+
 @app.command("download-model")
 def download_model(
     model_id: str = typer.Argument(

src/sqlite_rag/engine.py

@@ -3,6 +3,7 @@
 import sqlite3
 from pathlib import Path
 
+from sqlite_rag.logger import Logger
 from sqlite_rag.models.document_result import DocumentResult
 
 from .chunker import Chunker
@@ -19,6 +20,7 @@ def __init__(self, conn: sqlite3.Connection, settings: Settings, chunker: Chunke
         self._conn = conn
         self._settings = settings
         self._chunker = chunker
+        self._logger = Logger()
 
     def load_model(self):
         """Load the model model from the specified path
@@ -77,6 +79,9 @@ def quantize(self) -> None:
 
         cursor.execute("SELECT vector_quantize('chunks', 'embedding');")
 
+        self._conn.commit()
+        self._logger.debug("Quantization completed.")
+
     def quantize_preload(self) -> None:
         """Preload quantized vectors into memory for faster search."""
         cursor = self._conn.cursor()
@@ -89,6 +94,8 @@ def quantize_cleanup(self) -> None:
 
         cursor.execute("SELECT vector_quantize_cleanup('chunks', 'embedding');")
 
+        self._conn.commit()
+
     def create_new_context(self) -> None:
         """"""
         cursor = self._conn.cursor()

src/sqlite_rag/sqliterag.py

@@ -32,13 +32,6 @@ def _ensure_initialized(self):
         if not self.ready:
             self._engine.load_model()
 
-        if self._settings.quantize_scan:
-            # TODO: quantize again if already quantized?
-            # TODO: DO NOT REPEAT FOR NOTHING, it takes time
-            self._engine.quantize()
-        else:
-            self._engine.quantize_cleanup()
-
         self.ready = True
 
     @staticmethod
@@ -64,7 +57,7 @@ def add(
         self,
         path: str,
         recursive: bool = False,
-        absolute_paths: bool = True,
+        use_absolute_paths: bool = True,
         metadata: dict = {},
     ) -> int:
         """Add the file content into the database"""
@@ -81,36 +74,36 @@ def add(
 
         processed = 0
         self._logger.info(f"Processing {len(files_to_process)} files...")
-        for file_path in files_to_process:
-            content = FileReader.parse_file(file_path)
+        try:
+            for file_path in files_to_process:
+                content = FileReader.parse_file(file_path)
 
-            if not content:
-                self._logger.warning(f"Skipping empty file: {file_path}")
-                continue
+                if not content:
+                    self._logger.warning(f"Skipping empty file: {file_path}")
+                    continue
 
-            uri = (
-                str(file_path.absolute())
-                if absolute_paths
-                else str(file_path.relative_to(parent))
-            )
-            document = Document(content=content, uri=uri, metadata=metadata)
+                uri = (
+                    str(file_path.absolute())
+                    if use_absolute_paths
+                    else str(file_path.relative_to(parent))
+                )
+                document = Document(content=content, uri=uri, metadata=metadata)
 
-            exists = self._repository.document_exists_by_hash(document.hash())
-            if exists:
-                self._logger.info(f"Unchanged: {file_path}")
-                continue
+                exists = self._repository.document_exists_by_hash(document.hash())
+                if exists:
+                    self._logger.info(f"Unchanged: {file_path}")
+                    continue
 
-            self._logger.info(f"Processing: {file_path}")
-            document = self._engine.process(document)
+                self._logger.info(f"Processing: {file_path}")
+                document = self._engine.process(document)
 
-            self._repository.add_document(document)
+                self._repository.add_document(document)
 
-            # TODO: try/expect and run in the finally block?
+                processed += 1
+        finally:
             if self._settings.quantize_scan:
                 self._engine.quantize()
 
-            processed += 1
-
         self._engine.free_context()
 
         return processed
@@ -247,7 +240,7 @@ def search(
         if new_context:
             self._engine.create_new_context()
 
-        if self._settings.quantize_preload:
+        if self._settings.quantize_scan and self._settings.quantize_preload:
             self._engine.quantize_preload()
 
         return self._engine.search(query, limit=top_k)
@@ -257,7 +250,17 @@ def get_settings(self) -> dict:
         versions = self._engine.versions()
         return {**versions, **asdict(self._settings)}
 
-    def destroy(self) -> None:
+    def quantize_vectors(self) -> None:
+        """Quantize vectors for faster search"""
+        self._ensure_initialized()
+        self._engine.quantize()
+
+    def quantize_cleanup(self) -> None:
+        """Clean up quantization structures"""
+        self._ensure_initialized()
+        self._engine.quantize_cleanup()
+
+    def close(self) -> None:
         """Free up resources"""
         if self._conn:
             self._engine.close()

tests/integration/test_cli.py

@@ -46,7 +46,7 @@ def test_search_exact_match(self):
                         "set",
                         "--model-path-or-name",
                         str(model_path),
-                        "--other-vector-config",
+                        "--other-vector-options",
                         "distance=cosine",
                     ],
                 )
@@ -95,11 +95,11 @@ def test_set_settings(self, temp_dir):
                 "set",
                 "--model-path-or-name",
                 model_path,
-                "--other-vector-config",
+                "--other-vector-options",
                 "distance=L2",
             ],
         )
         assert result.exit_code == 0
 
         assert f"model_path_or_name: {model_path}" in result.stdout
-        assert "other_vector_config: distance=L2" in result.stdout
+        assert "other_vector_options: distance=L2" in result.stdout

tests/test_sqlite_rag.py

@@ -3,6 +3,8 @@
 import tempfile
 from pathlib import Path
 
+import pytest
+
 from sqlite_rag import SQLiteRag
 
 
@@ -94,7 +96,7 @@ def test_add_with_absolute_paths_option_true(self):
 
         rag = SQLiteRag.create(":memory:")
 
-        rag.add(temp_file_path, absolute_paths=True)
+        rag.add(temp_file_path, use_absolute_paths=True)
 
         conn = rag._conn
         cursor = conn.execute("SELECT uri FROM documents")
@@ -109,7 +111,7 @@ def test_add_with_absolute_paths_option_false(self):
 
         rag = SQLiteRag.create(":memory:")
 
-        rag.add(str(temp_file_path), absolute_paths=False)
+        rag.add(str(temp_file_path), use_absolute_paths=False)
 
         conn = rag._conn
         cursor = conn.execute("SELECT uri FROM documents")
@@ -540,3 +542,41 @@ def test_search_exact_match(self):
         assert expected_string == results[0].document.content
         assert 1 == results[0].vec_rank
         assert 0.0 == results[0].vec_distance
+
+    @pytest.mark.parametrize(
+        "quantize_scan", [True, False], ids=["quantize", "no-quantize"]
+    )
+    def test_search_samples_exact_match_by_scan_type(self, quantize_scan: bool):
+        # Test that searching for exact content from sample files returns distance 0
+        # FTS not included in the combined score
+        settings = {
+            "other_vector_options": "distance=cosine",
+            "weight_fts": 0.0,
+            "quantize_scan": quantize_scan,
+        }
+
+        temp_file_path = os.path.join(tempfile.mkdtemp(), "mydb.db")
+        rag = SQLiteRag.create(temp_file_path, settings=settings)
+
+        # Index all sample files
+        samples_dir = Path(__file__).parent / "assets" / "samples"
+        rag.add(str(samples_dir))
+
+        # Get all sample files and test each one
+        sample_files = list(samples_dir.glob("*.txt"))
+
+        for sample_file in sample_files:
+            file_content = sample_file.read_text(encoding="utf-8")
+
+            # Search for the exact content
+            results = rag.search(file_content, top_k=2)
+
+            assert len(results) == 2
+
+            first_result = results[0]
+            assert first_result.vec_distance == 0.0
+            assert first_result.document.content == file_content
+
+            # Second result should have distance > 0
+            second_result = results[1]
+            assert second_result.vec_distance and second_result.vec_distance > 0.0