fix(context): contex creation

Author: Daniele Briggi

src/sqlite_rag/database.py

@@ -91,7 +91,7 @@ def _create_schema(conn: sqlite3.Connection, settings: Settings):
 
         cursor.execute(
             f"""
-            SELECT vector_init('chunks', 'embedding', 'type={settings.vector_type},dimension={settings.embedding_dim},{settings.other_vector_config}');
+            SELECT vector_init('chunks', 'embedding', 'type={settings.vector_type},dimension={settings.embedding_dim},{settings.other_vector_options}');
         """
         )
 

src/sqlite_rag/engine.py

@@ -1,3 +1,4 @@
+import json
 import re
 import sqlite3
 from pathlib import Path
@@ -38,7 +39,7 @@ def load_model(self):
 
         self._conn.execute(
             "SELECT llm_model_load(?, ?);",
-            (self._settings.model_path_or_name, self._settings.model_config),
+            (self._settings.model_path_or_name, self._settings.model_options),
         )
 
     def process(self, document: Document) -> Document:
@@ -53,8 +54,6 @@ 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,)
@@ -94,7 +93,9 @@ def create_new_context(self) -> None:
         """"""
         cursor = self._conn.cursor()
 
-        cursor.execute("SELECT llm_context_create(?);", (self._settings.model_config,))
+        cursor.execute(
+            "SELECT llm_context_create(?);", (self._settings.model_context_options,)
+        )
 
     def free_context(self) -> None:
         """"""
@@ -117,8 +118,9 @@ def search(self, query: str, limit: int = 10) -> list[DocumentResult]:
         )
 
         cursor = self._conn.cursor()
+        # TODO: understand how to sort results depending on the distance metric
+        # Eg, for cosine distance, higher is better (closer to 1)
         cursor.execute(
-            # TODO: use vector_convert_XXX to convert the query to the correct type
             f"""
             -- sqlite-vector KNN vector search results
             WITH vec_matches AS (
@@ -159,6 +161,7 @@ def search(self, query: str, limit: int = 10) -> list[DocumentResult]:
                 documents.id,
                 documents.uri,
                 documents.content as document_content,
+                documents.metadata,
                 chunks.content AS snippet,
                 vec_rank,
                 fts_rank,
@@ -170,7 +173,7 @@ def search(self, query: str, limit: int = 10) -> list[DocumentResult]:
                 JOIN documents ON documents.id = chunks.document_id
             ORDER BY combined_rank DESC
             ;
-            """,
+            """,  # nosec B608
             {
                 "query": query,
                 "query_embedding": query_embedding,
@@ -188,6 +191,7 @@ def search(self, query: str, limit: int = 10) -> list[DocumentResult]:
                     id=row["id"],
                     uri=row["uri"],
                     content=row["document_content"],
+                    metadata=json.loads(row["metadata"]) if row["metadata"] else {},
                 ),
                 snippet=row["snippet"],
                 vec_rank=row["vec_rank"],

src/sqlite_rag/settings.py

@@ -14,14 +14,21 @@ class Settings:
     model_path_or_name: str = (
         "./models/Qwen/Qwen3-Embedding-0.6B-GGUF/Qwen3-Embedding-0.6B-f16.gguf"
     )
-    model_config: str = "n_ctx=128,pooling_type=last,normalize_embedding=1"
+    # See: https://github.com/sqliteai/sqlite-ai/blob/main/API.md#llm_model_loadpath-text-options-text
+    model_options: str = ""
+    # See: https://github.com/sqliteai/sqlite-ai/blob/main/API.md#llm_context_createoptions-text
+    model_context_options: str = (
+        "generate_embedding=1,normalize_embedding=1,pooling_type=mean"
+    )
 
     vector_type: str = "FLOAT16"
     embedding_dim: int = 1024
-    other_vector_config: str = "distance=cosine"  # e.g. distance=metric,other=value,...
+    other_vector_options: str = (
+        "distance=cosine"  # e.g. distance=metric,other=value,...
+    )
 
     chunk_size: int = 128
-    # Token overlap between chunks
+    # Tokens overlap between chunks
     chunk_overlap: int = 20
 
     #

src/sqlite_rag/sqliterag.py

@@ -111,6 +111,8 @@ def add(
 
             processed += 1
 
+        self._engine.free_context()
+
         return processed
 
     def add_text(
@@ -120,6 +122,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)
 
@@ -128,6 +131,8 @@ def add_text(
         if self._settings.quantize_scan:
             self._engine.quantize()
 
+        self._engine.free_context()
+
     def list_documents(self) -> list[Document]:
         """List all documents in the database"""
         self._ensure_initialized()
@@ -206,6 +211,8 @@ def rebuild(self, remove_missing: bool = False) -> dict:
             if self._settings.quantize_scan:
                 self._engine.quantize()
 
+        self._engine.free_context()
+
         return {
             "total": total_docs,
             "reprocessed": reprocessed,

tests/test_settings.py

@@ -9,7 +9,7 @@ def test_store_settings(self, db_conn):
         settings_manager = SettingsManager(db_conn[0])
         settings = Settings(
             model_path_or_name="test_model",
-            model_config="test_config",
+            model_options="test_config",
             embedding_dim=768,
             vector_type="test_store",
             chunk_overlap=100,
@@ -23,7 +23,7 @@ def test_store_settings(self, db_conn):
 
         assert stored_settings is not None
         assert stored_settings.model_path_or_name == "test_model"
-        assert stored_settings.model_config == "test_config"
+        assert stored_settings.model_options == "test_config"
         assert stored_settings.embedding_dim == 768
         assert stored_settings.vector_type == "test_store"
         assert stored_settings.chunk_overlap == 100
@@ -34,7 +34,7 @@ def test_store_settings_when_exist(self, db_conn):
         settings_manager = SettingsManager(db_conn[0])
         settings = Settings(
             model_path_or_name="test_model",
-            model_config="test_config",
+            model_options="test_config",
             embedding_dim=768,
             vector_type="test_store",
             chunk_overlap=100,
@@ -47,7 +47,7 @@ def test_store_settings_when_exist(self, db_conn):
         # Store again with different values
         new_settings = Settings(
             model_path_or_name="new_model",
-            model_config="new_config",
+            model_options="new_config",
             embedding_dim=512,
             vector_type="new_store",
             chunk_overlap=50,
@@ -60,7 +60,7 @@ def test_store_settings_when_exist(self, db_conn):
 
         assert stored_settings is not None
         assert stored_settings.model_path_or_name == "new_model"
-        assert stored_settings.model_config == "new_config"
+        assert stored_settings.model_options == "new_config"
         assert stored_settings.embedding_dim == 512
         assert stored_settings.vector_type == "new_store"
         assert stored_settings.chunk_overlap == 50
@@ -82,7 +82,7 @@ def test_load_settings_with_defaults(self, db_conn):
 
         assert loaded_settings is not None
         assert loaded_settings.model_path_or_name == settings.model_path_or_name
-        assert loaded_settings.model_config == settings.model_config
+        assert loaded_settings.model_options == settings.model_options
         assert loaded_settings.embedding_dim == settings.embedding_dim
         assert loaded_settings.vector_type == settings.vector_type
         assert loaded_settings.chunk_overlap == settings.chunk_overlap