haiku.rag 0.10.2__py3-none-any.whl → 0.14.0__py3-none-any.whl

haiku/rag/embeddings/__init__.py

@@ -1,35 +0,0 @@
-from haiku.rag.config import Config
-from haiku.rag.embeddings.base import EmbedderBase
-from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
-
-
-def get_embedder() -> EmbedderBase:
-    """
-    Factory function to get the appropriate embedder based on the configuration.
-    """
-
-    if Config.EMBEDDINGS_PROVIDER == "ollama":
-        return OllamaEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
-
-    if Config.EMBEDDINGS_PROVIDER == "voyageai":
-        try:
-            from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
-        except ImportError:
-            raise ImportError(
-                "VoyageAI embedder requires the 'voyageai' package. "
-                "Please install haiku.rag with the 'voyageai' extra: "
-                "uv pip install haiku.rag[voyageai]"
-            )
-        return VoyageAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
-
-    if Config.EMBEDDINGS_PROVIDER == "openai":
-        from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
-
-        return OpenAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
-
-    if Config.EMBEDDINGS_PROVIDER == "vllm":
-        from haiku.rag.embeddings.vllm import Embedder as VllmEmbedder
-
-        return VllmEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
-
-    raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDINGS_PROVIDER}")

haiku/rag/embeddings/base.py

@@ -1,15 +0,0 @@
-from haiku.rag.config import Config
-
-
-class EmbedderBase:
-    _model: str = Config.EMBEDDINGS_MODEL
-    _vector_dim: int = Config.EMBEDDINGS_VECTOR_DIM
-
-    def __init__(self, model: str, vector_dim: int):
-        self._model = model
-        self._vector_dim = vector_dim
-
-    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-        raise NotImplementedError(
-            "Embedder is an abstract class. Please implement the embed method in a subclass."
-        )

haiku/rag/embeddings/ollama.py

@@ -1,17 +0,0 @@
-from openai import AsyncOpenAI
-
-from haiku.rag.config import Config
-from haiku.rag.embeddings.base import EmbedderBase
-
-
-class Embedder(EmbedderBase):
-    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-        client = AsyncOpenAI(base_url=f"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy")
-        response = await client.embeddings.create(
-            model=self._model,
-            input=text,
-        )
-        if isinstance(text, str):
-            return response.data[0].embedding
-        else:
-            return [item.embedding for item in response.data]

haiku/rag/embeddings/openai.py

@@ -1,16 +0,0 @@
-from openai import AsyncOpenAI
-
-from haiku.rag.embeddings.base import EmbedderBase
-
-
-class Embedder(EmbedderBase):
-    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-        client = AsyncOpenAI()
-        response = await client.embeddings.create(
-            model=self._model,
-            input=text,
-        )
-        if isinstance(text, str):
-            return response.data[0].embedding
-        else:
-            return [item.embedding for item in response.data]

haiku/rag/embeddings/vllm.py

@@ -1,19 +0,0 @@
-from openai import AsyncOpenAI
-
-from haiku.rag.config import Config
-from haiku.rag.embeddings.base import EmbedderBase
-
-
-class Embedder(EmbedderBase):
-    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-        client = AsyncOpenAI(
-            base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy"
-        )
-        response = await client.embeddings.create(
-            model=self._model,
-            input=text,
-        )
-        if isinstance(text, str):
-            return response.data[0].embedding
-        else:
-            return [item.embedding for item in response.data]

haiku/rag/embeddings/voyageai.py

@@ -1,17 +0,0 @@
-try:
-    from voyageai.client import Client  # type: ignore
-
-    from haiku.rag.embeddings.base import EmbedderBase
-
-    class Embedder(EmbedderBase):
-        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-            client = Client()
-            if isinstance(text, str):
-                res = client.embed([text], model=self._model, output_dtype="float")
-                return res.embeddings[0]  # type: ignore[return-value]
-            else:
-                res = client.embed(text, model=self._model, output_dtype="float")
-                return res.embeddings  # type: ignore[return-value]
-
-except ImportError:
-    pass

haiku/rag/logging.py

@@ -1,56 +0,0 @@
-import logging
-import warnings
-
-from rich.console import Console
-from rich.logging import RichHandler
-
-
-def get_logger() -> logging.Logger:
-    """Return the library logger configured with a Rich handler."""
-    logger = logging.getLogger("haiku.rag")
-
-    handler = RichHandler(
-        console=Console(stderr=True),
-        rich_tracebacks=True,
-    )
-    formatter = logging.Formatter("%(message)s")
-    handler.setFormatter(formatter)
-
-    logger.setLevel(logging.INFO)
-
-    # Remove any existing handlers to avoid duplicates on reconfiguration
-    for hdlr in logger.handlers[:]:
-        logger.removeHandler(hdlr)
-
-    logger.addHandler(handler)
-    # Do not let messages propagate to the root logger
-    logger.propagate = False
-    return logger
-
-
-def configure_cli_logging(level: int = logging.INFO) -> logging.Logger:
-    """Configure logging for CLI runs.
-
-    - Silence ALL non-haiku.rag loggers by detaching root handlers and setting
-      their level to ERROR.
-    - Attach a Rich handler only to the "haiku.rag" logger.
-    - Prevent propagation so only our logger prints in the CLI.
-    """
-    # Silence root logger completely
-    root = logging.getLogger()
-    for hdlr in root.handlers[:]:
-        root.removeHandler(hdlr)
-    root.setLevel(logging.ERROR)
-
-    # Optionally silence some commonly noisy libraries explicitly as a safeguard
-    for noisy in ("httpx", "httpcore", "docling", "urllib3", "asyncio"):
-        logging.getLogger(noisy).setLevel(logging.ERROR)
-        logging.getLogger(noisy).propagate = False
-
-    # Configure and return our app logger
-    logger = get_logger()
-    logger.setLevel(level)
-    logger.propagate = False
-
-    warnings.filterwarnings("ignore")
-    return logger

haiku/rag/mcp.py

@@ -1,156 +0,0 @@
-from pathlib import Path
-from typing import Any
-
-from fastmcp import FastMCP
-from pydantic import BaseModel
-
-from haiku.rag.client import HaikuRAG
-
-
-class SearchResult(BaseModel):
-    document_id: str
-    content: str
-    score: float
-
-
-class DocumentResult(BaseModel):
-    id: str | None
-    content: str
-    uri: str | None = None
-    title: str | None = None
-    metadata: dict[str, Any] = {}
-    created_at: str
-    updated_at: str
-
-
-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,
-        title: str | None = None,
-    ) -> str | None:
-        """Add a document to the RAG system from a file path."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document_from_source(
-                    Path(file_path), title=title, metadata=metadata or {}
-                )
-                return document.id
-        except Exception:
-            return None
-
-    @mcp.tool()
-    async def add_document_from_url(
-        url: str, metadata: dict[str, Any] | None = None, title: str | None = None
-    ) -> str | None:
-        """Add a document to the RAG system from a URL."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document_from_source(
-                    url, title=title, metadata=metadata or {}
-                )
-                return document.id
-        except Exception:
-            return None
-
-    @mcp.tool()
-    async def add_document_from_text(
-        content: str,
-        uri: str | None = None,
-        metadata: dict[str, Any] | None = None,
-        title: str | None = None,
-    ) -> str | None:
-        """Add a document to the RAG system from text content."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document(
-                    content, uri, title=title, metadata=metadata or {}
-                )
-                return document.id
-        except Exception:
-            return None
-
-    @mcp.tool()
-    async def search_documents(query: str, limit: int = 5) -> list[SearchResult]:
-        """Search the RAG system for documents using hybrid search (vector similarity + full-text search)."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                results = await rag.search(query, limit)
-
-                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,
-                            content=chunk.content,
-                            score=score,
-                        )
-                    )
-
-                return search_results
-        except Exception:
-            return []
-
-    @mcp.tool()
-    async def get_document(document_id: str) -> DocumentResult | None:
-        """Get a document by its ID."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                document = await rag.get_document_by_id(document_id)
-
-                if document is None:
-                    return None
-
-                return DocumentResult(
-                    id=document.id,
-                    content=document.content,
-                    uri=document.uri,
-                    title=document.title,
-                    metadata=document.metadata,
-                    created_at=str(document.created_at),
-                    updated_at=str(document.updated_at),
-                )
-        except Exception:
-            return None
-
-    @mcp.tool()
-    async def list_documents(
-        limit: int | None = None, offset: int | None = None
-    ) -> list[DocumentResult]:
-        """List all documents with optional pagination."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                documents = await rag.list_documents(limit, offset)
-
-                return [
-                    DocumentResult(
-                        id=doc.id,
-                        content=doc.content,
-                        uri=doc.uri,
-                        title=doc.title,
-                        metadata=doc.metadata,
-                        created_at=str(doc.created_at),
-                        updated_at=str(doc.updated_at),
-                    )
-                    for doc in documents
-                ]
-        except Exception:
-            return []
-
-    @mcp.tool()
-    async def delete_document(document_id: str) -> bool:
-        """Delete a document by its ID."""
-        try:
-            async with HaikuRAG(db_path) as rag:
-                return await rag.delete_document(document_id)
-        except Exception:
-            return False
-
-    return mcp

haiku/rag/migration.py

@@ -1,316 +0,0 @@
-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
-
-            # Load the sqlite-vec extension
-            try:
-                import sqlite_vec  # type: ignore
-
-                sqlite_conn.enable_load_extension(True)
-                sqlite_vec.load(sqlite_conn)
-                self.console.print("[cyan]Loaded sqlite-vec extension[/cyan]")
-            except Exception as e:
-                self.console.print(
-                    f"[yellow]Warning: Could not load sqlite-vec extension: {e}[/yellow]"
-                )
-                self.console.print(
-                    "[yellow]Install sqlite-vec with[/yellow]\n[green]uv pip install sqlite-vec [/green]"
-                )
-                exit(1)
-
-            # 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 and cleanup using centralized vacuum
-            self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
-            try:
-                lance_store.vacuum()
-                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 using the sqlite-vec virtual table
-        embeddings_map = {}
-        try:
-            # Use the virtual table to get embeddings properly
-            cursor.execute("""
-                SELECT chunk_id, embedding
-                FROM chunk_embeddings
-            """)
-
-            for row in cursor.fetchall():
-                chunk_id = row[0]
-                embedding_blob = row[1]
-                if embedding_blob and chunk_id not in embeddings_map:
-                    embeddings_map[chunk_id] = embedding_blob
-
-        except sqlite3.OperationalError as e:
-            self.console.print(
-                f"[yellow]Warning: Could not extract embeddings from virtual table: {e}[/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()