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

haiku/rag/migration.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""
+Migration script to migrate from SQLite to LanceDB.
+
+This script will:
+1. Read data from an existing SQLite database
+2. Create a new LanceDB database with the same data
+3. Preserve all documents, chunks, embeddings, and settings
+"""
+
+import json
+import sqlite3
+import struct
+from pathlib import Path
+from uuid import uuid4
+
+from rich.console import Console
+from rich.progress import Progress, TaskID
+
+from haiku.rag.store.engine import Store
+
+
+def deserialize_sqlite_embedding(data: bytes) -> list[float]:
+    """Deserialize sqlite-vec embedding from bytes."""
+    if not data:
+        return []
+    # sqlite-vec stores embeddings as float32 arrays
+    num_floats = len(data) // 4
+    return list(struct.unpack(f"{num_floats}f", data))
+
+
+class SQLiteToLanceDBMigrator:
+    """Migrates data from SQLite to LanceDB."""
+
+    def __init__(self, sqlite_path: Path, lancedb_path: Path):
+        self.sqlite_path = sqlite_path
+        self.lancedb_path = lancedb_path
+        self.console = Console()
+
+    def migrate(self) -> bool:
+        """Perform the migration."""
+        try:
+            self.console.print(
+                f"[blue]Starting migration from {self.sqlite_path} to {self.lancedb_path}[/blue]"
+            )
+
+            # Check if SQLite database exists
+            if not self.sqlite_path.exists():
+                self.console.print(
+                    f"[red]SQLite database not found: {self.sqlite_path}[/red]"
+                )
+                return False
+
+            # Connect to SQLite database
+            sqlite_conn = sqlite3.connect(self.sqlite_path)
+            sqlite_conn.row_factory = sqlite3.Row
+
+            # Create LanceDB store
+            lance_store = Store(self.lancedb_path, skip_validation=True)
+
+            with Progress() as progress:
+                # Migrate documents
+                doc_task = progress.add_task(
+                    "[green]Migrating documents...", total=None
+                )
+                document_id_mapping = self._migrate_documents(
+                    sqlite_conn, lance_store, progress, doc_task
+                )
+
+                # Migrate chunks and embeddings
+                chunk_task = progress.add_task(
+                    "[yellow]Migrating chunks and embeddings...", total=None
+                )
+                self._migrate_chunks(
+                    sqlite_conn, lance_store, progress, chunk_task, document_id_mapping
+                )
+
+                # Migrate settings
+                settings_task = progress.add_task(
+                    "[blue]Migrating settings...", total=None
+                )
+                self._migrate_settings(
+                    sqlite_conn, lance_store, progress, settings_task
+                )
+
+            sqlite_conn.close()
+
+            # Optimize the chunks table after migration
+            self.console.print("[blue]Optimizing LanceDB...[/blue]")
+            try:
+                lance_store.chunks_table.optimize()
+                self.console.print("[green]✅ Optimization completed[/green]")
+            except Exception as e:
+                self.console.print(
+                    f"[yellow]Warning: Optimization failed: {e}[/yellow]"
+                )
+
+            lance_store.close()
+
+            self.console.print("[green]✅ Migration completed successfully![/green]")
+            self.console.print(
+                f"[green]✅ Migrated {len(document_id_mapping)} documents[/green]"
+            )
+            return True
+
+        except Exception as e:
+            self.console.print(f"[red]❌ Migration failed: {e}[/red]")
+            import traceback
+
+            self.console.print(f"[red]{traceback.format_exc()}[/red]")
+            return False
+
+    def _migrate_documents(
+        self,
+        sqlite_conn: sqlite3.Connection,
+        lance_store: Store,
+        progress: Progress,
+        task: TaskID,
+    ) -> dict[int, str]:
+        """Migrate documents from SQLite to LanceDB and return ID mapping."""
+        cursor = sqlite_conn.cursor()
+        cursor.execute(
+            "SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY id"
+        )
+
+        documents = []
+        id_mapping = {}  # Maps old integer ID to new UUID
+
+        for row in cursor.fetchall():
+            new_uuid = str(uuid4())
+            id_mapping[row["id"]] = new_uuid
+
+            doc_data = {
+                "id": new_uuid,
+                "content": row["content"],
+                "uri": row["uri"],
+                "metadata": json.loads(row["metadata"]) if row["metadata"] else {},
+                "created_at": row["created_at"],
+                "updated_at": row["updated_at"],
+            }
+            documents.append(doc_data)
+
+        # Batch insert documents to LanceDB
+        if documents:
+            from haiku.rag.store.engine import DocumentRecord
+
+            doc_records = [
+                DocumentRecord(
+                    id=doc["id"],
+                    content=doc["content"],
+                    uri=doc["uri"],
+                    metadata=json.dumps(doc["metadata"]),
+                    created_at=doc["created_at"],
+                    updated_at=doc["updated_at"],
+                )
+                for doc in documents
+            ]
+            lance_store.documents_table.add(doc_records)
+
+        progress.update(task, completed=len(documents), total=len(documents))
+        return id_mapping
+
+    def _migrate_chunks(
+        self,
+        sqlite_conn: sqlite3.Connection,
+        lance_store: Store,
+        progress: Progress,
+        task: TaskID,
+        document_id_mapping: dict[int, str],
+    ):
+        """Migrate chunks and embeddings from SQLite to LanceDB."""
+        cursor = sqlite_conn.cursor()
+
+        # Get chunks first
+        cursor.execute("""
+            SELECT id, document_id, content, metadata
+            FROM chunks
+            ORDER BY id
+        """)
+
+        chunks_data = cursor.fetchall()
+
+        # Get embeddings separately to avoid vec0 virtual table issues
+        embeddings_map = {}
+        try:
+            # Try to get embeddings from the vec0 tables directly
+            cursor.execute("""
+                SELECT
+                    r.chunk_id,
+                    v.vectors
+                FROM chunk_embeddings_rowids r
+                JOIN chunk_embeddings_vector_chunks00 v ON r.rowid = v.rowid
+            """)
+
+            for row in cursor.fetchall():
+                chunk_id = row[0]
+                vectors_blob = row[1]
+                if vectors_blob and chunk_id not in embeddings_map:
+                    embeddings_map[chunk_id] = vectors_blob
+
+        except sqlite3.OperationalError as e:
+            self.console.print(
+                f"[yellow]Warning: Could not extract embeddings: {e}[/yellow]"
+            )
+            self.console.print(
+                "[yellow]Continuing migration without embeddings...[/yellow]"
+            )
+
+        chunks = []
+        for row in chunks_data:
+            # Generate new UUID for chunk
+            chunk_uuid = str(uuid4())
+
+            # Map the old document_id to new UUID
+            document_uuid = document_id_mapping.get(row["document_id"])
+            if not document_uuid:
+                self.console.print(
+                    f"[yellow]Warning: Document ID {row['document_id']} not found in mapping for chunk {row['id']}[/yellow]"
+                )
+                continue
+
+            # Get embedding for this chunk
+            embedding = []
+            embedding_blob = embeddings_map.get(row["id"])
+            if embedding_blob:
+                try:
+                    embedding = deserialize_sqlite_embedding(embedding_blob)
+                except Exception as e:
+                    self.console.print(
+                        f"[yellow]Warning: Failed to deserialize embedding for chunk {row['id']}: {e}[/yellow]"
+                    )
+                    # Generate a zero vector of the expected dimension
+                    embedding = [0.0] * lance_store.embedder._vector_dim
+            else:
+                # No embedding found, generate zero vector
+                embedding = [0.0] * lance_store.embedder._vector_dim
+
+            chunk_data = {
+                "id": chunk_uuid,
+                "document_id": document_uuid,
+                "content": row["content"],
+                "metadata": json.loads(row["metadata"]) if row["metadata"] else {},
+                "vector": embedding,
+            }
+            chunks.append(chunk_data)
+
+        # Batch insert chunks to LanceDB
+        if chunks:
+            chunk_records = [
+                lance_store.ChunkRecord(
+                    id=chunk["id"],
+                    document_id=chunk["document_id"],
+                    content=chunk["content"],
+                    metadata=json.dumps(chunk["metadata"]),
+                    vector=chunk["vector"],
+                )
+                for chunk in chunks
+            ]
+            lance_store.chunks_table.add(chunk_records)
+
+        progress.update(task, completed=len(chunks), total=len(chunks))
+
+    def _migrate_settings(
+        self,
+        sqlite_conn: sqlite3.Connection,
+        lance_store: Store,
+        progress: Progress,
+        task: TaskID,
+    ):
+        """Migrate settings from SQLite to LanceDB."""
+        cursor = sqlite_conn.cursor()
+
+        try:
+            cursor.execute("SELECT id, settings FROM settings WHERE id = 1")
+            row = cursor.fetchone()
+
+            if row:
+                settings_data = json.loads(row["settings"]) if row["settings"] else {}
+
+                # Update the existing settings in LanceDB (use string ID)
+                lance_store.settings_table.update(
+                    where="id = 'settings'",
+                    values={"settings": json.dumps(settings_data)},
+                )
+
+                progress.update(task, completed=1, total=1)
+            else:
+                progress.update(task, completed=0, total=0)
+
+        except sqlite3.OperationalError:
+            # Settings table doesn't exist in old SQLite database
+            self.console.print(
+                "[yellow]No settings table found in SQLite database[/yellow]"
+            )
+            progress.update(task, completed=0, total=0)
+
+
+async def migrate_sqlite_to_lancedb(
+    sqlite_path: Path, lancedb_path: Path | None = None
+) -> bool:
+    """
+    Migrate an existing SQLite database to LanceDB.
+
+    Args:
+        sqlite_path: Path to the existing SQLite database
+        lancedb_path: Path for the new LanceDB database (optional, will auto-generate if not provided)
+
+    Returns:
+        True if migration was successful, False otherwise
+    """
+    if lancedb_path is None:
+        # Auto-generate LanceDB path
+        lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
+
+    migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
+    return migrator.migrate()

haiku/rag/qa/__init__.py

@@ -1,44 +1,15 @@
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
-from haiku.rag.qa.base import QuestionAnswerAgentBase
-from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent
+from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(
-    client: HaikuRAG, model: str = "", use_citations: bool = False
-) -> QuestionAnswerAgentBase:
-    """
-    Factory function to get the appropriate QA agent based on the configuration.
-    """
-    if Config.QA_PROVIDER == "ollama":
-        return QuestionAnswerOllamaAgent(
-            client, model or Config.QA_MODEL, use_citations
-        )
+def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+    provider = Config.QA_PROVIDER
+    model_name = Config.QA_MODEL
 
-    if Config.QA_PROVIDER == "openai":
-        try:
-            from haiku.rag.qa.openai import QuestionAnswerOpenAIAgent
-        except ImportError:
-            raise ImportError(
-                "OpenAI QA agent requires the 'openai' package. "
-                "Please install haiku.rag with the 'openai' extra:"
-                "uv pip install haiku.rag[openai]"
-            )
-        return QuestionAnswerOpenAIAgent(
-            client, model or Config.QA_MODEL, use_citations
-        )
-
-    if Config.QA_PROVIDER == "anthropic":
-        try:
-            from haiku.rag.qa.anthropic import QuestionAnswerAnthropicAgent
-        except ImportError:
-            raise ImportError(
-                "Anthropic QA agent requires the 'anthropic' package. "
-                "Please install haiku.rag with the 'anthropic' extra:"
-                "uv pip install haiku.rag[anthropic]"
-            )
-        return QuestionAnswerAnthropicAgent(
-            client, model or Config.QA_MODEL, use_citations
-        )
-
-    raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}")
+    return QuestionAnswerAgent(
+        client=client,
+        provider=provider,
+        model=model_name,
+        use_citations=use_citations,
+    )

haiku/rag/qa/agent.py

@@ -0,0 +1,76 @@
+from pydantic import BaseModel, Field
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.ollama import OllamaProvider
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
+
+
+class SearchResult(BaseModel):
+    content: str = Field(description="The document text content")
+    score: float = Field(description="Relevance score (higher is more relevant)")
+    document_uri: str = Field(description="Source URI/path of the document")
+
+
+class Dependencies(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+    client: HaikuRAG
+
+
+class QuestionAnswerAgent:
+    def __init__(
+        self,
+        client: HaikuRAG,
+        provider: str,
+        model: str,
+        use_citations: bool = False,
+        q: float = 0.0,
+    ):
+        self._client = client
+
+        system_prompt = SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
+        model_obj = self._get_model(provider, model)
+
+        self._agent = Agent(
+            model=model_obj,
+            deps_type=Dependencies,
+            system_prompt=system_prompt,
+        )
+
+        @self._agent.tool
+        async def search_documents(
+            ctx: RunContext[Dependencies],
+            query: str,
+            limit: int = 3,
+        ) -> list[SearchResult]:
+            """Search the knowledge base for relevant documents."""
+            search_results = await ctx.deps.client.search(query, limit=limit)
+            expanded_results = await ctx.deps.client.expand_context(search_results)
+
+            return [
+                SearchResult(
+                    content=chunk.content,
+                    score=score,
+                    document_uri=chunk.document_uri or "",
+                )
+                for chunk, score in expanded_results
+            ]
+
+    def _get_model(self, provider: str, model: str):
+        """Get the appropriate model object for the provider."""
+        if provider == "ollama":
+            return OpenAIModel(
+                model_name=model,
+                provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+            )
+        else:
+            # For all other providers, use the provider:model format
+            return f"{provider}:{model}"
+
+    async def answer(self, question: str) -> str:
+        """Answer a question using the RAG system."""
+        deps = Dependencies(client=self._client)
+        result = await self._agent.run(question, deps=deps)
+        return result.output

haiku/rag/qa/prompts.py

@@ -18,6 +18,7 @@ Guidelines:
 - Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
 """
 
 SYSTEM_PROMPT_WITH_CITATIONS = """
@@ -55,4 +56,5 @@ Citations:
 - /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
 """

haiku/rag/reranking/__init__.py

@@ -31,10 +31,4 @@ def get_reranker() -> RerankerBase | None:
         except ImportError:
             return None
 
-    if Config.RERANK_PROVIDER == "ollama":
-        from haiku.rag.reranking.ollama import OllamaReranker
-
-        _reranker = OllamaReranker()
-        return _reranker
-
     return None