36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
|
"""NL2SQL 向量化适配器，复用公共 LLM 客户端。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from config.settings import settings
|
||
|
|
from tool.llm import llm
|
||
|
|
|
||
|
|
logger = logging.getLogger("nl2sql.embedding")
|
||
|
|
EMBEDDING_BATCH_SIZE = 10
|
||
|
|
|
||
|
|
|
||
|
|
class EmbeddingError(RuntimeError):
|
||
|
|
"""向量服务返回无效数据时抛出的异常。"""
|
||
|
|
|
||
|
|
|
||
|
|
async def embed_texts(texts: list[str], *, client=None) -> list[list[float]]:
|
||
|
|
if not texts:
|
||
|
|
return []
|
||
|
|
provider = client or llm
|
||
|
|
dimension = settings.llm.embed_dimensions
|
||
|
|
vectors: list[list[float]] = []
|
||
|
|
for start in range(0, len(texts), EMBEDDING_BATCH_SIZE):
|
||
|
|
batch = texts[start : start + EMBEDDING_BATCH_SIZE]
|
||
|
|
try:
|
||
|
|
batch_vectors = await provider.embed(batch)
|
||
|
|
except Exception as exc: # noqa: BLE001 向量服务异常统一转换
|
||
|
|
logger.exception("NL2SQL embedding provider failed")
|
||
|
|
raise EmbeddingError("Embedding service unavailable") from exc
|
||
|
|
if len(batch_vectors) != len(batch) or any(
|
||
|
|
len(vector) != dimension for vector in batch_vectors
|
||
|
|
):
|
||
|
|
raise EmbeddingError(f"Embedding dimension must be {dimension}")
|
||
|
|
vectors.extend(batch_vectors)
|
||
|
|
return vectors
|