feat(cli): implement add with recursion, add_text and list

Author: Daniele Briggi

src/cli.py

@@ -1,6 +1,8 @@
 #!/usr/bin/env python3
+from pathlib import Path
 import shlex
 import sys
+from typing import Optional
 
 import typer
 
@@ -25,18 +27,22 @@ def __call__(self, *args, **kwds):
 
 
 @app.command()
-def add(path: str):
+def add(
+    path: str = typer.Argument(..., help="File or directory path to add"),
+    recursive: bool = typer.Option(
+        False, "-r", "--recursive", help="Recursively add all files in directories"
+    ),
+):
     """Add a file path to the database"""
     rag = SQLiteRag()
-    rag.add(path)
+    rag.add(path, recursively=recursive)
 
 
 @app.command()
-def add_text(text: str):
+def add_text(text: str, uri: Optional[str] = None, metadata: dict = {}):
     """Add a text to the database"""
-    # rag = SQLiteRag()
-    # rag.add_text(text)
-    pass
+    rag = SQLiteRag()
+    rag.add_text(text, uri=uri, metadata=metadata)
 
 
 @app.command("list")

src/engine.py

@@ -1,7 +1,7 @@
 import sqlite3
-from anyio import Path
-from huggingface_hub import hf_hub_download
-from repository import Repository
+from pathlib import Path
+
+from models.chunk import Chunk
 from settings import Settings
 
 
@@ -13,7 +13,11 @@ def __init__(self, conn: sqlite3.Connection, settings: Settings):
     def load_model(self):
         """Load the model model from the specified path
         or download it from Hugging Face if not found."""
-        pass
+
+        model_path = Path(self.settings.model_path_or_name)
+        if not model_path.exists():
+            raise FileNotFoundError(f"Model file not found at {model_path}")
+
         # model_path = self.settings.model_path_or_name
         # if not Path(self.settings.model_path_or_name).exists():
         #     # check if exists locally or try to download it from Hugging Face
@@ -27,4 +31,25 @@ def load_model(self):
             f"SELECT llm_model_load('{self.settings.model_path_or_name}', '{self.settings.model_config}');"
         )
 
-    
+    def generate_embedding(self, chunks: list[Chunk]) -> list[Chunk]:
+        """Generate embedding for the given text."""
+        cursor = self._conn.cursor()
+
+        for chunk in chunks:
+            cursor.execute("SELECT llm_embed_generate(?)", (chunk.content,))
+            result = cursor.fetchone()
+
+            if result is None:
+                raise RuntimeError("Failed to generate embedding.")
+
+            chunk.embedding = result[0]
+
+        return chunks
+
+    def quantize(self) -> None:
+        """Quantitize stored vector for faster search via quantitize scan."""
+        cursor = self._conn.cursor()
+
+        cursor.execute("SELECT vector_quantize('chunks', 'embedding');")
+
+        self._conn.commit()

src/reader.py

@@ -39,6 +39,11 @@ class FileReader:
         ".yml",
     ]
 
+    @staticmethod
+    def is_supported(path: Path) -> bool:
+        """Check if the file extension is supported"""
+        return path.suffix.lower() in FileReader.extensions
+
     @staticmethod
     def parse_file(path: Path) -> str:
         try:

src/repository.py

@@ -1,5 +1,4 @@
 import json
-from pathlib import Path
 import sqlite3
 from uuid import uuid4
 
@@ -12,16 +11,6 @@ def __init__(self, conn: sqlite3.Connection, settings: Settings):
         self._conn = conn
         self.settings = settings
 
-    def load_model(self, path: str):
-        """Load the LLMA model"""
-        model_path = Path(path)
-        if not model_path.exists():
-            raise FileNotFoundError(f"Model file not found at {model_path}")
-
-        self._conn.execute(
-            f"SELECT llm_model_load('{model_path}', '{self.settings.model_config}');"
-        )
-
     def add_document(self, document: Document) -> str:
         """Add a text content to the database"""
         cursor = self._conn.cursor()
@@ -46,5 +35,25 @@ def add_document(self, document: Document) -> str:
             )
 
         self._conn.commit()
-        
+
         return document_id
+
+    def list_documents(self) -> list[Document]:
+        """List all documents in the database"""
+        cursor = self._conn.cursor()
+        cursor.execute("SELECT id, content, uri, metadata FROM documents")
+        rows = cursor.fetchall()
+
+        documents = []
+        for row in rows:
+            doc_id, content, uri, metadata = row
+            documents.append(
+                Document(
+                    id=doc_id,
+                    content=content,
+                    uri=uri,
+                    metadata=json.loads(metadata),
+                )
+            )
+
+        return documents

src/settings.py

@@ -6,7 +6,9 @@ def __init__(self, model_path_or_name: str, db_path: str = "sqliterag.db"):
         self.embedding_dim = 384  # Default for all-MiniLM-L6-v2
         self.vector_type = "FLOAT32"  # Default vector type for sqlite-vector
 
-        self.model_config = "n_ctx=0"  # See: https://github.com/sqliteai/sqlite-ai/blob/e172b9c9b9d76435be635d1e02c1e88b3681cc6e/src/sqlite-ai.c#L51-L57
-        
-        self.chunk_size = 1000  # Maximum tokens per chunk
-        self.chunk_overlap = 200  # Token overlap between chunks
+        self.model_config = "n_ctx=128"  # See: https://github.com/sqliteai/sqlite-ai/blob/e172b9c9b9d76435be635d1e02c1e88b3681cc6e/src/sqlite-ai.c#L51-L57
+
+        self.chunk_size = 256  # Maximum tokens per chunk
+        self.chunk_overlap = 32  # Token overlap between chunks
+
+        self.quantize_scan = True  # Whether to quantize the vector for faster search

src/sqliterag.py

@@ -1,3 +1,4 @@
+import glob
 from pathlib import Path
 import sqlite3
 from typing import Optional
@@ -36,18 +37,62 @@ def _ensure_initialized(self):
 
         self.ready = True
 
-    def add(self, path: str):
+    def add(self, path: str, recursively: bool = False):
         """Add the file content into the database"""
         self._ensure_initialized()
 
         if not Path(path).exists():
-            raise FileNotFoundError(f"File {path} does not exist.")
+            raise FileNotFoundError(f"{path} does not exist.")
 
-        # TODO: check the file extension
-        content = FileReader.parse_file(Path(path))
-        # TODO: include metadata extraction and mdx options (see our docsearch)
-        document = Document(content=content, uri=path)
+        files_to_process: list[Path] = []
+        path_obj = Path(path)
+
+        if path_obj.is_file():
+            files_to_process.append(path_obj)
+        elif path_obj.is_dir():
+            if recursively:
+                files_to_process = list(path_obj.rglob("*"))
+            else:
+                files_to_process = list(path_obj.glob("*"))
+
+            files_to_process = [
+                f
+                for f in files_to_process
+                if f.is_file() and FileReader.is_supported(f)
+            ]
+
+        for file_path in files_to_process:
+            # TODO: check the file extension
+            content = FileReader.parse_file(file_path)
+            # TODO: include metadata extraction and mdx options (see our docsearch)
+            document = Document(content=content, uri=str(file_path))
+            chunks = self._chunker.chunk(document.content)
+            chunks = self._engine.generate_embedding(chunks)
+            document.chunks = chunks
+
+            self._repository.add_document(document)
+
+        if self.settings.quantize_scan:
+            self._engine.quantize()
+
+    def add_text(
+        self, text: str, uri: Optional[str] = None, metadata: dict = {}
+    ) -> None:
+        """Add a text content into the database"""
+        self._ensure_initialized()
+
+        document = Document(content=text, uri=uri, metadata=metadata)
         chunks = self._chunker.chunk(document.content)
+        chunks = self._engine.generate_embedding(chunks)
         document.chunks = chunks
 
         self._repository.add_document(document)
+
+        if self.settings.quantize_scan:
+            self._engine.quantize()
+
+    def list_documents(self) -> list[Document]:
+        """List all documents in the database"""
+        self._ensure_initialized()
+
+        return self._repository.list_documents()

tests/conftest.py

@@ -10,7 +10,10 @@
 @pytest.fixture
 def db_conn():
     with tempfile.NamedTemporaryFile(suffix=".db") as tmp_db:
-        settings = Settings(model_path_or_name="all-MiniLM-L6-v2", db_path=tmp_db.name)
+        settings = Settings(
+            model_path_or_name="./capybarahermes-2.5-mistral-7b.Q4_K_M.gguf",
+            db_path=tmp_db.name,
+        )
 
     conn = sqlite3.connect(settings.db_path)
     Database.initialize(conn, settings)
@@ -23,5 +26,5 @@ def db_conn():
 @pytest.fixture
 def db_settings() -> Settings:
     with tempfile.NamedTemporaryFile(suffix=".db") as tmp_db:
-        settings = Settings(model_path_or_name="all-MiniLM-L6-v2", db_path=tmp_db.name)
+        settings = Settings(model_path_or_name="./capybarahermes-2.5-mistral-7b.Q4_K_M.gguf", db_path=tmp_db.name)
     return settings

tests/integration/test_cli.py

@@ -0,0 +1,19 @@
+from engine import Engine
+from models.chunk import Chunk
+
+
+class TestEngine:
+    def test_generate_embedding(self, db_conn):
+        conn, settings = db_conn
+        engine = Engine(conn, settings)
+        engine.load_model()
+
+        # Create a sample chunk
+        chunk = Chunk(content="This is a test chunk for embedding generation.")
+
+        # Generate embedding
+        result_chunks = engine.generate_embedding([chunk])
+
+        assert len(result_chunks) == 1
+        assert result_chunks[0].embedding is not None
+        assert isinstance(result_chunks[0].embedding, bytes)

tests/test_engine.py

@@ -76,3 +76,31 @@ def test_add_document_with_chunks(self, db_conn):
         assert chunk_rows[0][1] == b"\x00" * 384
         assert chunk_rows[1][0] == "Chunk 2 content"
         assert chunk_rows[1][1] == b"\x00" * 384
+
+    def test_list_documents(self, db_conn):
+        conn, settings = db_conn
+
+        repo = Repository(conn, settings)
+
+        doc1 = Document(
+            content="Document 1 content.", uri="doc1.txt", metadata={"author": "test1"}
+        )
+        doc2 = Document(
+            content="Document 2 content.", uri="doc2.txt", metadata={"author": "test2"}
+        )
+
+        repo.add_document(doc1)
+        repo.add_document(doc2)
+
+        documents = repo.list_documents()
+
+        assert len(documents) == 2
+        assert documents[0].id is not None
+        assert documents[0].uri == "doc1.txt"
+        assert documents[0].content == "Document 1 content."
+        assert documents[0].metadata == {"author": "test1"}
+
+        assert documents[1].id is not None
+        assert documents[1].content == "Document 2 content."
+        assert documents[1].uri == "doc2.txt"
+        assert documents[1].metadata == {"author": "test2"}