151 lines
4.8 KiB
Python
151 lines
4.8 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from pydantic_ai.embeddings import Embedder
|
|
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
|
|
from pydantic_ai.providers.ollama import OllamaProvider
|
|
from pydantic_ai.providers.openai import OpenAIProvider
|
|
|
|
from haiku.rag.config import AppConfig, Config
|
|
|
|
if TYPE_CHECKING:
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
|
|
class EmbedderWrapper:
|
|
"""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
|
|
|
|
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])
|
|
|
|
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]:
|
|
"""Prepare chunk content for embedding/FTS by adding context.
|
|
|
|
Prepends section headings to chunk content for better semantic search.
|
|
|
|
Args:
|
|
chunks: List of chunks to contextualize.
|
|
|
|
Returns:
|
|
List of contextualized text strings.
|
|
"""
|
|
texts = []
|
|
for chunk in chunks:
|
|
meta = chunk.get_chunk_metadata()
|
|
if meta.headings:
|
|
text = "\n".join(meta.headings) + "\n" + chunk.content
|
|
else:
|
|
text = chunk.content
|
|
texts.append(text)
|
|
return texts
|
|
|
|
|
|
EMBEDDING_BATCH_SIZE = 512
|
|
|
|
|
|
async def embed_chunks(
|
|
chunks: list["Chunk"], config: AppConfig = Config
|
|
) -> list["Chunk"]:
|
|
"""Generate embeddings for chunks.
|
|
|
|
Contextualizes chunks (prepends headings) before embedding for better
|
|
semantic search. Returns new Chunk objects with embeddings set.
|
|
|
|
Embeddings are generated in batches to avoid request size limits
|
|
and timeouts with large document sets.
|
|
|
|
Args:
|
|
chunks: List of chunks to embed.
|
|
config: Configuration for embedder selection.
|
|
|
|
Returns:
|
|
New list of Chunk objects with embedding field populated.
|
|
"""
|
|
if not chunks:
|
|
return []
|
|
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
embedder = get_embedder(config)
|
|
texts = contextualize(chunks)
|
|
|
|
# Batch embedding calls to avoid request size limits
|
|
all_embeddings: list[list[float]] = []
|
|
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
|
|
batch = texts[i : i + EMBEDDING_BATCH_SIZE]
|
|
batch_embeddings = await embedder.embed_documents(batch)
|
|
all_embeddings.extend(batch_embeddings)
|
|
|
|
return [
|
|
Chunk(
|
|
id=chunk.id,
|
|
document_id=chunk.document_id,
|
|
content=chunk.content,
|
|
metadata=chunk.metadata,
|
|
order=chunk.order,
|
|
document_uri=chunk.document_uri,
|
|
document_title=chunk.document_title,
|
|
document_meta=chunk.document_meta,
|
|
embedding=embedding,
|
|
)
|
|
for chunk, embedding in zip(chunks, all_embeddings)
|
|
]
|
|
|
|
|
|
def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
|
|
"""Factory function to get the appropriate embedder based on the configuration.
|
|
|
|
Args:
|
|
config: Configuration to use. Defaults to global Config.
|
|
|
|
Returns:
|
|
An embedder instance configured according to the config.
|
|
"""
|
|
embedding_model = config.embeddings.model
|
|
provider = embedding_model.provider
|
|
model_name = embedding_model.name
|
|
vector_dim = embedding_model.vector_dim
|
|
|
|
if provider == "ollama":
|
|
# Use model-level base_url if set, otherwise fall back to providers config
|
|
base_url = embedding_model.base_url or f"{config.providers.ollama.base_url}/v1"
|
|
model = OpenAIEmbeddingModel(
|
|
model_name,
|
|
provider=OllamaProvider(base_url=base_url),
|
|
)
|
|
return EmbedderWrapper(Embedder(model), vector_dim)
|
|
|
|
if provider == "openai":
|
|
if embedding_model.base_url:
|
|
model = OpenAIEmbeddingModel(
|
|
model_name,
|
|
provider=OpenAIProvider(base_url=embedding_model.base_url),
|
|
)
|
|
return EmbedderWrapper(Embedder(model), vector_dim)
|
|
return EmbedderWrapper(Embedder(f"openai:{model_name}"), vector_dim)
|
|
|
|
if provider == "voyageai":
|
|
return EmbedderWrapper(Embedder(f"voyageai:{model_name}"), vector_dim)
|
|
|
|
if provider == "cohere":
|
|
return EmbedderWrapper(Embedder(f"cohere:{model_name}"), vector_dim)
|
|
|
|
if provider == "sentence-transformers":
|
|
return EmbedderWrapper(
|
|
Embedder(f"sentence-transformers:{model_name}"), vector_dim
|
|
)
|
|
|
|
raise ValueError(f"Unsupported embedding provider: {provider}")
|