haiku.rag 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

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 = {}

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

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)

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

{haiku_rag-0.4.1.dist-info → haiku_rag-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.4.1
+Version: 0.4.2
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT

{haiku_rag-0.4.1.dist-info → haiku_rag-0.4.2.dist-info}/RECORD

@@ -2,7 +2,7 @@ haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haiku/rag/app.py,sha256=FpLVyP1-zAq_XPmU8CPVLkuIAeuhBOGvMqhYS8RbN40,7649
 haiku/rag/chunker.py,sha256=MbCtP66OfTFoIBvqmVT9T9c87fozsYYzAQzJJEfPBVI,1812
 haiku/rag/cli.py,sha256=oCj65JcV2MEhzA2okbLHAK1I0FrClIKuYZx2jtbjbqE,5628
-haiku/rag/client.py,sha256=gqHsRAZqM5s4-c-RjeR4HanOKqPuK0z_MtsfmZMvT-w,12553
+haiku/rag/client.py,sha256=MZNIpMm6MS3P6vjLqiCztT2dBOM7-bZOosX5IpbHJbI,12724
 haiku/rag/config.py,sha256=_Ss54kmfxVAJupExLKaYjYUlFxJgb7hEEdbG4-isapY,1662
 haiku/rag/logging.py,sha256=zTTGpGq5tPdcd7RpCbd9EGw1IZlQDbYkrCg9t9pqRc4,580
 haiku/rag/mcp.py,sha256=tMN6fNX7ZtAER1R6DL1GkC9HZozTC4HzuQs199p7icI,4551
@@ -27,17 +27,17 @@ haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
 haiku/rag/store/engine.py,sha256=4ouAD0s-TFwEoEHjVVw_KnV6aaw5nwhe9fdT8PRXfok,6061
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
-haiku/rag/store/models/chunk.py,sha256=lmbPOOTz-N4PXhrA5XCUxyRcSTZBo135fqkV1mwnGcE,309
+haiku/rag/store/models/chunk.py,sha256=9-vIxW75-kMTelIhgVIMd_WhP-Drc1q65vjaWMP8w1E,364
 haiku/rag/store/models/document.py,sha256=TVXVY-nQs-1vCORQEs9rA7zOtndeGC4dgCoujLAS054,396
 haiku/rag/store/repositories/__init__.py,sha256=uIBhxjQh-4o3O-ck8b7BQ58qXQTuJdPvrDIHVhY5T1A,263
 haiku/rag/store/repositories/base.py,sha256=cm3VyQXhtxvRfk1uJHpA0fDSxMpYN-mjQmRiDiLsQ68,1008
-haiku/rag/store/repositories/chunk.py,sha256=gik7ZPOK3gCoG6tU1pGueAZBPmJxIb7obYFUhwINrYg,16497
-haiku/rag/store/repositories/document.py,sha256=xpWOpjHFbhVwNJ1gpusEKNY6l_Qyibg9y_bdHCwcfpk,7133
+haiku/rag/store/repositories/chunk.py,sha256=UyvHhKb1ESZePoTp2GneAARdfKoocEdfPOwgWPPQ0v8,16878
+haiku/rag/store/repositories/document.py,sha256=fXIWevJaOe6x2cK4u9cQxiEGD0ntKQb9y3VRqklQypE,7920
 haiku/rag/store/repositories/settings.py,sha256=dme3_ulQdQvyF9daavSjAd-SjZ5hh0MJoxP7iXgap-A,2492
 haiku/rag/store/upgrades/__init__.py,sha256=kKS1YWT_P-CYKhKtokOLTIFNKf9jlfjFFr8lyIMeogM,100
 haiku/rag/store/upgrades/v0_3_4.py,sha256=GLogKZdZ40NX1vBHKdOJju7fFzNUCHoEnjSZg17Hm2U,663
-haiku_rag-0.4.1.dist-info/METADATA,sha256=Vqg_r9uBqdKh3V4dUgPGzx40cNUtXodIELD9_sU2xYs,4235
-haiku_rag-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.4.1.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.4.1.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.4.1.dist-info/RECORD,,
+haiku_rag-0.4.2.dist-info/METADATA,sha256=0ctTSGB6uqGl2INUcNxnCphxwrLTlK7KVdKKXXB35mg,4235
+haiku_rag-0.4.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.4.2.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.4.2.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.4.2.dist-info/RECORD,,