刘金宝
23 lines
1015 B
Python
23 lines
1015 B
Python
import time
|
|
|
|
from app.core.config import settings
|
|
from app.integrations.embedding_adapter import OpenAICompatibleEmbeddingClient
|
|
|
|
|
|
class EmbeddingService:
|
|
"""Embedding 服务：按配置批量调用向量模型，控制批大小和耗时统计。"""
|
|
|
|
def __init__(self, client: OpenAICompatibleEmbeddingClient | None = None) -> None:
|
|
self.client = client or OpenAICompatibleEmbeddingClient()
|
|
|
|
async def embed_texts(self, texts: list[str]) -> tuple[list[list[float]], int]:
|
|
"""批量向量化：按 EMBEDDING_BATCH_SIZE 分批生成向量并返回总耗时。"""
|
|
start = time.perf_counter()
|
|
vectors: list[list[float]] = []
|
|
batch_size = max(1, settings.embedding_batch_size)
|
|
for index in range(0, len(texts), batch_size):
|
|
batch = texts[index : index + batch_size]
|
|
batch_vectors, _usage = await self.client.embed_texts(batch)
|
|
vectors.extend(batch_vectors)
|
|
return vectors, int((time.perf_counter() - start) * 1000)
|