haiku.rag 0.7.2__tar.gz → 0.7.3__tar.gz

{haiku_rag-0.7.2 → haiku_rag-0.7.3}/PKG-INFO

@@ -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 → haiku_rag-0.7.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.7.2"
+version = "0.7.3"
 description = "Retrieval Augmented Generation (RAG) with LanceDB"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }

{haiku_rag-0.7.2 → haiku_rag-0.7.3}/src/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-0.7.3/src/haiku/rag/embeddings/ollama.py

@@ -0,0 +1,17 @@
+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[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-0.7.2 → haiku_rag-0.7.3}/src/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-0.7.2 → haiku_rag-0.7.3}/src/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-0.7.3/src/haiku/rag/embeddings/voyageai.py

@@ -0,0 +1,17 @@
+try:
+    from voyageai.client import Client  # type: ignore
+
+    from haiku.rag.embeddings.base import EmbedderBase
+
+    class Embedder(EmbedderBase):
+        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+            client = Client()
+            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-0.7.2 → haiku_rag-0.7.3}/src/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 → haiku_rag-0.7.3}/tests/test_embedder.py

@@ -28,7 +28,13 @@ async def test_ollama_embedder():
         "Python is my favorite programming language.",
         "I love to travel and see new places.",
     ]
-    embeddings = [np.array(await embedder.embed(phrase)) for phrase in phrases]
+
+    # Test batch embedding
+    embeddings = await embedder.embed(phrases)
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == 3
+    assert all(isinstance(emb, list) for emb in embeddings)
+    embeddings = [np.array(emb) for emb in embeddings]
 
     test_phrase = "I am going for a camping trip."
     test_embedding = await embedder.embed(test_phrase)
@@ -58,7 +64,13 @@ async def test_openai_embedder():
         "Python is my favorite programming language.",
         "I love to travel and see new places.",
     ]
-    embeddings = [np.array(await embedder.embed(phrase)) for phrase in phrases]
+
+    # Test batch embedding
+    embeddings = await embedder.embed(phrases)
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == 3
+    assert all(isinstance(emb, list) for emb in embeddings)
+    embeddings = [np.array(emb) for emb in embeddings]
 
     test_phrase = "I am going for a camping trip."
     test_embedding = await embedder.embed(test_phrase)
@@ -91,7 +103,13 @@ async def test_voyageai_embedder():
             "Python is my favorite programming language.",
             "I love to travel and see new places.",
         ]
-        embeddings = [np.array(await embedder.embed(phrase)) for phrase in phrases]
+
+        # Test batch embedding
+        embeddings = await embedder.embed(phrases)
+        assert isinstance(embeddings, list)
+        assert len(embeddings) == 3
+        assert all(isinstance(emb, list) for emb in embeddings)
+        embeddings = [np.array(emb) for emb in embeddings]
 
         test_phrase = "I am going for a camping trip."
         test_embedding = await embedder.embed(test_phrase)
@@ -126,7 +144,13 @@ async def test_vllm_embedder():
         "Python is my favorite programming language.",
         "I love to travel and see new places.",
     ]
-    embeddings = [np.array(await embedder.embed(phrase)) for phrase in phrases]
+
+    # Test batch embedding
+    embeddings = await embedder.embed(phrases)
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == 3
+    assert all(isinstance(emb, list) for emb in embeddings)
+    embeddings = [np.array(emb) for emb in embeddings]
 
     test_phrase = "I am going for a camping trip."
     test_embedding = await embedder.embed(test_phrase)

{haiku_rag-0.7.2 → haiku_rag-0.7.3}/uv.lock

@@ -951,7 +951,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.7.2"
+version = "0.7.3"
 source = { editable = "." }
 dependencies = [
     { name = "docling" },

haiku_rag-0.7.2/src/haiku/rag/embeddings/ollama.py

@@ -1,11 +0,0 @@
-from ollama import AsyncClient
-
-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"])

haiku_rag-0.7.2/src/haiku/rag/embeddings/voyageai.py

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