Fix embedding 400 errors

This commit is contained in:
holenzzz 2026-02-26 12:43:38 +08:00
parent 4d5f38beae
commit 9cd6470849
4 changed files with 27 additions and 12 deletions

View File

@ -191,13 +191,13 @@ class MilvusKB(KnowledgeBase):
def _get_async_embedding_function(self, embed_info: dict):
"""获取 embedding 函数"""
embedding_model = self._get_async_embedding(embed_info)
return partial(embedding_model.abatch_encode, batch_size=40)
return partial(embedding_model.abatch_encode, batch_size=10)
def _get_embedding_function(self, embed_info: dict):
"""获取 embedding 函数"""
embedding_model = self._get_async_embedding(embed_info)
return partial(embedding_model.batch_encode, batch_size=40)
return partial(embedding_model.batch_encode, batch_size=10)
async def _get_milvus_collection(self, db_id: str):
"""获取或创建 Milvus 集合"""

View File

@ -488,7 +488,7 @@ class UploadGraphService:
logger.error(f"加载图数据库信息失败:{e}")
return False
async def aget_embedding(self, text, batch_size=40):
async def aget_embedding(self, text, batch_size=10):
if isinstance(text, list):
outputs = await self.embed_model.abatch_encode(text, batch_size=batch_size)
return outputs

View File

@ -46,7 +46,7 @@ class BaseEmbeddingModel(ABC):
"""等同于aencode"""
return await self.aencode(queries)
def batch_encode(self, messages: list[str], batch_size: int = 40) -> list[list[float]]:
def batch_encode(self, messages: list[str], batch_size: int = 10) -> list[list[float]]:
# logger.info(f"Batch encoding {len(messages)} messages")
data = []
task_id = None
@ -67,24 +67,39 @@ class BaseEmbeddingModel(ABC):
return data
async def abatch_encode(self, messages: list[str], batch_size: int = 40) -> list[list[float]]:
async def abatch_encode(self, messages: list[str], batch_size: int = 10) -> list[list[float]]:
data = []
task_id = None
if len(messages) > batch_size:
task_id = hashstr(messages)
self.embed_state[task_id] = {"status": "in-progress", "total": len(messages), "progress": 0}
tasks = []
#保留原有逻辑:
#使用 asyncio.gather 并发执行所有 embedding 批次请求:
# tasks = []
# for i in range(0, len(messages), batch_size):
# group_msg = messages[i : i + batch_size]
# tasks.append(self.aencode(group_msg))
# results = await asyncio.gather(*tasks)
# for res in results:
# data.extend(res)
# if task_id:
# self.embed_state[task_id]["progress"] = len(messages)
# self.embed_state[task_id]["status"] = "completed"
# return data
for i in range(0, len(messages), batch_size):
group_msg = messages[i : i + batch_size]
tasks.append(self.aencode(group_msg))
results = await asyncio.gather(*tasks)
for res in results:
logger.info(f"Async encoding [{i}/{len(messages)}] messages (bsz={batch_size})")
res = await self.aencode(group_msg)
data.extend(res)
if task_id:
self.embed_state[task_id]["progress"] = i + len(group_msg)
if task_id:
self.embed_state[task_id]["progress"] = len(messages)
self.embed_state[task_id]["status"] = "completed"
return data

View File

@ -336,7 +336,7 @@ class EvaluationService:
# Currently, we re-calculate embeddings for ALL chunks in the KB for every benchmark generation.
# This is inefficient for large KBs (O(N) embedding calls).
# Optimization: Reuse existing embeddings from Vector DB if embedding_model_id matches the KB's embedding model.
embeddings = await embed_model.abatch_encode(contents, batch_size=40)
embeddings = await embed_model.abatch_encode(contents, batch_size=10)
norms = [math.sqrt(sum(x * x for x in vec)) or 1.0 for vec in embeddings]
def cosine(a, b, na, nb):