diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 75d39f77..8c031aba 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/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( diff --git a/haiku_rag_slim/haiku/rag/embeddings/voyageai.py b/haiku_rag_slim/haiku/rag/embeddings/voyageai.py index 4d6af089..ff6f16b0 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/voyageai.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 40cdadae..b5e7799e 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/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