haiku.rag 0.1.0__py3-none-any.whl

haiku/rag/embeddings/voyageai.py

@@ -0,0 +1,17 @@
+try:
+    from voyageai.client import Client  # type: ignore
+
+    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 = Client()
+            res = client.embed([text], model=self._model, output_dtype="float")
+            return res.embeddings[0]  # type: ignore[return-value]
+
+except ImportError:
+    pass

haiku/rag/mcp.py

@@ -0,0 +1,141 @@
+from pathlib import Path
+from typing import Any, Literal
+
+from fastmcp import FastMCP
+from pydantic import BaseModel
+
+from haiku.rag.client import HaikuRAG
+
+
+class SearchResult(BaseModel):
+    document_id: int
+    content: str
+    score: float
+
+
+class DocumentResult(BaseModel):
+    id: int | None
+    content: str
+    uri: str | None = None
+    metadata: dict[str, Any] = {}
+    created_at: str
+    updated_at: str
+
+
+def create_mcp_server(db_path: Path | Literal[":memory:"]) -> 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:
+        """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), 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
+    ) -> int | 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, 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
+    ) -> int | 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, 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:
+                    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: int) -> 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,
+                    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,
+                        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: int) -> 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/reader.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import ClassVar
+
+from markitdown import MarkItDown
+
+
+class FileReader:
+    extensions: ClassVar[list[str]] = [
+        ".astro",
+        ".c",
+        ".cpp",
+        ".css",
+        ".csv",
+        ".docx",
+        ".go",
+        ".h",
+        ".hpp",
+        ".html",
+        ".java",
+        ".js",
+        ".json",
+        ".kt",
+        ".md",
+        ".mdx",
+        ".mjs",
+        ".mp3",
+        ".pdf",
+        ".php",
+        ".pptx",
+        ".py",
+        ".rb",
+        ".rs",
+        ".svelte",
+        ".swift",
+        ".ts",
+        ".tsx",
+        ".txt",
+        ".vue",
+        ".wav",
+        ".xml",
+        ".xlsx",
+        ".yaml",
+        ".yml",
+    ]
+
+    @staticmethod
+    def parse_file(path: Path) -> str:
+        try:
+            reader = MarkItDown()
+            return reader.convert(path).text_content
+        except Exception:
+            raise ValueError(f"Failed to parse file: {path}")

haiku/rag/store/__init__.py

@@ -0,0 +1,4 @@
+from .engine import Store
+from .models import Chunk, Document
+
+__all__ = ["Store", "Chunk", "Document"]

haiku/rag/store/engine.py

@@ -0,0 +1,80 @@
+import sqlite3
+import struct
+from pathlib import Path
+from typing import Literal
+
+import sqlite_vec
+
+from haiku.rag.embeddings import get_embedder
+
+
+class Store:
+    def __init__(self, db_path: Path | Literal[":memory:"]):
+        self.db_path: Path | Literal[":memory:"] = db_path
+        self._connection = self.create_db()
+
+    def create_db(self) -> sqlite3.Connection:
+        """Create the database and tables with sqlite-vec support for embeddings."""
+        db = sqlite3.connect(self.db_path)
+        db.enable_load_extension(True)
+        sqlite_vec.load(db)
+
+        # Create documents table
+        db.execute("""
+            CREATE TABLE IF NOT EXISTS documents (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                content TEXT NOT NULL,
+                uri TEXT,
+                metadata TEXT DEFAULT '{}',
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+
+        # Create chunks table
+        db.execute("""
+            CREATE TABLE IF NOT EXISTS chunks (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                document_id INTEGER NOT NULL,
+                content TEXT NOT NULL,
+                metadata TEXT DEFAULT '{}',
+                FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE
+            )
+        """)
+
+        # Create vector table for chunk embeddings
+        embedder = get_embedder()
+        db.execute(f"""
+            CREATE VIRTUAL TABLE IF NOT EXISTS chunk_embeddings USING vec0(
+                chunk_id INTEGER PRIMARY KEY,
+                embedding FLOAT[{embedder._vector_dim}]
+            )
+        """)
+
+        # Create FTS5 table for full-text search
+        db.execute("""
+            CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
+                content,
+                content='chunks',
+                content_rowid='id'
+            )
+        """)
+
+        # Create indexes for better performance
+        db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
+        )
+
+        db.commit()
+        return db
+
+    @staticmethod
+    def serialize_embedding(embedding: list[float]) -> bytes:
+        """Serialize a list of floats to bytes for sqlite-vec storage."""
+        return struct.pack(f"{len(embedding)}f", *embedding)
+
+    def close(self):
+        """Close the database connection if it's an in-memory database."""
+        if self._connection is not None:
+            self._connection.close()
+            self._connection = None

haiku/rag/store/models/__init__.py

@@ -0,0 +1,4 @@
+from .chunk import Chunk
+from .document import Document
+
+__all__ = ["Chunk", "Document"]

haiku/rag/store/models/chunk.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+
+class Chunk(BaseModel):
+    """
+    Represents a document with an ID, content, and metadata.
+    """
+
+    id: int | None = None
+    document_id: int
+    content: str
+    metadata: dict = {}

haiku/rag/store/models/document.py

@@ -0,0 +1,16 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+
+class Document(BaseModel):
+    """
+    Represents a document with an ID, content, and metadata.
+    """
+
+    id: int | None = None
+    content: str
+    uri: str | None = None
+    metadata: dict = {}
+    created_at: datetime = Field(default_factory=datetime.now)
+    updated_at: datetime = Field(default_factory=datetime.now)

haiku/rag/store/repositories/__init__.py

@@ -0,0 +1,5 @@
+from haiku.rag.store.repositories.base import BaseRepository
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+
+__all__ = ["BaseRepository", "DocumentRepository", "ChunkRepository"]

haiku/rag/store/repositories/base.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from haiku.rag.store.engine import Store
+
+T = TypeVar("T")
+
+
+class BaseRepository(ABC, Generic[T]):
+    """Base repository interface for database operations."""
+
+    def __init__(self, store: Store):
+        self.store = store
+
+    @abstractmethod
+    async def create(self, entity: T) -> T:
+        """Create a new entity in the database."""
+        pass
+
+    @abstractmethod
+    async def get_by_id(self, entity_id: int) -> T | None:
+        """Get an entity by its ID."""
+        pass
+
+    @abstractmethod
+    async def update(self, entity: T) -> T:
+        """Update an existing entity."""
+        pass
+
+    @abstractmethod
+    async def delete(self, entity_id: int) -> bool:
+        """Delete an entity by its ID."""
+        pass
+
+    @abstractmethod
+    async def list_all(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[T]:
+        """List all entities with optional pagination."""
+        pass