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

haiku/rag/config.py

@@ -33,6 +33,9 @@ class AppConfig(BaseModel):
     CONTEXT_CHUNK_RADIUS: int = 0
 
     OLLAMA_BASE_URL: str = "http://localhost:11434"
+    VLLM_EMBEDDINGS_BASE_URL: str = ""
+    VLLM_RERANK_BASE_URL: str = ""
+    VLLM_QA_BASE_URL: str = ""
 
     # Provider keys
     VOYAGE_API_KEY: str = ""

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

@@ -0,0 +1,19 @@
+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.VLLM_EMBEDDINGS_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/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/qa/agent.py

@@ -2,6 +2,7 @@ from pydantic import BaseModel, Field
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
@@ -65,6 +66,13 @@ class QuestionAnswerAgent:
                 model_name=model,
                 provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
             )
+        elif provider == "vllm":
+            return OpenAIChatModel(
+                model_name=model,
+                provider=OpenAIProvider(
+                    base_url=f"{Config.VLLM_QA_BASE_URL}/v1", api_key="none"
+                ),
+            )
         else:
             # For all other providers, use the provider:model format
             return f"{provider}:{model}"

haiku/rag/reranking/vllm.py

@@ -0,0 +1,44 @@
+import httpx
+
+from haiku.rag.config import Config
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class VLLMReranker(RerankerBase):
+    def __init__(self, model: str):
+        self._model = model
+        self._base_url = Config.VLLM_RERANK_BASE_URL
+
+    async def rerank(
+        self, query: str, chunks: list[Chunk], top_n: int = 10
+    ) -> list[tuple[Chunk, float]]:
+        if not chunks:
+            return []
+
+        # Prepare documents for reranking
+        documents = [chunk.content for chunk in chunks]
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self._base_url}/v1/rerank",
+                json={"model": self._model, "query": query, "documents": documents},
+                headers={
+                    "accept": "application/json",
+                    "Content-Type": "application/json",
+                },
+            )
+            response.raise_for_status()
+
+            result = response.json()
+
+            # Extract scores and pair with chunks
+            scored_chunks = []
+            for item in result.get("results", []):
+                index = item["index"]
+                score = item["relevance_score"]
+                scored_chunks.append((chunks[index], score))
+
+            # Sort by score (descending) and return top_n
+            scored_chunks.sort(key=lambda x: x[1], reverse=True)
+            return scored_chunks[:top_n]

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.1.dist-info → haiku_rag-0.7.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.7.1
+Version: 0.7.3
 Summary: Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -47,10 +47,10 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 ## Features
 
 - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
-- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI
+- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
 - **Multiple QA providers**: Any provider/model supported by Pydantic AI
 - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
-- **Reranking**: Default search result reranking with MixedBread AI or Cohere
+- **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
 - **Question answering**: Built-in QA agents on your documents
 - **File monitoring**: Auto-index files when run as server
 - **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs

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

@@ -3,7 +3,7 @@ haiku/rag/app.py,sha256=GmuZxH7BMutWt8Mdu0RSateRBaKiqXh7Z9tV7cZX6n0,7655
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
 haiku/rag/cli.py,sha256=UY9Vh5RsIxSCV14eQbNOiwToKmbFAvqTOAnxjieaYBs,6399
 haiku/rag/client.py,sha256=N4zkWjE9Rsw9YgPvNo83xptHUQR2ognfOnjkoV_w6hc,20999
-haiku/rag/config.py,sha256=9Mv0QJ3c6VF1oVRSXJlSsG234dCd_sKnJO-ybMaTpDQ,1690
+haiku/rag/config.py,sha256=3H41da9BU1R1y2JJHD0cOSErX_VSM1UXA7M2JSOxFXE,1795
 haiku/rag/logging.py,sha256=DOQi9QMpQRl8h17Vu4nQh8HxpHdeIu29n8-HZaT3SRQ,786
 haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
 haiku/rag/migration.py,sha256=gWxQwiKo0YulRhogYz4K8N98kHN9LQXIx9FeTmT24v4,10915
@@ -11,29 +11,31 @@ 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/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=gZ12vLUSHHCMl0HyPoLlPDbhUWoyEUydXG7u8lG1eqg,2602
+haiku/rag/qa/agent.py,sha256=15-jMuF08U0uxGdqgQysKMZLr8BUWssI76PtyQ2Ngd8,2912
 haiku/rag/qa/prompts.py,sha256=xdT4cyrOrAK9UDgVqyev1wHF49jD57Bh40gx2sH4NPI,3341
 haiku/rag/reranking/__init__.py,sha256=IRXHs4qPu6VbGJQpzSwhgtVWWumURH_vEoVFE-extlo,894
 haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
 haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
 haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
+haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
 haiku/rag/store/engine.py,sha256=XHGo5Xl-dCFdQHrOdMo64xVK5n0k8-LoUl5V-tlA0HI,7131
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
 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.1.dist-info/METADATA,sha256=yIz4nATa_b2vDstrCZe9CzV_77Vw74QIhRWqCZCdaes,4597
-haiku_rag-0.7.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.7.1.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.7.1.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.7.1.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,,