From 74eea9284ab6651bdf56db34e9ec9ab1980dca04 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 26 Jul 2025 10:42:27 +0300 Subject: [PATCH] create_document now optionally accepts pre-computed chunks --- docs/python.md | 25 ++++++++++++ src/haiku/rag/client.py | 9 ++++- src/haiku/rag/store/models/chunk.py | 3 +- src/haiku/rag/store/repositories/chunk.py | 14 +++++-- src/haiku/rag/store/repositories/document.py | 26 ++++++++++--- tests/test_client.py | 40 ++++++++++++++++++++ 6 files changed, 106 insertions(+), 11 deletions(-) diff --git a/docs/python.md b/docs/python.md index dc8121b5..8368bcc0 100644 --- a/docs/python.md +++ b/docs/python.md @@ -27,6 +27,31 @@ doc = await client.create_document( ) ``` +With custom externally generated chunks: +```python +from haiku.rag.store.models.chunk import Chunk + +# Create custom chunks with optional embeddings +chunks = [ + Chunk( + content="This is the first chunk", + metadata={"section": "intro"} + ), + Chunk( + content="This is the second chunk", + metadata={"section": "body"}, + embedding=[0.1] * 1024 # Optional pre-computed embedding + ), +] + +doc = await client.create_document( + content="Full document content", + uri="doc://custom", + metadata={"source": "manual"}, + chunks=chunks # Use provided chunks instead of auto-generating +) +``` + From file: ```python doc = await client.create_document_from_source("path/to/document.pdf") diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 2b04f78f..e50ce992 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -50,7 +50,11 @@ class HaikuRAG: return False async def create_document( - self, content: str, uri: str | None = None, metadata: dict | None = None + self, + content: str, + uri: str | None = None, + metadata: dict | None = None, + chunks: list[Chunk] | None = None, ) -> Document: """Create a new document with optional URI and metadata. @@ -58,6 +62,7 @@ class HaikuRAG: content: The text content of the document. uri: Optional URI identifier for the document. metadata: Optional metadata dictionary. + chunks: Optional list of pre-created chunks to use instead of generating new ones. Returns: The created Document instance. @@ -67,7 +72,7 @@ class HaikuRAG: uri=uri, metadata=metadata or {}, ) - return await self.document_repository.create(document) + return await self.document_repository.create(document, chunks) async def create_document_from_source( self, source: str | Path, metadata: dict = {} diff --git a/src/haiku/rag/store/models/chunk.py b/src/haiku/rag/store/models/chunk.py index bd26aaed..c128802c 100644 --- a/src/haiku/rag/store/models/chunk.py +++ b/src/haiku/rag/store/models/chunk.py @@ -7,8 +7,9 @@ class Chunk(BaseModel): """ id: int | None = None - document_id: int + document_id: int | None = None content: str metadata: dict = {} document_uri: str | None = None document_meta: dict = {} + embedding: list[float] | None = None diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 4261cfeb..3a5c1f78 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -18,6 +18,8 @@ class ChunkRepository(BaseRepository[Chunk]): """Create a chunk in the database.""" if self.store._connection is None: raise ValueError("Store connection is not available") + if entity.document_id is None: + raise ValueError("Chunk must have a document_id to be created") cursor = self.store._connection.cursor() cursor.execute( @@ -34,9 +36,15 @@ class ChunkRepository(BaseRepository[Chunk]): entity.id = cursor.lastrowid - # Generate and store embedding - embedding = await self.embedder.embed(entity.content) - serialized_embedding = self.store.serialize_embedding(embedding) + # Generate and store embedding - use existing one if provided + if entity.embedding is not None: + # Use the provided embedding + serialized_embedding = self.store.serialize_embedding(entity.embedding) + else: + # Generate embedding from content + embedding = await self.embedder.embed(entity.content) + serialized_embedding = self.store.serialize_embedding(embedding) + cursor.execute( """ INSERT INTO chunk_embeddings (chunk_id, embedding) diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index ab5cf8e8..dd5c17c8 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -1,8 +1,12 @@ import json +from typing import TYPE_CHECKING from haiku.rag.store.models.document import Document from haiku.rag.store.repositories.base import BaseRepository +if TYPE_CHECKING: + from haiku.rag.store.models.chunk import Chunk + class DocumentRepository(BaseRepository[Document]): """Repository for Document database operations.""" @@ -16,7 +20,9 @@ class DocumentRepository(BaseRepository[Document]): chunk_repository = ChunkRepository(store) self.chunk_repository = chunk_repository - async def create(self, entity: Document) -> Document: + async def create( + self, entity: Document, chunks: list["Chunk"] | None = None + ) -> Document: """Create a document with its chunks and embeddings.""" if self.store._connection is None: raise ValueError("Store connection is not available") @@ -46,10 +52,20 @@ class DocumentRepository(BaseRepository[Document]): assert document_id is not None, "Failed to create document in database" entity.id = document_id - # Create chunks and embeddings using ChunkRepository - await self.chunk_repository.create_chunks_for_document( - document_id, entity.content, commit=False - ) + # Create chunks - either use provided chunks or generate from content + if chunks is not None: + # Use provided chunks, but update their document_id and set order from list position + for order, chunk in enumerate(chunks): + chunk.document_id = document_id + # Ensure order is set from list position + chunk.metadata = chunk.metadata.copy() if chunk.metadata else {} + chunk.metadata["order"] = order + await self.chunk_repository.create(chunk, commit=False) + else: + # Create chunks and embeddings using ChunkRepository + await self.chunk_repository.create_chunks_for_document( + document_id, entity.content, commit=False + ) cursor.execute("COMMIT") return entity diff --git a/tests/test_client.py b/tests/test_client.py index c5aaeb1c..e2bb3b14 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -7,6 +7,7 @@ import pytest from datasets import Dataset from haiku.rag.client import HaikuRAG +from haiku.rag.store.models.chunk import Chunk @pytest.mark.asyncio @@ -449,3 +450,42 @@ async def test_client_async_context_manager(): # Context manager should have automatically closed the connection # We can't easily test that the connection is closed without accessing internals, # but the test passing means the context manager methods work correctly + + +@pytest.mark.asyncio +async def test_client_create_document_with_custom_chunks(): + """Test creating a document with pre-created chunks.""" + async with HaikuRAG(":memory:") as client: + # Create some custom chunks with and without embeddings + chunks = [ + Chunk(content="This is the first chunk", metadata={"custom": "metadata1"}), + Chunk( + content="This is the second chunk", + metadata={"custom": "metadata2"}, + embedding=[0.1] * 1024, + ), # With embedding + Chunk(content="This is the third chunk", metadata={"custom": "metadata3"}), + ] + + # Create document with custom chunks + document = await client.create_document( + content="Full document content", chunks=chunks + ) + + assert document.id is not None + assert document.content == "Full document content" + + # Verify the chunks were created correctly + doc_chunks = await client.chunk_repository.get_by_document_id(document.id) + assert len(doc_chunks) == 3 + + # Check chunks have correct content, document_id, and order from list position + for i, chunk in enumerate(doc_chunks): + assert chunk.document_id == document.id + assert chunk.content == chunks[i].content + assert ( + chunk.metadata["order"] == i + ) # Order should be set from list position + assert ( + chunk.metadata["custom"] == f"metadata{i + 1}" + ) # Original metadata preserved