deploy: auto commit server changes 2026-06-11 20:55:54

server/public/dongqiudi-crawler/scripts/kb_worker.py

@@ -84,7 +84,7 @@ def vector_norm(vector: Iterable[float]) -> float:
     total = 0.0
     for value in vector:
         total += float(value) * float(value)
-    return math.sqrt(total) if total > 0 else 0.0
+    return round(math.sqrt(total), 8) if total > 0 else 0.0
 
 
 def json_loads_maybe(value: Any, default: Any) -> Any:
@@ -532,16 +532,21 @@ class EmbeddingClient:
         self.base_url = (os.environ.get("EMBEDDING_BASE_URL") or DEFAULT_EMBEDDING_BASE_URL).rstrip("/")
         self.model = os.environ.get("EMBEDDING_MODEL") or DEFAULT_EMBEDDING_MODEL
         self.dimension = int(os.environ.get("EMBEDDING_DIM") or DEFAULT_EMBEDDING_DIM)
+        self.batch_size = max(1, int(os.environ.get("EMBEDDING_BATCH_SIZE") or 10))
 
     async def embed(self, chunks: List[str]) -> Dict[str, Any]:
         if not self.api_key or not self.base_url or not self.model:
             return {"success": False, "error": "Embedding API配置未完成", "embeddings": []}
-        payload: Dict[str, Any] = {"model": self.model, "input": chunks}
-        if self.dimension > 0:
-            payload["dimensions"] = self.dimension
         start = time.monotonic()
+        embeddings: List[List[float]] = []
+        tokens = 0
         try:
             async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session:
+                for offset in range(0, len(chunks), self.batch_size):
+                    batch = chunks[offset:offset + self.batch_size]
+                    payload: Dict[str, Any] = {"model": self.model, "input": batch}
+                    if self.dimension > 0:
+                        payload["dimensions"] = self.dimension
                     async with session.post(
                         self.base_url + "/v1/embeddings",
                         json=payload,
@@ -556,18 +561,18 @@ class EmbeddingClient:
                         if resp.status != 200 or not isinstance(data.get("data"), list):
                             message = ((data.get("error") or {}).get("message") if isinstance(data, dict) else "") or f"HTTP {resp.status}"
                             return {"success": False, "error": message, "embeddings": [], "cost_ms": int((time.monotonic() - start) * 1000)}
-                    embeddings = []
                         for row in data.get("data") or []:
                             embedding = row.get("embedding") if isinstance(row, dict) else None
                             if isinstance(embedding, list) and embedding:
                                 embeddings.append([float(value) for value in embedding])
+                        tokens += int(((data.get("usage") or {}).get("total_tokens") or 0) if isinstance(data, dict) else 0)
             return {
-                        "success": bool(embeddings),
-                        "error": "" if embeddings else "Embedding结果为空",
+                "success": len(embeddings) == len(chunks),
+                "error": "" if len(embeddings) == len(chunks) else "Embedding结果数量不匹配",
                 "embeddings": embeddings,
                 "model": self.model,
                 "dimension": len(embeddings[0]) if embeddings else 0,
-                        "tokens": int(((data.get("usage") or {}).get("total_tokens") or 0) if isinstance(data, dict) else 0),
+                "tokens": tokens,
                 "cost_ms": int((time.monotonic() - start) * 1000),
             }
         except Exception as exc: