chore: implement list command

- reorganize project

Author: Daniele Briggi

.devcontainer/devcontainer.json

@@ -14,5 +14,6 @@
                 "hbenl.vscode-test-explorer"
             ]
         }
-    }
+    },
+    "runArgs": ["--network=host"]
 }

pyproject.toml

@@ -24,11 +24,11 @@ dev = [
 ]
 
 [project.scripts]
-sqlite-rag = "src.cli:cli"
+sqlite-rag = "sqlite_rag.cli:cli"
 
 [tool.setuptools.packages.find]
-where = ["."]
-include = ["src*"]
+where = ["src"]
+include = ["sqlite_rag*"]
 
 [tool.isort]
 # make the tools compatible to each other

src/__init__.py

@@ -0,0 +1,19 @@
+"""
+SQLite RAG - Hybrid search with SQLite AI and SQLite Vector
+
+This package provides both a library interface and CLI for document 
+indexing and semantic search using SQLite with AI and vector extensions.
+"""
+
+from .sqliterag import SQLiteRag
+from .models.document import Document
+from .models.chunk import Chunk
+from .settings import Settings
+
+__version__ = "0.1.0"
+__all__ = [
+    "SQLiteRag",
+    "Document", 
+    "Chunk",
+    "Settings",
+]

src/engine.py

@@ -1,8 +1,8 @@
 import sqlite3
 from typing import List
 
-from models.chunk import Chunk
-from settings import Settings
+from .models.chunk import Chunk
+from .settings import Settings
 import math
 
 

src/sqlite_rag/__init__.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
-from pathlib import Path
 import shlex
 import sys
+from pathlib import Path
 from typing import Optional
 
 import typer
 
-from sqliterag import SQLiteRag
+from .sqliterag import SQLiteRag
 
 
 class CLI:
@@ -39,16 +39,39 @@ def add(
 
 
 @app.command()
-def add_text(text: str, uri: Optional[str] = None, metadata: dict = {}):
+def add_text(text: str, uri: Optional[str] = None):
     """Add a text to the database"""
     rag = SQLiteRag()
-    rag.add_text(text, uri=uri, metadata=metadata)
+    rag.add_text(text, uri=uri, metadata={})
 
 
 @app.command("list")
 def list_documents():
     """List all documents in the database"""
-    pass
+    rag = SQLiteRag()
+    documents = rag.list_documents()
+
+    if not documents:
+        typer.echo("No documents found in the database.")
+        return
+
+    # Print stats
+    typer.echo(f"Total documents: {len(documents)}")
+    typer.echo("Documents:")
+    typer.echo("-" * 106)
+
+    typer.echo(f"{'ID':<36} {'URI/Content':<50} {'Created At':<20}")
+    typer.echo("-" * 106)
+
+    for doc in documents:
+        # Show URI if available, otherwise show first chars of content
+        uri_or_content = doc.uri or (
+            doc.content[:47] + "..." if len(doc.content) > 47 else doc.content
+        )
+        created_at = (
+            doc.created_at.strftime("%Y-%m-%d %H:%M:%S") if doc.created_at else "N/A"
+        )
+        typer.echo(f"{doc.id:<36} {uri_or_content:<50} {created_at:<20}")
 
 
 @app.command()

src/chunker.py src/sqlite_rag/chunker.pysrc/chunker.py renamed to src/sqlite_rag/chunker.py

@@ -1,7 +1,7 @@
 import sqlite3
 from pathlib import Path
 
-from settings import Settings
+from .settings import Settings
 
 
 class Database:
@@ -12,9 +12,9 @@ def initialize(conn: sqlite3.Connection, settings: Settings) -> sqlite3.Connecti
         """Initialize the database with extensions and schema"""
         conn.enable_load_extension(True)
         try:
-            conn.load_extension(str(Path(__file__).parent.parent / "extensions" / "ai"))
+            conn.load_extension(str(Path(__file__).parent.parent.parent / "extensions" / "ai"))
             conn.load_extension(
-                str(Path(__file__).parent.parent / "extensions" / "vector")
+                str(Path(__file__).parent.parent.parent / "extensions" / "vector")
             )
         except sqlite3.OperationalError as e:
             raise RuntimeError(

src/cli.py src/sqlite_rag/cli.pysrc/cli.py renamed to src/sqlite_rag/cli.py

@@ -0,0 +1,153 @@
+import sqlite3
+from pathlib import Path
+
+from .chunker import Chunker
+from .models.chunk import Chunk
+from .models.document import Document
+from .settings import Settings
+
+
+class Engine:
+    def __init__(self, conn: sqlite3.Connection, settings: Settings, chunker: Chunker):
+        self._conn = conn
+        self.settings = settings
+        self._chunker = chunker
+
+    def load_model(self):
+        """Load the model model from the specified path
+        or download it from Hugging Face if not found."""
+
+        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
+        #     model_path = hf_hub_download(
+        #         repo_id=self.settings.model_path_or_name,
+        #         filename="model-q4_0.gguf",  # GGUF format
+        #         cache_dir="./models"
+        #     )
+
+        self._conn.execute(
+            f"SELECT llm_model_load('{self.settings.model_path_or_name}', '{self.settings.model_config}');"
+        )
+
+    def process(self, document: Document) -> Document:
+        chunks = self._chunker.chunk(document.content)
+        chunks = self.generate_embedding(chunks)
+        document.chunks = chunks
+        return document
+
+    # TODO: better to get a list of str and return a list of embeddings?
+    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(?, 'json_output=0')", (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()
+
+    def search(self, query: str, limit: int = 10) -> list[Document]:
+        """Semantic search and full-text search sorted with Reciprocal Rank Fusion."""
+        cursor = self._conn.cursor()
+
+        query_embedding = self.generate_embedding([Chunk(content=query)])[0].embedding
+
+        cursor.execute(
+            # TODO: use vector_convert_XXX to convert the query to the correct type
+            """
+            -- sqlite-vector KNN vector search results
+            WITH vec_matches AS (
+                SELECT
+                    row_number() over (order by v.distance) AS rank_number,
+                    c.document_id,
+                    v.distance
+                FROM chunks AS c
+                    JOIN vector_quantize_scan('chunks', 'embedding', vector_convert_f32(:query_embedding), :k) AS v
+                    ON c.id = v.rowid
+            ),
+            -- Full-text search results
+            fts_matches AS (
+                SELECT
+                    row_number() over (order by rank) AS rank_number,
+                    c.document_id,
+                    rank AS score
+                FROM chunks_fts AS c_fts
+                    JOIN chunks AS c ON c_fts.rowid = c.id
+                WHERE c_fts.content MATCH :query
+                LIMIT :k
+            ),
+            -- combine FTS5 + vector search results with RRF
+            matches AS (
+                SELECT
+                    documents.id,
+                    documents.uri,
+                    documents.content,
+                    vec_matches.rank_number AS vec_rank,
+                    fts_matches.rank_number AS fts_rank,
+                    -- Reciprocal Rank Fusion score
+                    (
+                        COALESCE(1.0 / (:rrf_k + vec_matches.rank_number), 0.0) * :weight_vec +
+                        COALESCE(1.0 / (:rrf_k + fts_matches.rank_number), 0.0) * :weight_fts
+                    ) AS combined_rank,
+                    vec_matches.distance AS vec_distance,
+                    fts_matches.score AS fts_score
+                FROM vec_matches
+                    FULL OUTER JOIN fts_matches
+                        ON vec_matches.document_id = fts_matches.document_id
+                    JOIN documents ON documents.id = COALESCE(vec_matches.document_id, fts_matches.document_id)
+                ORDER BY combined_rank DESC
+            )
+            SELECT
+                id,
+                uri,
+                content,
+                vec_rank,
+                fts_rank,
+                combined_rank,
+                vec_distance,
+                fts_score
+            FROM matches;
+            """,
+            {
+                "query": f'"{query}"',
+                "query_embedding": query_embedding,
+                "k": limit,
+                # TODO: move to settings or costants
+                "rrf_k": 60,
+                "weight_fts": 1.0,
+                "weight_vec": 1.0,
+            },
+        )
+
+        rows = cursor.fetchall()
+        return [
+            Document(
+                id=row[0],
+                uri=row[1],
+                content=row[2],
+                vec_rank=row[3],
+                fts_rank=row[4],
+                combined_rank=row[5],
+                vec_distance=row[6],
+                fts_score=row[7],
+            )
+            for row in rows
+        ]