haiku.rag-slim 0.16.0__py3-none-any.whl

haiku/rag/store/repositories/chunk.py

@@ -0,0 +1,442 @@
+import inspect
+import json
+import logging
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from lancedb.query import LanceQueryBuilder
+
+from lancedb.rerankers import RRFReranker
+
+from haiku.rag.store.engine import DocumentRecord, Store
+from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.utils import load_callable
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+logger = logging.getLogger(__name__)
+
+
+class ChunkRepository:
+    """Repository for Chunk operations."""
+
+    def __init__(self, store: Store) -> None:
+        self.store = store
+        self.embedder = store.embedder
+
+    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, with_position=True, remove_stop_words=False
+            )
+        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 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)
+        order_val = int(entity.order)
+
+        chunk_record = self.store.ChunkRecord(
+            id=chunk_id,
+            document_id=entity.document_id,
+            content=entity.content,
+            metadata=json.dumps(
+                {k: v for k, v in entity.metadata.items() if k != "order"}
+            ),
+            order=order_val,
+            vector=embedding,
+        )
+
+        self.store.chunks_table.add([chunk_record])
+
+        entity.id = chunk_id
+        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]
+        md = json.loads(chunk_record.metadata)
+        return Chunk(
+            id=chunk_record.id,
+            document_id=chunk_record.document_id,
+            content=chunk_record.content,
+            metadata=md,
+            order=chunk_record.order,
+        )
+
+    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)
+        order_val = int(entity.order)
+
+        self.store.chunks_table.update(
+            where=f"id = '{entity.id}'",
+            values={
+                "document_id": entity.document_id,
+                "content": entity.content,
+                "metadata": json.dumps(
+                    {k: v for k, v in entity.metadata.items() if k != "order"}
+                ),
+                "order": order_val,
+                "vector": embedding,
+            },
+        )
+        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))
+
+        chunks: list[Chunk] = []
+        for rec in results:
+            md = json.loads(rec.metadata)
+            chunks.append(
+                Chunk(
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=md,
+                    order=rec.order,
+                )
+            )
+        return chunks
+
+    async def create_chunks_for_document(
+        self, document_id: str, document: "DoclingDocument"
+    ) -> list[Chunk]:
+        """Create chunks and embeddings for a document from DoclingDocument."""
+        # Lazy imports to avoid loading docling during module import
+        from haiku.rag.chunker import chunker
+        from haiku.rag.utils import text_to_docling_document
+
+        # Optionally preprocess markdown before chunking
+        processed_document = document
+        preprocessor_path = self.store._config.processing.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)
+
+        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 with phrase query support
+        self.store.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
+
+    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",
+        filter: str | None = None,
+    ) -> 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).
+            filter: Optional SQL WHERE clause to filter documents before searching chunks.
+
+        Returns:
+            List of (chunk, score) tuples ordered by relevance.
+        """
+        if not query.strip():
+            return []
+        filtered_doc_ids = None
+        if filter:
+            # We perform filtering as a two-step process, first filtering documents, then
+            # filtering chunks based on those document IDs.
+            # This is because LanceDB does not support joins directly in search queries.
+            docs_df = (
+                self.store.documents_table.search()
+                .select(["id"])
+                .where(filter)
+                .to_pandas()
+            )
+            # Early exit if no documents match the filter
+            if docs_df.empty:
+                return []
+            # Keep as pandas Series for efficient vectorized operations
+            filtered_doc_ids = docs_df["id"]
+
+        # Prepare search query based on search type
+        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"
+            )
+
+        elif search_type == "fts":
+            results = self.store.chunks_table.search(query, query_type="fts")
+
+        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)
+            )
+
+        # Apply filtering if needed (common for all search types)
+        if filtered_doc_ids is not None:
+            chunks_df = results.to_pandas()
+            filtered_chunks_df = chunks_df.loc[
+                chunks_df["document_id"].isin(filtered_doc_ids)
+            ].head(limit)
+            return await self._process_search_results(filtered_chunks_df)
+
+        # No filtering needed, apply limit and return
+        results = results.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_title = doc_results[0].title if doc_results else None
+        doc_meta = doc_results[0].metadata if doc_results else "{}"
+
+        chunks: list[Chunk] = []
+        for rec in results:
+            md = json.loads(rec.metadata)
+            chunks.append(
+                Chunk(
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=md,
+                    order=rec.order,
+                    document_uri=doc_uri,
+                    document_title=doc_title,
+                    document_meta=json.loads(doc_meta),
+                )
+            )
+
+        chunks.sort(key=lambda c: c.order)
+        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.order
+
+        # Fetch chunks for the same document and filter by order proximity
+        all_chunks = await self.get_by_document_id(chunk.document_id)
+
+        adjacent_chunks: list[Chunk] = []
+        for c in all_chunks:
+            c_order = c.order
+            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: "pd.DataFrame | LanceQueryBuilder"
+    ) -> list[tuple[Chunk, float]]:
+        """Process search results into chunks with document info and scores.
+
+        Args:
+            query_result: Either a pandas DataFrame or a LanceDB query result
+        """
+        import pandas as pd
+
+        def extract_scores(df: pd.DataFrame) -> list[float]:
+            """Extract scores from DataFrame columns based on search type."""
+            if "_distance" in df.columns:
+                # Vector search - convert distance to similarity
+                return ((df["_distance"] + 1).rdiv(1)).clip(lower=0.0).tolist()
+            elif "_relevance_score" in df.columns:
+                # Hybrid search - relevance score (higher is better)
+                return df["_relevance_score"].tolist()
+            elif "_score" in df.columns:
+                # FTS search - score (higher is better)
+                return df["_score"].tolist()
+            else:
+                raise ValueError("Unknown search result format, cannot extract scores")
+
+        # Convert everything to DataFrame for uniform processing
+        if isinstance(query_result, pd.DataFrame):
+            df = query_result
+        else:
+            # Convert LanceDB query result to DataFrame
+            df = query_result.to_pandas()
+
+        # Extract scores
+        scores = extract_scores(df)
+
+        # Convert DataFrame rows to ChunkRecords
+        pydantic_results = [
+            self.store.ChunkRecord(
+                id=str(row["id"]),
+                document_id=str(row["document_id"]),
+                content=str(row["content"]),
+                metadata=str(row["metadata"]),
+                order=int(row["order"]) if "order" in row else 0,
+            )
+            for _, row in df.iterrows()
+        ]
+
+        # 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:
+            # Use IN clause for efficient batch lookup
+            id_list = "', '".join(document_ids)
+            where_clause = f"id IN ('{id_list}')"
+            doc_results = list(
+                self.store.documents_table.search()
+                .where(where_clause)
+                .to_pydantic(DocumentRecord)
+            )
+            documents_map = {doc.id: doc for doc in doc_results}
+
+        # Build final results with document info
+        chunks_with_scores = []
+        for i, chunk_record in enumerate(pydantic_results):
+            doc = documents_map.get(chunk_record.document_id)
+            chunk = Chunk(
+                id=chunk_record.id,
+                document_id=chunk_record.document_id,
+                content=chunk_record.content,
+                metadata=json.loads(chunk_record.metadata),
+                order=chunk_record.order,
+                document_uri=doc.uri if doc else None,
+                document_title=doc.title if doc else None,
+                document_meta=json.loads(doc.metadata if doc else "{}"),
+            )
+            score = scores[i] if i < len(scores) else 1.0
+            chunks_with_scores.append((chunk, score))
+
+        return chunks_with_scores

haiku/rag/store/repositories/document.py

@@ -0,0 +1,261 @@
+import asyncio
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+from haiku.rag.store.engine import DocumentRecord, Store
+from haiku.rag.store.models.document import Document
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+    from haiku.rag.store.models.chunk import Chunk
+
+
+class DocumentRepository:
+    """Repository for Document operations."""
+
+    def __init__(self, store: Store) -> None:
+        self.store = store
+        self._chunk_repository = None
+
+    @property
+    def chunk_repository(self):
+        """Lazy-load ChunkRepository when needed."""
+        if self._chunk_repository is None:
+            from haiku.rag.store.repositories.chunk import ChunkRepository
+
+            self._chunk_repository = ChunkRepository(self.store)
+        return self._chunk_repository
+
+    def _record_to_document(self, record: DocumentRecord) -> Document:
+        """Convert a DocumentRecord to a Document model."""
+        return Document(
+            id=record.id,
+            content=record.content,
+            uri=record.uri,
+            title=record.title,
+            metadata=json.loads(record.metadata),
+            created_at=datetime.fromisoformat(record.created_at)
+            if record.created_at
+            else datetime.now(),
+            updated_at=datetime.fromisoformat(record.updated_at)
+            if record.updated_at
+            else datetime.now(),
+        )
+
+    async def create(self, entity: Document) -> Document:
+        """Create a document in the database."""
+        # Generate new UUID
+        doc_id = str(uuid4())
+
+        # Create timestamp
+        now = datetime.now().isoformat()
+
+        # Create document record
+        doc_record = DocumentRecord(
+            id=doc_id,
+            content=entity.content,
+            uri=entity.uri,
+            title=entity.title,
+            metadata=json.dumps(entity.metadata),
+            created_at=now,
+            updated_at=now,
+        )
+
+        # Add to table
+        self.store.documents_table.add([doc_record])
+
+        entity.id = doc_id
+        entity.created_at = datetime.fromisoformat(now)
+        entity.updated_at = datetime.fromisoformat(now)
+        return entity
+
+    async def get_by_id(self, entity_id: str) -> Document | None:
+        """Get a document by its ID."""
+        results = list(
+            self.store.documents_table.search()
+            .where(f"id = '{entity_id}'")
+            .limit(1)
+            .to_pydantic(DocumentRecord)
+        )
+
+        if not results:
+            return None
+
+        return self._record_to_document(results[0])
+
+    async def update(self, entity: Document) -> Document:
+        """Update an existing document."""
+        assert entity.id, "Document ID is required for update"
+
+        # Update timestamp
+        now = datetime.now().isoformat()
+        entity.updated_at = datetime.fromisoformat(now)
+
+        # Update the record
+        self.store.documents_table.update(
+            where=f"id = '{entity.id}'",
+            values={
+                "content": entity.content,
+                "uri": entity.uri,
+                "title": entity.title,
+                "metadata": json.dumps(entity.metadata),
+                "updated_at": now,
+            },
+        )
+
+        return entity
+
+    async def delete(self, entity_id: str) -> bool:
+        """Delete a document by its ID."""
+        # Check if document exists
+        doc = await self.get_by_id(entity_id)
+        if doc is None:
+            return False
+
+        # Delete associated chunks first
+        await self.chunk_repository.delete_by_document_id(entity_id)
+
+        # Delete the document
+        self.store.documents_table.delete(f"id = '{entity_id}'")
+        return True
+
+    async def list_all(
+        self,
+        limit: int | None = None,
+        offset: int | None = None,
+        filter: str | None = None,
+    ) -> list[Document]:
+        """List all documents with optional pagination and filtering.
+
+        Args:
+            limit: Maximum number of documents to return.
+            offset: Number of documents to skip.
+            filter: Optional SQL WHERE clause to filter documents.
+
+        Returns:
+            List of Document instances matching the criteria.
+        """
+        query = self.store.documents_table.search()
+
+        if filter is not None:
+            query = query.where(filter)
+        if offset is not None:
+            query = query.offset(offset)
+        if limit is not None:
+            query = query.limit(limit)
+
+        results = list(query.to_pydantic(DocumentRecord))
+        return [self._record_to_document(doc) for doc in results]
+
+    async def get_by_uri(self, uri: str) -> Document | None:
+        """Get a document by its URI."""
+        results = list(
+            self.store.documents_table.search()
+            .where(f"uri = '{uri}'")
+            .limit(1)
+            .to_pydantic(DocumentRecord)
+        )
+
+        if not results:
+            return None
+
+        return self._record_to_document(results[0])
+
+    async def delete_all(self) -> None:
+        """Delete all documents from the database."""
+        # Delete all chunks first
+        await self.chunk_repository.delete_all()
+
+        # Get count before deletion
+        count = len(
+            list(
+                self.store.documents_table.search().limit(1).to_pydantic(DocumentRecord)
+            )
+        )
+        if count > 0:
+            # Drop and recreate table to clear all data
+            self.store.db.drop_table("documents")
+            self.store.documents_table = self.store.db.create_table(
+                "documents", schema=DocumentRecord
+            )
+
+    async def _create_and_chunk(
+        self,
+        entity: Document,
+        docling_document: "DoclingDocument | None",
+        chunks: list["Chunk"] | None = None,
+    ) -> Document:
+        """Create a document with its chunks and embeddings."""
+        # Snapshot table versions for versioned rollback (if supported)
+        versions = self.store.current_table_versions()
+
+        # Create the document
+        created_doc = await self.create(entity)
+
+        # Attempt to create chunks; on failure, prefer version rollback
+        try:
+            # Create chunks if not provided
+            if chunks is None:
+                assert docling_document is not None, (
+                    "docling_document is required when chunks are not provided"
+                )
+                assert created_doc.id is not None, (
+                    "Document ID should not be None after creation"
+                )
+                await self.chunk_repository.create_chunks_for_document(
+                    created_doc.id, docling_document
+                )
+            else:
+                # Use provided chunks, set order from list position
+                assert created_doc.id is not None, (
+                    "Document ID should not be None after creation"
+                )
+                for order, chunk in enumerate(chunks):
+                    chunk.document_id = created_doc.id
+                    chunk.order = order
+                    await self.chunk_repository.create(chunk)
+
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
+            return created_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise
+
+    async def _update_and_rechunk(
+        self, entity: Document, docling_document: "DoclingDocument"
+    ) -> Document:
+        """Update a document and regenerate its chunks."""
+        assert entity.id is not None, "Document ID is required for update"
+
+        # Snapshot table versions for versioned rollback
+        versions = self.store.current_table_versions()
+
+        # Delete existing chunks before writing new ones
+        await self.chunk_repository.delete_by_document_id(entity.id)
+
+        try:
+            # Update the document
+            updated_doc = await self.update(entity)
+
+            # Create new chunks
+            assert updated_doc.id is not None, (
+                "Document ID should not be None after update"
+            )
+            await self.chunk_repository.create_chunks_for_document(
+                updated_doc.id, docling_document
+            )
+
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
+            return updated_doc
+        except Exception:
+            # Roll back to the captured versions and re-raise
+            self.store.restore_table_versions(versions)
+            raise