Differentiate between embedding a query and a document

haiku_rag_slim/haiku/rag/embeddings/__init__.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, overload
+from typing import TYPE_CHECKING
 
 from pydantic_ai.embeddings import Embedder
 from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
@@ -12,27 +12,23 @@ if TYPE_CHECKING:
 
 
 class EmbedderWrapper:
-    """Wrapper around pydantic-ai Embedder to provide simple embed() interface."""
+    """Wrapper around pydantic-ai Embedder with explicit query/document methods."""
 
     def __init__(self, embedder: Embedder, vector_dim: int):
         self._embedder = embedder
         self._vector_dim = vector_dim
 
-    @overload
-    async def embed(self, text: str) -> list[float]: ...
+    async def embed_query(self, text: str) -> list[float]:
+        """Embed a search query."""
+        result = await self._embedder.embed_query(text)
+        return list(result.embeddings[0])
 
-    @overload
-    async def embed(self, text: list[str]) -> list[list[float]]: ...
-
-    async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
-        if isinstance(text, str):
-            result = await self._embedder.embed_query(text)
-            return list(result.embeddings[0])
-        else:
-            if not text:
-                return []
-            result = await self._embedder.embed_documents(text)
-            return [list(e) for e in result.embeddings]
+    async def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents/chunks for indexing."""
+        if not texts:
+            return []
+        result = await self._embedder.embed_documents(texts)
+        return [list(e) for e in result.embeddings]
 
 
 def contextualize(chunks: list["Chunk"]) -> list[str]:
@@ -80,7 +76,7 @@ async def embed_chunks(
 
     embedder = get_embedder(config)
     texts = contextualize(chunks)
-    embeddings = await embedder.embed(texts)
+    embeddings = await embedder.embed_documents(texts)
 
     return [
         Chunk(

haiku_rag_slim/haiku/rag/embeddings/voyageai.py

@@ -1,27 +1,33 @@
 try:
-    from typing import overload
-
     from voyageai.client import Client  # type: ignore
 
-    from haiku.rag.embeddings.base import EmbedderBase
+    from haiku.rag.config import AppConfig
 
-    class Embedder(EmbedderBase):
-        @overload
-        async def embed(self, text: str) -> list[float]: ...
+    class Embedder:
+        """VoyageAI embedder with explicit query/document methods."""
 
-        @overload
-        async def embed(self, text: list[str]) -> list[list[float]]: ...
+        def __init__(self, model: str, vector_dim: int, config: AppConfig):
+            self._model = model
+            self._vector_dim = vector_dim
+            self._config = config
 
-        async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
+        async def embed_query(self, text: str) -> list[float]:
+            """Embed a search query."""
             client = Client()
-            if not text:
+            res = client.embed(
+                [text], model=self._model, input_type="query", output_dtype="float"
+            )
+            return res.embeddings[0]  # type: ignore[return-value]
+
+        async def embed_documents(self, texts: list[str]) -> list[list[float]]:
+            """Embed documents/chunks for indexing."""
+            if not texts:
                 return []
-            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]
+            client = Client()
+            res = client.embed(
+                texts, model=self._model, input_type="document", output_dtype="float"
+            )
+            return res.embeddings  # type: ignore[return-value]
 
 except ImportError:
     pass

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -245,7 +245,7 @@ class ChunkRepository:
 
         # Prepare search query based on search type
         if search_type == "vector":
-            query_embedding = await self.embedder.embed(query)
+            query_embedding = await self.embedder.embed_query(query)
             vector_query = cast(
                 "LanceVectorQueryBuilder",
                 self.store.chunks_table.search(
@@ -260,7 +260,7 @@ class ChunkRepository:
             results = self.store.chunks_table.search(query, query_type="fts")
 
         else:  # hybrid (default)
-            query_embedding = await self.embedder.embed(query)
+            query_embedding = await self.embedder.embed_query(query)
             # Create RRF reranker
             reranker = RRFReranker()
             # Perform native hybrid search with RRF reranking