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

haiku/rag/store/engine.py

@@ -1,230 +0,0 @@
-import json
-import logging
-from datetime import timedelta
-from importlib import metadata
-from pathlib import Path
-from uuid import uuid4
-
-import lancedb
-from lancedb.pydantic import LanceModel, Vector
-from pydantic import Field
-
-from haiku.rag.config import Config
-from haiku.rag.embeddings import get_embedder
-
-logger = logging.getLogger(__name__)
-
-
-class DocumentRecord(LanceModel):
-    id: str = Field(default_factory=lambda: str(uuid4()))
-    content: str
-    uri: str | None = None
-    metadata: str = Field(default="{}")
-    created_at: str = Field(default_factory=lambda: "")
-    updated_at: str = Field(default_factory=lambda: "")
-
-
-def create_chunk_model(vector_dim: int):
-    """Create a ChunkRecord model with the specified vector dimension.
-
-    This creates a model with proper vector typing for LanceDB.
-    """
-
-    class ChunkRecord(LanceModel):
-        id: str = Field(default_factory=lambda: str(uuid4()))
-        document_id: str
-        content: str
-        metadata: str = Field(default="{}")
-        vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
-
-    return ChunkRecord
-
-
-class SettingsRecord(LanceModel):
-    id: str = Field(default="settings")
-    settings: str = Field(default="{}")
-
-
-class Store:
-    def __init__(self, db_path: Path, skip_validation: bool = False):
-        self.db_path: Path = db_path
-        self.embedder = get_embedder()
-
-        # Create the ChunkRecord model with the correct vector dimension
-        self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
-
-        # Connect to LanceDB
-        self.db = self._connect_to_lancedb(db_path)
-
-        # Initialize tables
-        self.create_or_update_db()
-
-        # Validate config compatibility after connection is established
-        if not skip_validation:
-            self._validate_configuration()
-
-    def vacuum(self) -> None:
-        """Optimize and clean up old versions across all tables to reduce disk usage."""
-        if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
-            return
-
-        # Perform maintenance per table using optimize() with cleanup_older_than 0
-        for table in [self.documents_table, self.chunks_table, self.settings_table]:
-            table.optimize(cleanup_older_than=timedelta(0))
-
-    def _connect_to_lancedb(self, db_path: Path):
-        """Establish connection to LanceDB (local, cloud, or object storage)."""
-        # Check if we have cloud configuration
-        if self._has_cloud_config():
-            return lancedb.connect(
-                uri=Config.LANCEDB_URI,
-                api_key=Config.LANCEDB_API_KEY,
-                region=Config.LANCEDB_REGION,
-            )
-        else:
-            # Local file system connection
-            return lancedb.connect(db_path)
-
-    def _has_cloud_config(self) -> bool:
-        """Check if cloud configuration is complete."""
-        return bool(
-            Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION
-        )
-
-    def _validate_configuration(self) -> None:
-        """Validate that the configuration is compatible with the database."""
-        from haiku.rag.store.repositories.settings import SettingsRepository
-
-        settings_repo = SettingsRepository(self)
-        settings_repo.validate_config_compatibility()
-
-    def create_or_update_db(self):
-        """Create the database tables."""
-
-        # Get list of existing tables
-        existing_tables = self.db.table_names()
-
-        # Create or get documents table
-        if "documents" in existing_tables:
-            self.documents_table = self.db.open_table("documents")
-        else:
-            self.documents_table = self.db.create_table(
-                "documents", schema=DocumentRecord
-            )
-
-        # Create or get chunks table
-        if "chunks" in existing_tables:
-            self.chunks_table = self.db.open_table("chunks")
-        else:
-            self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
-            # Create FTS index on the new table
-            self.chunks_table.create_fts_index("content", replace=True)
-
-        # Create or get settings table
-        if "settings" in existing_tables:
-            self.settings_table = self.db.open_table("settings")
-        else:
-            self.settings_table = self.db.create_table(
-                "settings", schema=SettingsRecord
-            )
-            # Save current settings to the new database
-            settings_data = Config.model_dump(mode="json")
-            self.settings_table.add(
-                [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
-            )
-
-        # Set current version in settings
-        current_version = metadata.version("haiku.rag")
-        self.set_haiku_version(current_version)
-
-        # Check if we need to perform upgrades
-        try:
-            existing_settings = list(
-                self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
-            )
-            if existing_settings:
-                db_version = self.get_haiku_version()  # noqa: F841
-                # TODO: Add upgrade logic here similar to SQLite version when needed
-        except Exception:
-            # Settings table might not exist yet in fresh databases
-            pass
-
-    def get_haiku_version(self) -> str:
-        """Returns the user version stored in settings."""
-        settings_records = list(
-            self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
-        )
-        if settings_records:
-            settings = (
-                json.loads(settings_records[0].settings)
-                if settings_records[0].settings
-                else {}
-            )
-            return settings.get("version", "0.0.0")
-        return "0.0.0"
-
-    def set_haiku_version(self, version: str) -> None:
-        """Updates the user version in settings."""
-        settings_records = list(
-            self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
-        )
-        if settings_records:
-            # Only write if version actually changes to avoid creating new table versions
-            current = (
-                json.loads(settings_records[0].settings)
-                if settings_records[0].settings
-                else {}
-            )
-            if current.get("version") != version:
-                current["version"] = version
-                self.settings_table.update(
-                    where="id = 'settings'",
-                    values={"settings": json.dumps(current)},
-                )
-        else:
-            # Create new settings record
-            settings_data = Config.model_dump(mode="json")
-            settings_data["version"] = version
-            self.settings_table.add(
-                [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
-            )
-
-    def recreate_embeddings_table(self) -> None:
-        """Recreate the chunks table with current vector dimensions."""
-        # Drop and recreate chunks table
-        try:
-            self.db.drop_table("chunks")
-        except Exception:
-            pass
-
-        # Update the ChunkRecord model with new vector dimension
-        self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
-        self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
-
-        # Create FTS index on the new table
-        self.chunks_table.create_fts_index("content", replace=True)
-
-    def close(self):
-        """Close the database connection."""
-        # LanceDB connections are automatically managed
-        pass
-
-    def current_table_versions(self) -> dict[str, int]:
-        """Capture current versions of key tables for rollback using LanceDB's API."""
-        return {
-            "documents": int(self.documents_table.version),
-            "chunks": int(self.chunks_table.version),
-            "settings": int(self.settings_table.version),
-        }
-
-    def restore_table_versions(self, versions: dict[str, int]) -> bool:
-        """Restore tables to the provided versions using LanceDB's API."""
-        self.documents_table.restore(int(versions["documents"]))
-        self.chunks_table.restore(int(versions["chunks"]))
-        self.settings_table.restore(int(versions["settings"]))
-        return True
-
-    @property
-    def _connection(self):
-        """Compatibility property for repositories expecting _connection."""
-        return self

haiku/rag/store/models/__init__.py

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

haiku/rag/store/models/chunk.py

@@ -1,15 +0,0 @@
-from pydantic import BaseModel
-
-
-class Chunk(BaseModel):
-    """
-    Represents a chunk with content, metadata, and optional document information.
-    """
-
-    id: str | None = None
-    document_id: str | None = None
-    content: str
-    metadata: dict = {}
-    document_uri: str | None = None
-    document_meta: dict = {}
-    embedding: list[float] | None = None

haiku/rag/store/models/document.py

@@ -1,16 +0,0 @@
-from datetime import datetime
-
-from pydantic import BaseModel, Field
-
-
-class Document(BaseModel):
-    """
-    Represents a document with an ID, content, and metadata.
-    """
-
-    id: str | 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

@@ -1,9 +0,0 @@
-from haiku.rag.store.repositories.chunk import ChunkRepository
-from haiku.rag.store.repositories.document import DocumentRepository
-from haiku.rag.store.repositories.settings import SettingsRepository
-
-__all__ = [
-    "ChunkRepository",
-    "DocumentRepository",
-    "SettingsRepository",
-]

haiku/rag/store/repositories/chunk.py

@@ -1,399 +0,0 @@
-import asyncio
-import inspect
-import json
-import logging
-from uuid import uuid4
-
-from docling_core.types.doc.document import DoclingDocument
-from lancedb.rerankers import RRFReranker
-
-from haiku.rag.chunker import chunker
-from haiku.rag.config import Config
-from haiku.rag.embeddings import get_embedder
-from haiku.rag.store.engine import DocumentRecord, Store
-from haiku.rag.store.models.chunk import Chunk
-from haiku.rag.utils import load_callable, text_to_docling_document
-
-logger = logging.getLogger(__name__)
-
-
-class ChunkRepository:
-    """Repository for Chunk operations."""
-
-    def __init__(self, store: Store) -> None:
-        self.store = store
-        self.embedder = get_embedder()
-        self._optimize_lock = asyncio.Lock()
-
-    def _ensure_fts_index(self) -> None:
-        """Ensure FTS index exists on the content column."""
-        try:
-            self.store.chunks_table.create_fts_index("content", replace=True)
-        except Exception as e:
-            # Log the error but don't fail - FTS might already exist
-            logger.debug(f"FTS index creation skipped: {e}")
-
-    async def _optimize(self) -> None:
-        """Optimize the chunks table to refresh indexes."""
-        # Skip optimization for LanceDB Cloud as it handles this automatically
-        if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
-            return
-
-        async with self._optimize_lock:
-            try:
-                self.store.chunks_table.optimize()
-            except (RuntimeError, OSError) as e:
-                # Handle "too many open files" and other resource errors gracefully
-                logger.debug(
-                    f"Table optimization skipped due to resource constraints: {e}"
-                )
-
-    async def create(self, entity: Chunk) -> Chunk:
-        """Create a chunk in the database."""
-        assert entity.document_id, "Chunk must have a document_id to be created"
-
-        chunk_id = str(uuid4())
-
-        # Generate embedding if not provided
-        if entity.embedding is not None:
-            embedding = entity.embedding
-        else:
-            embedding = await self.embedder.embed(entity.content)
-        chunk_record = self.store.ChunkRecord(
-            id=chunk_id,
-            document_id=entity.document_id,
-            content=entity.content,
-            metadata=json.dumps(entity.metadata),
-            vector=embedding,
-        )
-
-        self.store.chunks_table.add([chunk_record])
-
-        entity.id = chunk_id
-
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
-        return entity
-
-    async def get_by_id(self, entity_id: str) -> Chunk | None:
-        """Get a chunk by its ID."""
-        results = list(
-            self.store.chunks_table.search()
-            .where(f"id = '{entity_id}'")
-            .limit(1)
-            .to_pydantic(self.store.ChunkRecord)
-        )
-
-        if not results:
-            return None
-
-        chunk_record = results[0]
-        return Chunk(
-            id=chunk_record.id,
-            document_id=chunk_record.document_id,
-            content=chunk_record.content,
-            metadata=json.loads(chunk_record.metadata) if chunk_record.metadata else {},
-        )
-
-    async def update(self, entity: Chunk) -> Chunk:
-        """Update an existing chunk."""
-        assert entity.id, "Chunk ID is required for update"
-
-        embedding = await self.embedder.embed(entity.content)
-
-        self.store.chunks_table.update(
-            where=f"id = '{entity.id}'",
-            values={
-                "document_id": entity.document_id,
-                "content": entity.content,
-                "metadata": json.dumps(entity.metadata),
-                "vector": embedding,
-            },
-        )
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
-        return entity
-
-    async def delete(self, entity_id: str) -> bool:
-        """Delete a chunk by its ID."""
-        chunk = await self.get_by_id(entity_id)
-        if chunk is None:
-            return False
-
-        self.store.chunks_table.delete(f"id = '{entity_id}'")
-        return True
-
-    async def list_all(
-        self, limit: int | None = None, offset: int | None = None
-    ) -> list[Chunk]:
-        """List all chunks with optional pagination."""
-        query = self.store.chunks_table.search()
-
-        if offset is not None:
-            query = query.offset(offset)
-        if limit is not None:
-            query = query.limit(limit)
-
-        results = list(query.to_pydantic(self.store.ChunkRecord))
-
-        return [
-            Chunk(
-                id=chunk.id,
-                document_id=chunk.document_id,
-                content=chunk.content,
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
-            )
-            for chunk in results
-        ]
-
-    async def create_chunks_for_document(
-        self, document_id: str, document: DoclingDocument
-    ) -> list[Chunk]:
-        """Create chunks and embeddings for a document from DoclingDocument."""
-        # Optionally preprocess markdown before chunking
-        processed_document = document
-        preprocessor_path = Config.MARKDOWN_PREPROCESSOR
-        if preprocessor_path:
-            try:
-                pre_fn = load_callable(preprocessor_path)
-                markdown = document.export_to_markdown()
-                result = pre_fn(markdown)
-                if inspect.isawaitable(result):
-                    result = await result  # type: ignore[assignment]
-                processed_markdown = result
-                if not isinstance(processed_markdown, str):
-                    raise ValueError("Preprocessor must return a markdown string")
-                processed_document = text_to_docling_document(
-                    processed_markdown, name="content.md"
-                )
-            except Exception as e:
-                logger.error(
-                    f"Failed to apply MARKDOWN_PREPROCESSOR '{preprocessor_path}': {e}. Proceeding without preprocessing."
-                )
-                raise e
-
-        chunk_texts = await chunker.chunk(processed_document)
-
-        embeddings = await self.embedder.embed(chunk_texts)
-
-        # Prepare all chunk records for batch insertion
-        chunk_records = []
-        created_chunks = []
-
-        for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)):
-            chunk_id = str(uuid4())
-
-            chunk_record = self.store.ChunkRecord(
-                id=chunk_id,
-                document_id=document_id,
-                content=chunk_text,
-                metadata=json.dumps({"order": order}),
-                vector=embedding,
-            )
-            chunk_records.append(chunk_record)
-
-            chunk = Chunk(
-                id=chunk_id,
-                document_id=document_id,
-                content=chunk_text,
-                metadata={"order": order},
-            )
-            created_chunks.append(chunk)
-
-        # Batch insert all chunks at once
-        if chunk_records:
-            self.store.chunks_table.add(chunk_records)
-
-        # Force optimization once at the end for bulk operations
-        await self._optimize()
-        return created_chunks
-
-    async def delete_all(self) -> None:
-        """Delete all chunks from the database."""
-        # Drop and recreate table to clear all data
-        self.store.db.drop_table("chunks")
-        self.store.chunks_table = self.store.db.create_table(
-            "chunks", schema=self.store.ChunkRecord
-        )
-        # Create FTS index on the new table
-        self.store.chunks_table.create_fts_index("content", replace=True)
-
-    async def delete_by_document_id(self, document_id: str) -> bool:
-        """Delete all chunks for a document."""
-        chunks = await self.get_by_document_id(document_id)
-
-        if not chunks:
-            return False
-
-        self.store.chunks_table.delete(f"document_id = '{document_id}'")
-        return True
-
-    async def search(
-        self, query: str, limit: int = 5, search_type: str = "hybrid"
-    ) -> list[tuple[Chunk, float]]:
-        """Search for relevant chunks using the specified search method.
-
-        Args:
-            query: The search query string.
-            limit: Maximum number of results to return.
-            search_type: Type of search - "vector", "fts", or "hybrid" (default).
-
-        Returns:
-            List of (chunk, score) tuples ordered by relevance.
-        """
-        if not query.strip():
-            return []
-
-        if search_type == "vector":
-            query_embedding = await self.embedder.embed(query)
-
-            results = self.store.chunks_table.search(
-                query_embedding, query_type="vector", vector_column_name="vector"
-            ).limit(limit)
-
-            return await self._process_search_results(results)
-
-        elif search_type == "fts":
-            results = self.store.chunks_table.search(query, query_type="fts").limit(
-                limit
-            )
-            return await self._process_search_results(results)
-
-        else:  # hybrid (default)
-            query_embedding = await self.embedder.embed(query)
-
-            # Create RRF reranker
-            reranker = RRFReranker()
-
-            # Perform native hybrid search with RRF reranking
-            results = (
-                self.store.chunks_table.search(query_type="hybrid")
-                .vector(query_embedding)
-                .text(query)
-                .rerank(reranker)
-                .limit(limit)
-            )
-            return await self._process_search_results(results)
-
-    async def get_by_document_id(self, document_id: str) -> list[Chunk]:
-        """Get all chunks for a specific document."""
-        results = list(
-            self.store.chunks_table.search()
-            .where(f"document_id = '{document_id}'")
-            .to_pydantic(self.store.ChunkRecord)
-        )
-
-        # Get document info
-        doc_results = list(
-            self.store.documents_table.search()
-            .where(f"id = '{document_id}'")
-            .limit(1)
-            .to_pydantic(DocumentRecord)
-        )
-
-        doc_uri = doc_results[0].uri if doc_results else None
-        doc_meta = doc_results[0].metadata if doc_results else "{}"
-
-        # Sort by order in metadata
-        chunks = [
-            Chunk(
-                id=chunk.id,
-                document_id=chunk.document_id,
-                content=chunk.content,
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
-                document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
-            )
-            for chunk in results
-        ]
-
-        chunks.sort(key=lambda c: c.metadata.get("order", 0))
-        return chunks
-
-    async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
-        """Get adjacent chunks before and after the given chunk within the same document."""
-        assert chunk.document_id, "Document id is required for adjacent chunk finding"
-
-        chunk_order = chunk.metadata.get("order")
-        if chunk_order is None:
-            return []
-
-        # Get all chunks for the document
-        all_chunks = await self.get_by_document_id(chunk.document_id)
-
-        # Filter to adjacent chunks
-        adjacent_chunks = []
-        for c in all_chunks:
-            c_order = c.metadata.get("order", 0)
-            if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
-                adjacent_chunks.append(c)
-
-        return adjacent_chunks
-
-    async def _process_search_results(self, query_result) -> list[tuple[Chunk, float]]:
-        """Process search results into chunks with document info and scores."""
-        chunks_with_scores = []
-
-        # Get both arrow and pydantic results to access scores
-        arrow_result = query_result.to_arrow()
-        pydantic_results = list(query_result.to_pydantic(self.store.ChunkRecord))
-
-        # Extract scores from arrow result based on search type
-        scores = []
-        column_names = arrow_result.column_names
-
-        if "_distance" in column_names:
-            # Vector search - distance (lower is better, convert to similarity)
-            distances = arrow_result.column("_distance").to_pylist()
-            scores = [max(0.0, 1.0 / (1.0 + dist)) for dist in distances]
-        elif "_relevance_score" in column_names:
-            # Hybrid search - relevance score (higher is better)
-            scores = arrow_result.column("_relevance_score").to_pylist()
-        elif "_score" in column_names:
-            # FTS search - score (higher is better)
-            scores = arrow_result.column("_score").to_pylist()
-        else:
-            raise ValueError("Unknown search result format, cannot extract scores")
-
-        # Collect all unique document IDs for batch lookup
-        document_ids = list(set(chunk.document_id for chunk in pydantic_results))
-
-        # Batch fetch all documents at once
-        documents_map = {}
-        if document_ids:
-            # Create a WHERE clause for all document IDs
-            where_clause = " OR ".join(f"id = '{doc_id}'" for doc_id in document_ids)
-            doc_results = list(
-                self.store.documents_table.search()
-                .where(where_clause)
-                .to_pydantic(DocumentRecord)
-            )
-            documents_map = {doc.id: doc for doc in doc_results}
-
-        for i, chunk_record in enumerate(pydantic_results):
-            # Get document info from pre-fetched map
-            doc = documents_map.get(chunk_record.document_id)
-            doc_uri = doc.uri if doc else None
-            doc_meta = doc.metadata if doc else "{}"
-
-            chunk = Chunk(
-                id=chunk_record.id,
-                document_id=chunk_record.document_id,
-                content=chunk_record.content,
-                metadata=json.loads(chunk_record.metadata)
-                if chunk_record.metadata
-                else {},
-                document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
-            )
-
-            # Get score from arrow result
-            score = scores[i] if i < len(scores) else 1.0
-
-            chunks_with_scores.append((chunk, score))
-
-        return chunks_with_scores