haha/sbnews
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
hajimi
2026-06-11 20:55:54 +08:00
parent e17a9832af
commit 6a1d60f8a5
1 changed files with 37 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/public/dongqiudi-crawler/scripts/kb_worker.py
+37 -32
View File
@@ -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(
self.base_url + "/v1/embeddings",
json=payload,
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
ssl=False,
) as resp:
body = await resp.text()
try:
data = json.loads(body)
except Exception:
data = {}
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])
return {
"success": bool(embeddings),
"error": "" if embeddings 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),
"cost_ms": int((time.monotonic() - start) * 1000),
}
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,
headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
ssl=False,
) as resp:
body = await resp.text()
try:
data = json.loads(body)
except Exception:
data = {}
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)}
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": 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": tokens,
"cost_ms": int((time.monotonic() - start) * 1000),
}
except Exception as exc:
return {"success": False, "error": f"Embedding请求失败: {exc}", "embeddings": []}