haiku.rag 0.1.0__py3-none-any.whl

haiku/rag/app.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+
+from rich.console import Console
+from rich.markdown import Markdown
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.store.models.document import Document
+
+
+class HaikuRAGApp:
+    def __init__(self, db_path: Path):
+        self.db_path = db_path
+        self.console = Console()
+
+    async def list_documents(self):
+        async with HaikuRAG(db_path=self.db_path) as self.client:
+            documents = await self.client.list_documents()
+            for doc in documents:
+                self._rich_print_document(doc, truncate=True)
+
+    async def add_document_from_text(self, text: str):
+        async with HaikuRAG(db_path=self.db_path) as self.client:
+            doc = await self.client.create_document(text)
+            self._rich_print_document(doc, truncate=True)
+            self.console.print(
+                f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
+            )
+
+    async def add_document_from_source(self, file_path: Path):
+        async with HaikuRAG(db_path=self.db_path) as self.client:
+            doc = await self.client.create_document_from_source(file_path)
+            self._rich_print_document(doc, truncate=True)
+            self.console.print(
+                f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
+            )
+
+    async def get_document(self, doc_id: int):
+        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:
+                self.console.print(f"[red]Document with id {doc_id} not found.[/red]")
+                return
+            self._rich_print_document(doc, truncate=False)
+
+    async def delete_document(self, doc_id: int):
+        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 with HaikuRAG(db_path=self.db_path) as self.client:
+            results = await self.client.search(query, limit=limit, k=k)
+            if not results:
+                self.console.print("[red]No results found.[/red]")
+                return
+            for chunk, score in results:
+                self._rich_print_search_result(chunk, score)
+
+    def _rich_print_document(self, doc: Document, truncate: bool = False):
+        """Format a document for display."""
+        if truncate:
+            content = doc.content.splitlines()
+            if len(content) > 3:
+                content = content[:3] + ["\n…"]
+            content = "\n".join(content)
+            content = Markdown(content)
+        else:
+            content = Markdown(doc.content)
+        self.console.print(
+            f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} [repr.attrib_name]uri[/repr.attrib_name]: {doc.uri} [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
+        )
+        self.console.print(
+            f"[repr.attrib_name]created at[/repr.attrib_name]: {doc.created_at} [repr.attrib_name]updated at[/repr.attrib_name]: {doc.updated_at}"
+        )
+        self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
+        self.console.print(content)
+        self.console.rule()
+
+    def _rich_print_search_result(self, chunk: Chunk, score: float):
+        """Format a search result chunk for display."""
+        content = Markdown(chunk.content)
+        self.console.print(
+            f"[repr.attrib_name]document_id[/repr.attrib_name]: {chunk.document_id} "
+            f"[repr.attrib_name]score[/repr.attrib_name]: {score:.4f}"
+        )
+        self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
+        self.console.print(content)
+        self.console.rule()
+
+    def serve(self, transport: str | None = None):
+        """Start the MCP server."""
+        from haiku.rag.mcp import create_mcp_server
+
+        server = create_mcp_server(self.db_path)
+
+        if transport == "stdio":
+            self.console.print("[green]Starting MCP server on stdio...[/green]")
+            server.run("stdio")
+        elif transport == "sse":
+            self.console.print(
+                "[green]Starting MCP server with streamable HTTP...[/green]"
+            )
+            server.run("sse")
+        else:
+            self.console.print("[green]Starting MCP server with HTTP...[/green]")
+            server.run("streamable-http")

haiku/rag/chunker.py

@@ -0,0 +1,76 @@
+from typing import ClassVar
+
+import tiktoken
+
+from haiku.rag.config import Config
+
+
+class Chunker:
+    """
+    A class that chunks text into smaller pieces for embedding and retrieval.
+
+    Parameters
+    ----------
+    chunk_size : int
+        The maximum size of a chunk in characters.
+    chunk_overlap : int
+        The number of characters of overlap between chunks.
+    """
+
+    encoder: ClassVar[tiktoken.Encoding] = tiktoken.encoding_for_model("gpt-4o")
+
+    def __init__(
+        self,
+        chunk_size: int = Config.CHUNK_SIZE,
+        chunk_overlap: int = Config.CHUNK_OVERLAP,
+    ):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    async def chunk(self, text: str) -> list[str]:
+        """
+        Split the text into chunks.
+
+        Parameters
+        ----------
+        text : str
+            The text to be split into chunks.
+
+        Returns
+        -------
+        list
+            A list of text chunks.
+        """
+        if not text:
+            return []
+
+        encoded_tokens = self.encoder.encode(text, disallowed_special=())
+
+        if self.chunk_size > len(encoded_tokens):
+            return [text]
+
+        chunks = []
+        i = 0
+        split_id_counter = 0
+        while i < len(encoded_tokens):
+            # Overlap
+            start_i = i
+            end_i = min(i + self.chunk_size, len(encoded_tokens))
+
+            chunk_tokens = encoded_tokens[start_i:end_i]
+            chunk_text = self.encoder.decode(chunk_tokens)
+
+            chunks.append(chunk_text)
+            split_id_counter += 1
+
+            # Exit loop if this was the last possible chunk
+            if end_i == len(encoded_tokens):
+                break
+
+            i += (
+                self.chunk_size - self.chunk_overlap
+            )  # Step forward, considering overlap
+        return chunks
+
+
+chunker = Chunker()

haiku/rag/cli.py

@@ -0,0 +1,153 @@
+import asyncio
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from haiku.rag.app import HaikuRAGApp
+from haiku.rag.utils import get_default_data_dir
+
+cli = typer.Typer(
+    context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
+)
+
+console = Console()
+event_loop = asyncio.get_event_loop()
+
+
+@cli.command("list", help="List all stored documents")
+def list_documents(
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.list_documents())
+
+
+@cli.command("add", help="Add a document from text input")
+def add_document_text(
+    text: str = typer.Argument(
+        help="The text content of the document to add",
+    ),
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.add_document_from_text(text=text))
+
+
+@cli.command("add-src", help="Add a document from a file path or URL")
+def add_document_src(
+    file_path: Path = typer.Argument(
+        help="The file path or URL of the document to add",
+    ),
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.add_document_from_source(file_path=file_path))
+
+
+@cli.command("get", help="Get and display a document by its ID")
+def get_document(
+    doc_id: int = typer.Argument(
+        help="The ID of the document to get",
+    ),
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.get_document(doc_id=doc_id))
+
+
+@cli.command("delete", help="Delete a document by its ID")
+def delete_document(
+    doc_id: int = typer.Argument(
+        help="The ID of the document to delete",
+    ),
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.delete_document(doc_id=doc_id))
+
+
+@cli.command("search", help="Search for documents by a query")
+def search(
+    query: str = typer.Argument(
+        help="The search query to use",
+    ),
+    limit: int = typer.Option(
+        5,
+        "--limit",
+        "-l",
+        help="Maximum number of results to return",
+    ),
+    k: int = typer.Option(
+        60,
+        "--k",
+        help="Reciprocal Rank Fusion k parameter",
+    ),
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+):
+    app = HaikuRAGApp(db_path=db)
+    event_loop.run_until_complete(app.search(query=query, limit=limit, k=k))
+
+
+@cli.command(
+    "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
+)
+def serve(
+    db: Path = typer.Option(
+        get_default_data_dir() / "haiku.rag.sqlite",
+        "--db",
+        help="Path to the SQLite database file",
+    ),
+    stdio: bool = typer.Option(
+        False,
+        "--stdio",
+        help="Run MCP server on stdio Transport",
+    ),
+    sse: bool = typer.Option(
+        False,
+        "--sse",
+        help="Run MCP server on SSE transport",
+    ),
+) -> None:
+    """Start the MCP server."""
+    if stdio and sse:
+        console.print("[red]Error: Cannot use both --stdio and --http options[/red]")
+        raise typer.Exit(1)
+
+    app = HaikuRAGApp(db_path=db)
+
+    transport = None
+    if stdio:
+        transport = "stdio"
+    elif sse:
+        transport = "sse"
+
+    app.serve(transport=transport)
+
+
+if __name__ == "__main__":
+    cli()

haiku/rag/client.py

@@ -0,0 +1,261 @@
+import hashlib
+import mimetypes
+import tempfile
+from pathlib import Path
+from typing import Literal
+from urllib.parse import urlparse
+
+import httpx
+
+from haiku.rag.config import Config
+from haiku.rag.reader import FileReader
+from haiku.rag.store.engine import Store
+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
+
+
+class HaikuRAG:
+    """High-level haiku-rag client."""
+
+    def __init__(
+        self,
+        db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
+        / "haiku.rag.sqlite",
+    ):
+        """Initialize the RAG client with a database path."""
+        if isinstance(db_path, Path):
+            if not db_path.parent.exists():
+                Path.mkdir(db_path.parent, parents=True)
+        self.store = Store(db_path)
+        self.document_repository = DocumentRepository(self.store)
+        self.chunk_repository = ChunkRepository(self.store)
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        self.close()
+        return False
+
+    async def create_document(
+        self, content: str, uri: str | None = None, metadata: dict | None = None
+    ) -> Document:
+        """Create a new document with optional URI and metadata."""
+        document = Document(
+            content=content,
+            uri=uri,
+            metadata=metadata or {},
+        )
+        return await self.document_repository.create(document)
+
+    async def create_document_from_source(
+        self, source: str | Path, metadata: dict = {}
+    ) -> Document:
+        """Create or update a document from a file path or URL.
+
+        Checks if a document with the same URI already exists:
+        - If MD5 is unchanged, returns existing document
+        - If MD5 changed, updates the document
+        - If no document exists, creates a new one
+
+        Args:
+            source: File path (as string or Path) or URL to parse
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Document instance (created, updated, or existing)
+
+        Raises:
+            ValueError: If the file/URL cannot be parsed or doesn't exist
+            httpx.RequestError: If URL request fails
+        """
+
+        # Check if it's a URL
+        source_str = str(source)
+        parsed_url = urlparse(source_str)
+        if parsed_url.scheme in ("http", "https"):
+            return await self._create_or_update_document_from_url(source_str, metadata)
+
+        # Handle as file path
+        source_path = Path(source) if isinstance(source, str) else source
+        if source_path.suffix.lower() not in FileReader.extensions:
+            raise ValueError(f"Unsupported file extension: {source_path.suffix}")
+
+        if not source_path.exists():
+            raise ValueError(f"File does not exist: {source_path}")
+
+        uri = str(source_path.resolve())
+        md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
+
+        # Check if document already exists
+        existing_doc = await self.get_document_by_uri(uri)
+        if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
+            # MD5 unchanged, return existing document
+            return existing_doc
+
+        content = FileReader.parse_file(source_path)
+
+        # Get content type from file extension
+        content_type, _ = mimetypes.guess_type(str(source_path))
+        if not content_type:
+            content_type = "application/octet-stream"
+
+        # Merge metadata with contentType and md5
+        metadata.update({"contentType": content_type, "md5": md5_hash})
+
+        if existing_doc:
+            # Update existing document
+            existing_doc.content = content
+            existing_doc.metadata = metadata
+            return await self.update_document(existing_doc)
+        else:
+            # Create new document
+            return await self.create_document(
+                content=content, uri=uri, metadata=metadata
+            )
+
+    async def _create_or_update_document_from_url(
+        self, url: str, metadata: dict = {}
+    ) -> Document:
+        """Create or update a document from a URL by downloading and parsing the content.
+
+        Checks if a document with the same URI already exists:
+        - If MD5 is unchanged, returns existing document
+        - If MD5 changed, updates the document
+        - If no document exists, creates a new one
+
+        Args:
+            url: URL to download and parse
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Document instance (created, updated, or existing)
+
+        Raises:
+            ValueError: If the content cannot be parsed
+            httpx.RequestError: If URL request fails
+        """
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url)
+            response.raise_for_status()
+
+            md5_hash = hashlib.md5(response.content).hexdigest()
+
+            # Check if document already exists
+            existing_doc = await self.get_document_by_uri(url)
+            if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
+                # MD5 unchanged, return existing document
+                return existing_doc
+
+            # Get content type to determine file extension
+            content_type = response.headers.get("content-type", "").lower()
+            file_extension = self._get_extension_from_content_type_or_url(
+                url, content_type
+            )
+
+            if file_extension not in FileReader.extensions:
+                raise ValueError(
+                    f"Unsupported content type/extension: {content_type}/{file_extension}"
+                )
+
+            # Create a temporary file with the appropriate extension
+            with tempfile.NamedTemporaryFile(
+                mode="wb", suffix=file_extension, delete=False
+            ) as temp_file:
+                temp_file.write(response.content)
+                temp_path = Path(temp_file.name)
+
+            try:
+                # Parse the content using FileReader
+                content = FileReader.parse_file(temp_path)
+
+                # Merge metadata with contentType and md5
+                metadata.update({"contentType": content_type, "md5": md5_hash})
+
+                if existing_doc:
+                    existing_doc.content = content
+                    existing_doc.metadata = metadata
+                    return await self.update_document(existing_doc)
+                else:
+                    return await self.create_document(
+                        content=content, uri=url, metadata=metadata
+                    )
+            finally:
+                # Clean up temporary file
+                temp_path.unlink(missing_ok=True)
+
+    def _get_extension_from_content_type_or_url(
+        self, url: str, content_type: str
+    ) -> str:
+        """Determine file extension from content type or URL."""
+        # Common content type mappings
+        content_type_map = {
+            "text/html": ".html",
+            "text/plain": ".txt",
+            "text/markdown": ".md",
+            "application/pdf": ".pdf",
+            "application/json": ".json",
+            "text/csv": ".csv",
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
+            "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
+        }
+
+        # Try content type first
+        for ct, ext in content_type_map.items():
+            if ct in content_type:
+                return ext
+
+        # Try URL extension
+        parsed_url = urlparse(url)
+        path = Path(parsed_url.path)
+        if path.suffix:
+            return path.suffix.lower()
+
+        # Default to .html for web content
+        return ".html"
+
+    async def get_document_by_id(self, document_id: int) -> Document | None:
+        """Get a document by its ID."""
+        return await self.document_repository.get_by_id(document_id)
+
+    async def get_document_by_uri(self, uri: str) -> Document | None:
+        """Get a document by its URI."""
+        return await self.document_repository.get_by_uri(uri)
+
+    async def update_document(self, document: Document) -> Document:
+        """Update an existing document."""
+        return await self.document_repository.update(document)
+
+    async def delete_document(self, document_id: int) -> bool:
+        """Delete a document by its ID."""
+        return await self.document_repository.delete(document_id)
+
+    async def list_documents(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[Document]:
+        """List all documents with optional pagination."""
+        return await self.document_repository.list_all(limit=limit, offset=offset)
+
+    async def search(
+        self, query: str, limit: int = 5, k: int = 60
+    ) -> list[tuple[Chunk, float]]:
+        """Search for relevant chunks using hybrid search (vector similarity + full-text search).
+
+        Args:
+            query: The search query string
+            limit: Maximum number of results to return
+            k: Parameter for Reciprocal Rank Fusion (default: 60)
+
+        Returns:
+            List of (chunk, score) tuples ordered by relevance
+        """
+        return await self.chunk_repository.search_chunks_hybrid(query, limit, k)
+
+    def close(self):
+        """Close the underlying store connection."""
+        self.store.close()

haiku/rag/config.py

@@ -0,0 +1,28 @@
+import os
+from pathlib import Path
+
+from dotenv import load_dotenv
+from pydantic import BaseModel
+
+from haiku.rag.utils import get_default_data_dir
+
+load_dotenv()
+
+
+class AppConfig(BaseModel):
+    ENV: str = "development"
+
+    DEFAULT_DATA_DIR: Path = get_default_data_dir()
+
+    EMBEDDING_PROVIDER: str = "ollama"
+    EMBEDDING_MODEL: str = "mxbai-embed-large"
+    EMBEDDING_VECTOR_DIM: int = 1024
+
+    CHUNK_SIZE: int = 256
+    CHUNK_OVERLAP: int = 32
+
+    OLLAMA_BASE_URL: str = "http://localhost:11434"
+
+
+# Expose Config object for app to import
+Config = AppConfig.model_validate(os.environ)

haiku/rag/embeddings/__init__.py

@@ -0,0 +1,24 @@
+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.EMBEDDING_PROVIDER == "ollama":
+        return OllamaEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM)
+
+    if Config.EMBEDDING_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 --extra voyageai"
+            )
+        return VoyageAIEmbedder(Config.EMBEDDING_MODEL, Config.EMBEDDING_VECTOR_DIM)
+    raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDING_PROVIDER}")

haiku/rag/embeddings/base.py

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

haiku/rag/embeddings/ollama.py

@@ -0,0 +1,14 @@
+from ollama import AsyncClient
+
+from haiku.rag.config import Config
+from haiku.rag.embeddings.base import EmbedderBase
+
+
+class Embedder(EmbedderBase):
+    _model: str = Config.EMBEDDING_MODEL
+    _vector_dim: int = 1024
+
+    async def embed(self, text: str) -> list[float]:
+        client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+        res = await client.embeddings(model=self._model, prompt=text)
+        return list(res["embedding"])