31 lines
957 B
Python
31 lines
957 B
Python
|
"""Embedding provider wrapper for Milvus ingestion and retrieval."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from tool.llm import llm
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger("rag.embedding")
|
||
|
|
EMBEDDING_DIMENSION = 768
|
||
|
|
|
||
|
|
|
||
|
|
class EmbeddingError(RuntimeError):
|
||
|
|
"""Raised when the embedding provider fails or returns invalid vectors."""
|
||
|
|
|
||
|
|
|
||
|
|
async def embed_texts(texts: list[str], *, client=None) -> list[list[float]]:
|
||
|
|
if not texts:
|
||
|
|
return []
|
||
|
|
provider = client or llm
|
||
|
|
try:
|
||
|
|
vectors = await provider.embed(texts)
|
||
|
|
except Exception as exc: # provider-specific exceptions are normalized here
|
||
|
|
logger.exception("embedding provider failed")
|
||
|
|
raise EmbeddingError("Embedding service unavailable") from exc
|
||
|
|
if len(vectors) != len(texts) or any(
|
||
|
|
len(vector) != EMBEDDING_DIMENSION for vector in vectors
|
||
|
|
):
|
||
|
|
raise EmbeddingError(f"Embedding dimension must be {EMBEDDING_DIMENSION}")
|
||
|
|
return vectors
|