feat(settings): improve

Author: Daniele Briggi

.devcontainer/devcontainer.json

@@ -15,5 +15,4 @@
             ]
         }
     },
-    "runArgs": ["--network=host"]
 }

src/sqlite_rag/cli.py

@@ -2,12 +2,12 @@
 import json
 import shlex
 import sys
-from dataclasses import replace
 from typing import Optional
 
 import typer
 
-from sqlite_rag.settings import Settings
+from sqlite_rag.database import Database
+from sqlite_rag.settings import SettingsManager
 
 from .sqliterag import SQLiteRag
 
@@ -40,7 +40,6 @@ def show_settings():
         typer.echo(f"  {key}: {value}")
 
 
-# TODO: separate store settings from SQLiteRag.create()?
 @app.command("set")
 def set_settings(
     model_path_or_name: Optional[str] = typer.Option(
@@ -106,9 +105,9 @@ def set_settings(
         show_settings()
         return
 
-    # Create new settings with updated fields
-    new_settings = replace(Settings(), **updates)
-    SQLiteRag.create(settings=new_settings)
+    conn = Database.new_connection()
+    settings_manager = SettingsManager(conn)
+    settings_manager.prepare_settings(updates)
 
     show_settings()
     typer.echo("Settings updated.")

src/sqlite_rag/engine.py

@@ -37,7 +37,8 @@ def load_model(self):
         #     )
 
         self._conn.execute(
-            f"SELECT llm_model_load('{self._settings.model_path_or_name}', '{self._settings.model_config}');"
+            "SELECT llm_model_load(?, ?);",
+            (self._settings.model_path_or_name, self._settings.model_config),
         )
 
     def process(self, document: Document) -> Document:
@@ -52,6 +53,8 @@ def generate_embedding(self, chunks: list[Chunk]) -> list[Chunk]:
         for chunk in chunks:
             cursor = self._conn.cursor()
 
+            self.create_new_context()
+
             try:
                 cursor.execute(
                     "SELECT llm_embed_generate(?) AS embedding", (chunk.content,)
@@ -75,23 +78,29 @@ def quantize(self) -> None:
 
         cursor.execute("SELECT vector_quantize('chunks', 'embedding');")
 
-        self._conn.commit()
-
     def quantize_preload(self) -> None:
         """Preload quantized vectors into memory for faster search."""
         cursor = self._conn.cursor()
 
         cursor.execute("SELECT vector_quantize_preload('chunks', 'embedding');")
 
-        self._conn.commit()
-
     def quantize_cleanup(self) -> None:
         """Clean up internal structures related to a previously quantized table/column."""
         cursor = self._conn.cursor()
 
         cursor.execute("SELECT vector_quantize_cleanup('chunks', 'embedding');")
 
-        self._conn.commit()
+    def create_new_context(self) -> None:
+        """"""
+        cursor = self._conn.cursor()
+
+        cursor.execute("SELECT llm_context_create(?);", (self._settings.model_config,))
+
+    def free_context(self) -> None:
+        """"""
+        cursor = self._conn.cursor()
+
+        cursor.execute("SELECT llm_context_free();")
 
     def search(self, query: str, limit: int = 10) -> list[DocumentResult]:
         """Semantic search and full-text search sorted with Reciprocal Rank Fusion."""

src/sqlite_rag/settings.py

@@ -1,6 +1,7 @@
 import json
 import sqlite3
-from dataclasses import asdict, dataclass, fields
+from dataclasses import asdict, dataclass, fields, replace
+from typing import Any, Optional
 
 
 @dataclass
@@ -54,6 +55,34 @@ def _ensure_table_exists(self):
         )
         self.connection.commit()
 
+    def prepare_settings(self, settings: Optional[dict[str, Any]]) -> Settings:
+        """Load, initialize or update settings.
+
+        If no settings are provided, load the last used settings or use defaults.
+        If settings are provided, check for critical changes and update them.
+        """
+        current_settings = self.load_settings()
+        if current_settings:
+            if settings:
+                new_settings = replace(current_settings, **settings)
+
+                if self.has_critical_changes(new_settings, current_settings):
+                    raise ValueError(
+                        "Critical settings changes detected. Please reset the database."
+                    )
+                # Update new settings
+                current_settings = self.store(new_settings)
+        elif settings:
+            # Store initial settings with customs
+            new_settings = replace(Settings(), **settings)
+            current_settings = self.store(new_settings)
+        else:
+            # Store default settings
+            new_settings = Settings()
+            current_settings = self.store(new_settings)
+
+        return current_settings
+
     def load_settings(self) -> Settings | None:
         cursor = self.connection.cursor()
 

src/sqlite_rag/sqliterag.py

@@ -1,5 +1,5 @@
 import sqlite3
-from dataclasses import asdict, replace
+from dataclasses import asdict
 from pathlib import Path
 from typing import Optional
 
@@ -51,30 +51,10 @@ def create(
         If no new settings are provided, it uses the default settings or load
         the settings used in the last execution."""
 
-        conn = sqlite3.connect(db_path)
-        conn.row_factory = sqlite3.Row
+        conn = Database.new_connection(db_path)
 
-        # Load, initialize or update settings
         settings_manager = SettingsManager(conn)
-        current_settings = settings_manager.load_settings()
-        if current_settings:
-            if settings:
-                settings = replace(current_settings, **asdict(settings))
-
-                if settings_manager.has_critical_changes(settings, current_settings):
-                    raise ValueError(
-                        "Critical settings changes detected. Please reset the database."
-                    )
-                # Update new settings
-                current_settings = settings_manager.store(settings)
-        elif settings:
-            # Store initial settings with customs
-            settings = replace(Settings(), **asdict(settings))
-            current_settings = settings_manager.store(settings)
-        else:
-            # Store default settings
-            settings = Settings()
-            current_settings = settings_manager.store(settings)
+        current_settings = settings_manager.prepare_settings(settings)
 
         Database.initialize(conn, current_settings)
 
@@ -97,6 +77,8 @@ def add(
 
         files_to_process = FileReader.collect_files(Path(path), recursive=recursive)
 
+        self._engine.create_new_context()
+
         processed = 0
         self._logger.info(f"Processing {len(files_to_process)} files...")
         for file_path in files_to_process:
@@ -138,6 +120,7 @@ def add_text(
         self._ensure_initialized()
 
         document = Document(content=text, uri=uri, metadata=metadata)
+        self._engine.create_new_context()
         document = self._engine.process(document)
 
         self._repository.add_document(document)
@@ -177,6 +160,8 @@ def rebuild(self, remove_missing: bool = False) -> dict:
         not_found = 0
         removed = 0
 
+        self._engine.create_new_context()
+
         for doc in documents:
             doc_id = doc.id or ""
 
@@ -218,8 +203,8 @@ def rebuild(self, remove_missing: bool = False) -> dict:
                 except Exception as e:
                     self._logger.error(f"Error processing text document {doc.id}: {e}")
 
-        if self._settings.quantize_scan:
-            self._engine.quantize()
+            if self._settings.quantize_scan:
+                self._engine.quantize()
 
         return {
             "total": total_docs,
@@ -247,9 +232,13 @@ def reset(self) -> bool:
             self._logger.error(f"Error during database reset: {e}")
             return False
 
-    def search(self, query: str, top_k: int = 10) -> list[DocumentResult]:
+    def search(
+        self, query: str, top_k: int = 10, new_context: bool = True
+    ) -> list[DocumentResult]:
         """Search for documents matching the query"""
         self._ensure_initialized()
+        if new_context:
+            self._engine.create_new_context()
 
         if self._settings.quantize_preload:
             self._engine.quantize_preload()

tests/integration/test_cli.py

@@ -2,12 +2,26 @@
 import tempfile
 from pathlib import Path
 
+from pytest import fixture
 from typer.testing import CliRunner
 
 from sqlite_rag.cli import app
 from sqlite_rag.settings import Settings
 
 
+@fixture
+def temp_dir():
+    """Change the current working directory in order to create
+    the default database in a temporary location."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        original_cwd = os.getcwd()
+        try:
+            os.chdir(tmpdir)
+            yield tmpdir
+        finally:
+            os.chdir(original_cwd)
+
+
 class TestCLI:
     def test_search_exact_match(self):
         # Use SQLiteRag to set up the test data directly
@@ -69,3 +83,23 @@ def test_search_exact_match(self):
             assert "Found 1 documents" in result.stdout
             # For exact match with cosine distance, we expect distance close to 0.0
             assert "0.000000" in result.stdout or "0.00000" in result.stdout
+
+    def test_set_settings(self, temp_dir):
+        runner = CliRunner()
+
+        model_path = "mypath/mymodel.gguf"
+
+        result = runner.invoke(
+            app,
+            [
+                "set",
+                "--model-path-or-name",
+                model_path,
+                "--other-vector-config",
+                "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

tests/test_settings.py

@@ -144,3 +144,64 @@ def test_has_critical_changes(self, db_conn):
             new_settings, current_settings
         )
         assert has_changes
+
+    def test_prepare_settings_with_no_existing_and_no_input(self, db_conn):
+        """Test prepare_settings returns default settings when no existing settings and no input"""
+        settings_manager = SettingsManager(db_conn[0])
+
+        result = settings_manager.prepare_settings(None)
+
+        defaults = Settings()
+        assert result.model_path_or_name == defaults.model_path_or_name
+        assert result.embedding_dim == defaults.embedding_dim
+        assert result.chunk_size == defaults.chunk_size
+
+    def test_prepare_settings_with_no_existing_and_custom_input(self, db_conn):
+        """Test prepare_settings stores and returns custom settings when no existing settings"""
+        settings_manager = SettingsManager(db_conn[0])
+
+        result = settings_manager.prepare_settings(
+            {"chunk_size": 5000, "quantize_scan": False}
+        )
+
+        assert result.chunk_size == 5000
+        assert result.quantize_scan is False
+        # Check defaults are preserved
+        defaults = Settings()
+        assert result.model_path_or_name == defaults.model_path_or_name
+
+    def test_prepare_settings_with_existing_and_no_input(self, db_conn):
+        """Test prepare_settings returns existing settings when they exist and no input provided"""
+        settings_manager = SettingsManager(db_conn[0])
+        existing = Settings(chunk_size=3000, quantize_scan=False)
+        settings_manager.store(existing)
+
+        result = settings_manager.prepare_settings(None)
+
+        assert result.chunk_size == 3000
+        assert result.quantize_scan is False
+
+    def test_prepare_settings_with_existing_and_non_critical_updates(self, db_conn):
+        """Test prepare_settings updates non-critical settings when existing settings present"""
+        settings_manager = SettingsManager(db_conn[0])
+        existing = Settings(chunk_size=3000, chunk_overlap=100)
+        settings_manager.store(existing)
+
+        result = settings_manager.prepare_settings(
+            {"chunk_size": 4000, "quantize_scan": False}
+        )
+
+        assert result.chunk_size == 4000
+        assert result.chunk_overlap == 100
+        assert result.quantize_scan is False
+
+    def test_prepare_settings_with_critical_changes_raises_error(self, db_conn):
+        """Test prepare_settings raises ValueError when critical settings change"""
+        settings_manager = SettingsManager(db_conn[0])
+        existing = Settings()
+        settings_manager.store(existing)
+
+        import pytest
+
+        with pytest.raises(ValueError, match="Critical settings changes detected"):
+            settings_manager.prepare_settings({"model_path_or_name": "new_model"})