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,44 +532,49 @@ 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:
-                async with session.post(
+                for offset in range(0, len(chunks), self.batch_size):
-                    self.base_url + "/v1/embeddings",
+                    batch = chunks[offset:offset + self.batch_size]
-                    json=payload,
+                    payload: Dict[str, Any] = {"model": self.model, "input": batch}
-                    headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+                    if self.dimension > 0:
-                    ssl=False,
+                        payload["dimensions"] = self.dimension
-                ) as resp:
+                    async with session.post(
-                    body = await resp.text()
+                        self.base_url + "/v1/embeddings",
-                    try:
+                        json=payload,
-                        data = json.loads(body)
+                        headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
-                    except Exception:
+                        ssl=False,
-                        data = {}
+                    ) as resp:
-                    if resp.status != 200 or not isinstance(data.get("data"), list):
+                        body = await resp.text()
-                        message = ((data.get("error") or {}).get("message") if isinstance(data, dict) else "") or f"HTTP {resp.status}"
+                        try:
-                        return {"success": False, "error": message, "embeddings": [], "cost_ms": int((time.monotonic() - start) * 1000)}
+                            data = json.loads(body)
-                    embeddings = []
+                        except Exception:
-                    for row in data.get("data") or []:
+                            data = {}
-                        embedding = row.get("embedding") if isinstance(row, dict) else None
+                        if resp.status != 200 or not isinstance(data.get("data"), list):
-                        if isinstance(embedding, list) and embedding:
+                            message = ((data.get("error") or {}).get("message") if isinstance(data, dict) else "") or f"HTTP {resp.status}"
-                            embeddings.append([float(value) for value in embedding])
+                            return {"success": False, "error": message, "embeddings": [], "cost_ms": int((time.monotonic() - start) * 1000)}
-                    return {
+                        for row in data.get("data") or []:
-                        "success": bool(embeddings),
+                            embedding = row.get("embedding") if isinstance(row, dict) else None
-                        "error": "" if embeddings else "Embedding结果为空",
+                            if isinstance(embedding, list) and embedding:
-                        "embeddings": embeddings,
+                                embeddings.append([float(value) for value in embedding])
-                        "model": self.model,
+                        tokens += int(((data.get("usage") or {}).get("total_tokens") or 0) if isinstance(data, dict) else 0)
-                        "dimension": len(embeddings[0]) if embeddings else 0,
+            return {
-                        "tokens": int(((data.get("usage") or {}).get("total_tokens") or 0) if isinstance(data, dict) else 0),
+                "success": len(embeddings) == len(chunks),
-                        "cost_ms": int((time.monotonic() - start) * 1000),
+                "error": "" if len(embeddings) == len(chunks) else "Embedding结果数量不匹配",
-                    }
+                "embeddings": embeddings,
+                "model": self.model,
+                "dimension": len(embeddings[0]) if embeddings else 0,
+                "tokens": tokens,
+                "cost_ms": int((time.monotonic() - start) * 1000),
+            }
         except Exception as exc:
             return {"success": False, "error": f"Embedding请求失败: {exc}", "embeddings": []}