haiku.rag 0.7.2__py3-none-any.whl → 0.7.3__py3-none-any.whl

haiku/rag/embeddings/base.py

@@ -9,7 +9,7 @@ class EmbedderBase:
         self._model = model
         self._vector_dim = vector_dim
 
-    async def embed(self, text: str) -> list[float]:
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         raise NotImplementedError(
             "Embedder is an abstract class. Please implement the embed method in a subclass."
         )

haiku/rag/embeddings/ollama.py

@@ -1,11 +1,17 @@
-from ollama import AsyncClient
+from openai import AsyncOpenAI
 
 from haiku.rag.config import Config
 from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
-    async def embed(self, text: str) -> list[float]:
-        client = AsyncClient(host=Config.OLLAMA_BASE_URL)
-        res = await client.embeddings(model=self._model, prompt=text)
-        return list(res["embedding"])
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        client = AsyncOpenAI(base_url=f"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy")
+        response = await client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/openai.py

@@ -4,10 +4,13 @@ from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
-    async def embed(self, text: str) -> list[float]:
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         client = AsyncOpenAI()
         response = await client.embeddings.create(
             model=self._model,
             input=text,
         )
-        return response.data[0].embedding
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/vllm.py

@@ -5,7 +5,7 @@ from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
-    async def embed(self, text: str) -> list[float]:
+    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         client = AsyncOpenAI(
             base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy"
         )
@@ -13,4 +13,7 @@ class Embedder(EmbedderBase):
             model=self._model,
             input=text,
         )
-        return response.data[0].embedding
+        if isinstance(text, str):
+            return response.data[0].embedding
+        else:
+            return [item.embedding for item in response.data]

haiku/rag/embeddings/voyageai.py

@@ -4,10 +4,14 @@ try:
     from haiku.rag.embeddings.base import EmbedderBase
 
     class Embedder(EmbedderBase):
-        async def embed(self, text: str) -> list[float]:
+        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
             client = Client()
-            res = client.embed([text], model=self._model, output_dtype="float")
-            return res.embeddings[0]  # type: ignore[return-value]
+            if isinstance(text, str):
+                res = client.embed([text], model=self._model, output_dtype="float")
+                return res.embeddings[0]  # type: ignore[return-value]
+            else:
+                res = client.embed(text, model=self._model, output_dtype="float")
+                return res.embeddings  # type: ignore[return-value]
 
 except ImportError:
     pass

haiku/rag/store/repositories/chunk.py

@@ -154,13 +154,7 @@ class ChunkRepository:
         """Create chunks and embeddings for a document from DoclingDocument."""
         chunk_texts = await chunker.chunk(document)
 
-        # Generate embeddings in parallel for all chunks
-        embeddings_tasks = []
-        for chunk_text in chunk_texts:
-            embeddings_tasks.append(self.embedder.embed(chunk_text))
-
-        # Wait for all embeddings to complete
-        embeddings = await asyncio.gather(*embeddings_tasks)
+        embeddings = await self.embedder.embed(chunk_texts)
 
         # Prepare all chunk records for batch insertion
         chunk_records = []

{haiku_rag-0.7.2.dist-info → haiku_rag-0.7.3.dist-info}/METADATA

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

{haiku_rag-0.7.2.dist-info → haiku_rag-0.7.3.dist-info}/RECORD

@@ -11,11 +11,11 @@ haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
 haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
 haiku/rag/utils.py,sha256=c8F0ECsFSqvQxzxINAOAnvShoOnJPLsOaNE3JEY2JSc,3230
 haiku/rag/embeddings/__init__.py,sha256=n7aHW3BxHlpGxU4ze4YYDOsljzFpEep8dwVE2n45JoE,1218
-haiku/rag/embeddings/base.py,sha256=NTQvuzbZPu0LBo5wAu3qGyJ4xXUaRAt1fjBO0ygWn_Y,465
-haiku/rag/embeddings/ollama.py,sha256=y6-lp0XpbnyIjoOEdtSzMdEVkU5glOwnWQ1FkpUZnpI,370
-haiku/rag/embeddings/openai.py,sha256=iA-DewCOSip8PLU_RhEJHFHBle4DtmCCIGNfGs58Wvk,357
-haiku/rag/embeddings/vllm.py,sha256=ymNDmpIvDWmkmce5j-TRc_QnJR4qwZCpnYA0tP3ab5o,480
-haiku/rag/embeddings/voyageai.py,sha256=0hiRTIqu-bpl-4OaCtMHvWfPdgbrzhnfZJowSV8pLRA,415
+haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
+haiku/rag/embeddings/ollama.py,sha256=LuLlHH6RGoO9_gFCIlbmesuXOj017gTw6z-p8Ez0CfE,595
+haiku/rag/embeddings/openai.py,sha256=fIFCk-jpUtaW0xsnrQnJ824O0UCjaGG2sgvBzREhilc,503
+haiku/rag/embeddings/vllm.py,sha256=vhaUnCn6VMkfSluLhWKtSV-sekFaPsp4pKo2N7-SBCY,626
+haiku/rag/embeddings/voyageai.py,sha256=UW-MW4tJKnPB6Fs2P7A3yt-ZeRm46H9npckchSriPX8,661
 haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
 haiku/rag/qa/agent.py,sha256=15-jMuF08U0uxGdqgQysKMZLr8BUWssI76PtyQ2Ngd8,2912
 haiku/rag/qa/prompts.py,sha256=xdT4cyrOrAK9UDgVqyev1wHF49jD57Bh40gx2sH4NPI,3341
@@ -30,12 +30,12 @@ haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnv
 haiku/rag/store/models/chunk.py,sha256=ZNyTfO6lh3rXWLVYO3TZcitbL4LSUGr42fR6jQQ5iQc,364
 haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
 haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
-haiku/rag/store/repositories/chunk.py,sha256=5S77mGh6pWxPHjaXriJGmvbSOhoNM8tLwygE2GXPlbU,13586
+haiku/rag/store/repositories/chunk.py,sha256=v4y4eh4yIf6zJaWfHxljvnmb12dmvwdinzmxQt8Lvhs,13343
 haiku/rag/store/repositories/document.py,sha256=lP8Lo82KTP-qwXFRpYZ46WjeAdAsHwZ5pJcrXdz4g0U,6988
 haiku/rag/store/repositories/settings.py,sha256=dqnAvm-98nQrWpLBbf9QghJw673QD80-iqQhRMP5t0c,5025
 haiku/rag/store/upgrades/__init__.py,sha256=wUiEoSiHTahvuagx93E4FB07v123AhdbOjwUkPusiIg,14
-haiku_rag-0.7.2.dist-info/METADATA,sha256=CLBIBBUYbvBbtynct-n_q7ZVO6cayLa7YUeFivrwEf4,4610
-haiku_rag-0.7.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.7.2.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.7.2.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.7.2.dist-info/RECORD,,
+haiku_rag-0.7.3.dist-info/METADATA,sha256=PAvA6VZuyZp9IekXhYCLWDxM1wMZMmujtntxZE2lBoE,4610
+haiku_rag-0.7.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.7.3.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.7.3.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.7.3.dist-info/RECORD,,