Mutual_Fund/rag/embedding.py

"""Embedding provider wrapper for Milvus ingestion and retrieval."""
from __future__ import annotations

import logging

from config.settings import settings
from tool.llm import llm


logger = logging.getLogger("rag.embedding")
# 单一来源：.env 的 LLM_EMBED_DIMENSIONS；Milvus 建表、入库校验、embeddings 请求参数均以此为准
EMBEDDING_DIMENSION = settings.llm.embed_dimensions


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