haiku.rag 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

haiku/rag/app.py

@@ -40,7 +40,7 @@ class HaikuRAGApp:
                 f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
             )
 
-    async def get_document(self, doc_id: int):
+    async def get_document(self, doc_id: str):
         async with HaikuRAG(db_path=self.db_path) as self.client:
             doc = await self.client.get_document_by_id(doc_id)
             if doc is None:
@@ -48,14 +48,14 @@ class HaikuRAGApp:
                 return
             self._rich_print_document(doc, truncate=False)
 
-    async def delete_document(self, doc_id: int):
+    async def delete_document(self, doc_id: str):
         async with HaikuRAG(db_path=self.db_path) as self.client:
             await self.client.delete_document(doc_id)
             self.console.print(f"[b]Document {doc_id} deleted successfully.[/b]")
 
-    async def search(self, query: str, limit: int = 5, k: int = 60):
+    async def search(self, query: str, limit: int = 5):
         async with HaikuRAG(db_path=self.db_path) as self.client:
-            results = await self.client.search(query, limit=limit, k=k)
+            results = await self.client.search(query, limit=limit)
             if not results:
                 self.console.print("[red]No results found.[/red]")
                 return

haiku/rag/cli.py

@@ -8,6 +8,7 @@ from rich.console import Console
 
 from haiku.rag.app import HaikuRAGApp
 from haiku.rag.config import Config
+from haiku.rag.migration import migrate_sqlite_to_lancedb
 from haiku.rag.utils import is_up_to_date
 
 if not Config.ENV == "development":
@@ -47,7 +48,7 @@ def main(
         help="Show version and exit",
     ),
 ):
-    """haiku.rag CLI - SQLite-based RAG system"""
+    """haiku.rag CLI - Vector database RAG system"""
     # Run version check before any command
     asyncio.run(check_version())
 
@@ -55,9 +56,9 @@ def main(
 @cli.command("list", help="List all stored documents")
 def list_documents(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -70,9 +71,9 @@ def add_document_text(
         help="The text content of the document to add",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -85,9 +86,9 @@ def add_document_src(
         help="The file path or URL of the document to add",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -96,13 +97,13 @@ def add_document_src(
 
 @cli.command("get", help="Get and display a document by its ID")
 def get_document(
-    doc_id: int = typer.Argument(
+    doc_id: str = typer.Argument(
         help="The ID of the document to get",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -111,13 +112,13 @@ def get_document(
 
 @cli.command("delete", help="Delete a document by its ID")
 def delete_document(
-    doc_id: int = typer.Argument(
+    doc_id: str = typer.Argument(
         help="The ID of the document to delete",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -135,19 +136,14 @@ def search(
         "-l",
         help="Maximum number of results to return",
     ),
-    k: int = typer.Option(
-        60,
-        "--k",
-        help="Reciprocal Rank Fusion k parameter",
-    ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.search(query=query, limit=limit, k=k))
+    asyncio.run(app.search(query=query, limit=limit))
 
 
 @cli.command("ask", help="Ask a question using the QA agent")
@@ -156,9 +152,9 @@ def ask(
         help="The question to ask",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
     cite: bool = typer.Option(
         False,
@@ -182,9 +178,9 @@ def settings():
 )
 def rebuild(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
 ):
     app = HaikuRAGApp(db_path=db)
@@ -196,9 +192,9 @@ def rebuild(
 )
 def serve(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.sqlite",
+        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
-        help="Path to the SQLite database file",
+        help="Path to the LanceDB database file",
     ),
     stdio: bool = typer.Option(
         False,
@@ -227,5 +223,20 @@ def serve(
     asyncio.run(app.serve(transport=transport))
 
 
+@cli.command("migrate", help="Migrate an SQLite database to LanceDB")
+def migrate(
+    sqlite_path: Path = typer.Argument(
+        help="Path to the SQLite database file to migrate",
+    ),
+):
+    # Generate LanceDB path in same parent directory
+    lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
+
+    success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path))
+
+    if not success:
+        raise typer.Exit(1)
+
+
 if __name__ == "__main__":
     cli()

haiku/rag/client.py

@@ -3,7 +3,6 @@ import mimetypes
 import tempfile
 from collections.abc import AsyncGenerator
 from pathlib import Path
-from typing import Literal
 from urllib.parse import urlparse
 
 import httpx
@@ -16,6 +15,7 @@ from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.store.repositories.settings import SettingsRepository
 from haiku.rag.utils import text_to_docling_document
 
 
@@ -24,19 +24,17 @@ class HaikuRAG:
 
     def __init__(
         self,
-        db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
-        / "haiku.rag.sqlite",
+        db_path: Path = Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         skip_validation: bool = False,
     ):
         """Initialize the RAG client with a database path.
 
         Args:
-            db_path: Path to the SQLite database file or ":memory:" for in-memory database.
+            db_path: Path to the database file.
             skip_validation: Whether to skip configuration validation on database load.
         """
-        if isinstance(db_path, Path):
-            if not db_path.parent.exists():
-                Path.mkdir(db_path.parent, parents=True)
+        if not db_path.parent.exists():
+            Path.mkdir(db_path.parent, parents=True)
         self.store = Store(db_path, skip_validation=skip_validation)
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
@@ -269,7 +267,7 @@ class HaikuRAG:
         # Default to .html for web content
         return ".html"
 
-    async def get_document_by_id(self, document_id: int) -> Document | None:
+    async def get_document_by_id(self, document_id: str) -> Document | None:
         """Get a document by its ID.
 
         Args:
@@ -300,7 +298,7 @@ class HaikuRAG:
             document, docling_document
         )
 
-    async def delete_document(self, document_id: int) -> bool:
+    async def delete_document(self, document_id: str) -> bool:
         """Delete a document by its ID."""
         return await self.document_repository.delete(document_id)
 
@@ -319,14 +317,14 @@ class HaikuRAG:
         return await self.document_repository.list_all(limit=limit, offset=offset)
 
     async def search(
-        self, query: str, limit: int = 5, k: int = 60
+        self, query: str, limit: int = 5, search_type: str = "hybrid"
     ) -> list[tuple[Chunk, float]]:
-        """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking.
+        """Search for relevant chunks using the specified search method with optional reranking.
 
         Args:
             query: The search query string.
             limit: Maximum number of results to return.
-            k: Parameter for Reciprocal Rank Fusion (default: 60).
+            search_type: Type of search - "vector", "fts", or "hybrid" (default).
 
         Returns:
             List of (chunk, score) tuples ordered by relevance.
@@ -335,12 +333,15 @@ class HaikuRAG:
         reranker = get_reranker()
 
         if reranker is None:
-            return await self.chunk_repository.search_chunks_hybrid(query, limit, k)
+            # No reranking - return direct search results
+            return await self.chunk_repository.search(query, limit, search_type)
 
         # Get more initial results (3X) for reranking
-        search_results = await self.chunk_repository.search_chunks_hybrid(
-            query, limit * 3, k
+        search_limit = limit * 3
+        search_results = await self.chunk_repository.search(
+            query, search_limit, search_type
         )
+
         # Apply reranking
         chunks = [chunk for chunk, _ in search_results]
         reranked_results = await reranker.rerank(query, chunks, top_n=limit)
@@ -493,7 +494,7 @@ class HaikuRAG:
         qa_agent = get_qa_agent(self, use_citations=cite)
         return await qa_agent.answer(question)
 
-    async def rebuild_database(self) -> AsyncGenerator[int, None]:
+    async def rebuild_database(self) -> AsyncGenerator[str, None]:
         """Rebuild the database by deleting all chunks and re-indexing all documents.
 
         For documents with URIs:
@@ -510,10 +511,8 @@ class HaikuRAG:
         self.store.recreate_embeddings_table()
 
         # Update settings to current config
-        from haiku.rag.store.repositories.settings import SettingsRepository
-
         settings_repo = SettingsRepository(self.store)
-        settings_repo.save()
+        settings_repo.save_current_settings()
 
         documents = await self.list_documents()
 
@@ -547,13 +546,10 @@ class HaikuRAG:
                 # Document without URI - re-create chunks from existing content
                 docling_document = text_to_docling_document(doc.content)
                 await self.chunk_repository.create_chunks_for_document(
-                    doc.id, docling_document, commit=False
+                    doc.id, docling_document
                 )
                 yield doc.id
 
-        if self.store._connection:
-            self.store._connection.commit()
-
     def close(self):
         """Close the underlying store connection."""
         self.store.close()

haiku/rag/config.py

@@ -12,6 +12,10 @@ load_dotenv()
 class AppConfig(BaseModel):
     ENV: str = "production"
 
+    LANCEDB_API_KEY: str = ""
+    LANCEDB_URI: str = ""
+    LANCEDB_REGION: str = ""
+
     DEFAULT_DATA_DIR: Path = get_default_data_dir()
     MONITOR_DIRECTORIES: list[Path] = []
 
@@ -19,8 +23,8 @@ class AppConfig(BaseModel):
     EMBEDDINGS_MODEL: str = "mxbai-embed-large"
     EMBEDDINGS_VECTOR_DIM: int = 1024
 
-    RERANK_PROVIDER: str = "ollama"
-    RERANK_MODEL: str = "qwen3"
+    RERANK_PROVIDER: str = ""
+    RERANK_MODEL: str = ""
 
     QA_PROVIDER: str = "ollama"
     QA_MODEL: str = "qwen3"

haiku/rag/logging.py

@@ -3,6 +3,10 @@ import logging
 from rich.console import Console
 from rich.logging import RichHandler
 
+logging.basicConfig(level=logging.DEBUG)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.getLogger("httpcore").setLevel(logging.WARNING)
+
 
 def get_logger() -> logging.Logger:
     logger = logging.getLogger("haiku.rag")

haiku/rag/mcp.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any
 
 from fastmcp import FastMCP
 from pydantic import BaseModel
@@ -8,13 +8,13 @@ from haiku.rag.client import HaikuRAG
 
 
 class SearchResult(BaseModel):
-    document_id: int
+    document_id: str
     content: str
     score: float
 
 
 class DocumentResult(BaseModel):
-    id: int | None
+    id: str | None
     content: str
     uri: str | None = None
     metadata: dict[str, Any] = {}
@@ -22,14 +22,14 @@ class DocumentResult(BaseModel):
     updated_at: str
 
 
-def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
+def create_mcp_server(db_path: Path) -> FastMCP:
     """Create an MCP server with the specified database path."""
     mcp = FastMCP("haiku-rag")
 
     @mcp.tool()
     async def add_document_from_file(
         file_path: str, metadata: dict[str, Any] | None = None
-    ) -> int | None:
+    ) -> str | None:
         """Add a document to the RAG system from a file path."""
         try:
             async with HaikuRAG(db_path) as rag:
@@ -43,7 +43,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
     @mcp.tool()
     async def add_document_from_url(
         url: str, metadata: dict[str, Any] | None = None
-    ) -> int | None:
+    ) -> str | None:
         """Add a document to the RAG system from a URL."""
         try:
             async with HaikuRAG(db_path) as rag:
@@ -55,7 +55,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
     @mcp.tool()
     async def add_document_from_text(
         content: str, uri: str | None = None, metadata: dict[str, Any] | None = None
-    ) -> int | None:
+    ) -> str | None:
         """Add a document to the RAG system from text content."""
         try:
             async with HaikuRAG(db_path) as rag:
@@ -73,6 +73,9 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
 
                 search_results = []
                 for chunk, score in results:
+                    assert chunk.document_id is not None, (
+                        "Chunk document_id should not be None in search results"
+                    )
                     search_results.append(
                         SearchResult(
                             document_id=chunk.document_id,
@@ -86,7 +89,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
             return []
 
     @mcp.tool()
-    async def get_document(document_id: int) -> DocumentResult | None:
+    async def get_document(document_id: str) -> DocumentResult | None:
         """Get a document by its ID."""
         try:
             async with HaikuRAG(db_path) as rag:
@@ -130,7 +133,7 @@ def create_mcp_server(db_path: Path | Literal[":memory:"]) -> FastMCP:
             return []
 
     @mcp.tool()
-    async def delete_document(document_id: int) -> bool:
+    async def delete_document(document_id: str) -> bool:
         """Delete a document by its ID."""
         try:
             async with HaikuRAG(db_path) as rag: