Files
Mutual_Fund/rag/embedding.py
T

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