deploy: auto commit server changes 2026-06-11 20:44:36
This commit is contained in:
hajimi
@@ -0,0 +1,694 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import hashlib
|
||||
import html
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import aiohttp
|
||||
import aiomysql
|
||||
import redis.asyncio as aioredis
|
||||
from loguru import logger
|
||||
|
||||
|
||||
APP_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(APP_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(APP_DIR))
|
||||
|
||||
from src.core.config import get_config, load_config # noqa: E402
|
||||
|
||||
|
||||
STREAM_KEY = "ai:kb:sync:stream"
|
||||
DEAD_STREAM_KEY = "ai:kb:sync:dead"
|
||||
GROUP_NAME = "ai_kb_workers"
|
||||
COLLECTION_NAME = "global"
|
||||
DEFAULT_CHUNK_SIZE = 600
|
||||
DEFAULT_CHUNK_OVERLAP = 80
|
||||
DEFAULT_EMBEDDING_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode"
|
||||
DEFAULT_EMBEDDING_MODEL = "text-embedding-v4"
|
||||
DEFAULT_EMBEDDING_DIM = 1024
|
||||
|
||||
|
||||
def now_ts() -> int:
|
||||
return int(time.time())
|
||||
|
||||
|
||||
def normalize_plain_text(text: Any) -> str:
|
||||
value = html.unescape(str(text or ""))
|
||||
value = re.sub(r"\s+", " ", value.strip(), flags=re.UNICODE)
|
||||
return value.strip()
|
||||
|
||||
|
||||
def normalize_html_text(text: Any) -> str:
|
||||
value = str(text or "")
|
||||
value = re.sub(r"(?i)<br\s*/?>|</p>|</div>|</li>", "\n", value)
|
||||
value = re.sub(r"<[^>]+>", "", value)
|
||||
return normalize_plain_text(value)
|
||||
|
||||
|
||||
def truncate_text(text: str, length: int) -> str:
|
||||
text = text.strip()
|
||||
if len(text) <= length:
|
||||
return text
|
||||
return text[:length] + "..."
|
||||
|
||||
|
||||
def split_text(text: str, chunk_size: int = DEFAULT_CHUNK_SIZE, overlap: int = DEFAULT_CHUNK_OVERLAP) -> List[str]:
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return []
|
||||
if len(text) <= chunk_size:
|
||||
return [text]
|
||||
chunks: List[str] = []
|
||||
start = 0
|
||||
step = max(1, chunk_size - overlap)
|
||||
while start < len(text):
|
||||
chunks.append(text[start : start + chunk_size].strip())
|
||||
if start + chunk_size >= len(text):
|
||||
break
|
||||
start += step
|
||||
return [chunk for chunk in chunks if chunk]
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def json_loads_maybe(value: Any, default: Any) -> Any:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, (dict, list)):
|
||||
return value
|
||||
try:
|
||||
return json.loads(str(value))
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def json_text(value: Any) -> str:
|
||||
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=str)
|
||||
|
||||
|
||||
class KbDatabase:
|
||||
def __init__(self):
|
||||
cfg = get_config().database
|
||||
self.cfg = cfg
|
||||
self.prefix = cfg.prefix
|
||||
self.pool: Optional[aiomysql.Pool] = None
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self.pool:
|
||||
return
|
||||
self.pool = await aiomysql.create_pool(
|
||||
host=self.cfg.host,
|
||||
port=self.cfg.port,
|
||||
user=self.cfg.username,
|
||||
password=self.cfg.password,
|
||||
db=self.cfg.database,
|
||||
charset=self.cfg.charset,
|
||||
autocommit=False,
|
||||
minsize=1,
|
||||
maxsize=max(2, min(10, self.cfg.pool_size)),
|
||||
)
|
||||
logger.info("KB worker MySQL pool ready: {}:{}/{}", self.cfg.host, self.cfg.port, self.cfg.database)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self.pool:
|
||||
self.pool.close()
|
||||
await self.pool.wait_closed()
|
||||
self.pool = None
|
||||
|
||||
def table(self, name: str) -> str:
|
||||
return f"`{self.prefix}{name}`"
|
||||
|
||||
async def fetchone(self, sql: str, args: tuple = ()) -> Optional[Dict[str, Any]]:
|
||||
await self.connect()
|
||||
assert self.pool is not None
|
||||
async with self.pool.acquire() as conn:
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await cur.execute(sql, args)
|
||||
return await cur.fetchone()
|
||||
|
||||
async def fetchall(self, sql: str, args: tuple = ()) -> List[Dict[str, Any]]:
|
||||
await self.connect()
|
||||
assert self.pool is not None
|
||||
async with self.pool.acquire() as conn:
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await cur.execute(sql, args)
|
||||
return list(await cur.fetchall())
|
||||
|
||||
async def build_payload(self, domain: str, subtype: str, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
key = f"{domain}:{subtype}"
|
||||
if key == "article:article":
|
||||
return await self.build_article_payload(source_id)
|
||||
if key == "post:post":
|
||||
return await self.build_post_payload(source_id)
|
||||
if key == "lottery:draw_result":
|
||||
return await self.build_lottery_draw_result_payload(source_id)
|
||||
if key == "lottery:ai_history":
|
||||
return await self.build_lottery_ai_history_payload(source_id)
|
||||
if key == "match:event":
|
||||
return await self.build_match_event_payload(source_id)
|
||||
return None
|
||||
|
||||
async def build_article_payload(self, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
row = await self.fetchone(
|
||||
f"SELECT * FROM {self.table('article')} WHERE id=%s AND is_show=1 AND delete_time IS NULL LIMIT 1",
|
||||
(source_id,),
|
||||
)
|
||||
if not row:
|
||||
return None
|
||||
cate = await self.fetchone(f"SELECT name FROM {self.table('article_cate')} WHERE id=%s LIMIT 1", (row.get("cid") or 0,))
|
||||
cate_name = str((cate or {}).get("name") or "")
|
||||
ext = json_loads_maybe(row.get("ext"), {})
|
||||
translated = str(ext.get("translated_content") or "").strip() if isinstance(ext, dict) else ""
|
||||
content = normalize_html_text(translated or row.get("content") or "")
|
||||
summary = normalize_plain_text(row.get("abstract") or row.get("desc") or "")
|
||||
keywords = list(dict.fromkeys(filter(None, [
|
||||
str(row.get("title") or ""),
|
||||
cate_name,
|
||||
str(row.get("author") or ""),
|
||||
str(row.get("category") or ""),
|
||||
])))
|
||||
return {
|
||||
"title": str(row.get("title") or ""),
|
||||
"summary": summary,
|
||||
"canonical_text": "\n\n".join(filter(None, [str(row.get("title") or ""), summary, content])).strip(),
|
||||
"keywords": ",".join(keywords),
|
||||
"metadata_json": {
|
||||
"cid": int(row.get("cid") or 0),
|
||||
"cate_name": cate_name,
|
||||
"author": str(row.get("author") or ""),
|
||||
"category": str(row.get("category") or ""),
|
||||
"article_id": int(row.get("article_id") or 0),
|
||||
"published_at": str(row.get("published_at") or ""),
|
||||
"is_video": int(row.get("is_video") or 0),
|
||||
},
|
||||
"source_created_at": int(row.get("create_time") or 0),
|
||||
"source_updated_at": int(row.get("update_time") or row.get("create_time") or 0),
|
||||
}
|
||||
|
||||
async def build_post_payload(self, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
row = await self.fetchone(
|
||||
f"SELECT * FROM {self.table('community_post')} WHERE id=%s AND status=1 AND delete_time IS NULL LIMIT 1",
|
||||
(source_id,),
|
||||
)
|
||||
if not row:
|
||||
return None
|
||||
tags = await self.fetchall(
|
||||
f"""
|
||||
SELECT t.name
|
||||
FROM {self.table('community_post_tag')} pt
|
||||
JOIN {self.table('community_tag')} t ON t.id = pt.tag_id
|
||||
WHERE pt.post_id=%s
|
||||
""",
|
||||
(source_id,),
|
||||
)
|
||||
tag_names = [str(item.get("name") or "") for item in tags if item.get("name")]
|
||||
ext = json_loads_maybe(row.get("ext"), {})
|
||||
segments = [
|
||||
str(row.get("content") or ""),
|
||||
str(ext.get("translated_content") or "") if isinstance(ext, dict) else "",
|
||||
str(ext.get("lottery_analysis_content") or "") if isinstance(ext, dict) else "",
|
||||
" ".join(tag_names),
|
||||
]
|
||||
canonical = normalize_plain_text("\n\n".join(filter(None, segments)))
|
||||
return {
|
||||
"title": truncate_text(canonical, 40) or f"社区帖子#{source_id}",
|
||||
"summary": truncate_text(canonical, 200),
|
||||
"canonical_text": canonical,
|
||||
"keywords": ",".join(dict.fromkeys(tag_names)),
|
||||
"metadata_json": {
|
||||
"post_type": int(row.get("post_type") or 0),
|
||||
"match_id": int(row.get("match_id") or 0),
|
||||
"is_paid": int(row.get("is_paid") or 0),
|
||||
"user_id": int(row.get("user_id") or 0),
|
||||
"tags": tag_names,
|
||||
},
|
||||
"source_created_at": int(row.get("create_time") or 0),
|
||||
"source_updated_at": int(row.get("update_time") or row.get("create_time") or 0),
|
||||
}
|
||||
|
||||
async def build_lottery_draw_result_payload(self, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
row = await self.fetchone(f"SELECT * FROM {self.table('lottery_draw_result')} WHERE id=%s LIMIT 1", (source_id,))
|
||||
if not row:
|
||||
return None
|
||||
game = await self.fetchone(
|
||||
f"SELECT name, template FROM {self.table('lottery_game')} WHERE code=%s LIMIT 1",
|
||||
(row.get("code") or "",),
|
||||
)
|
||||
game = game or {}
|
||||
numbers = [item for item in str(row.get("draw_code") or "").split(",") if item]
|
||||
title = f"{game.get('name') or row.get('code')} 第{row.get('issue')}期开奖"
|
||||
summary = f"{title} 于 {row.get('draw_time') or ''} 开奖，开奖号码为 {'、'.join(numbers)}。"
|
||||
trend = json_loads_maybe(row.get("trend"), {})
|
||||
canonical = "\n".join(filter(None, [
|
||||
summary,
|
||||
f"下一期期号：{row.get('next_issue')}" if row.get("next_issue") else "",
|
||||
f"下期开奖时间：{row.get('next_time')}" if row.get("next_time") else "",
|
||||
"走势数据：" + normalize_plain_text(json_text(trend)) if trend not in ({}, None) else "",
|
||||
])).strip()
|
||||
return {
|
||||
"title": title,
|
||||
"summary": summary,
|
||||
"canonical_text": canonical,
|
||||
"keywords": ",".join(filter(None, [
|
||||
str(row.get("code") or ""),
|
||||
str(game.get("name") or ""),
|
||||
str(game.get("template") or ""),
|
||||
str(row.get("issue") or ""),
|
||||
])),
|
||||
"metadata_json": {
|
||||
"code": str(row.get("code") or ""),
|
||||
"issue": str(row.get("issue") or ""),
|
||||
"template": str(game.get("template") or ""),
|
||||
"game_name": str(game.get("name") or ""),
|
||||
"draw_time": str(row.get("draw_time") or ""),
|
||||
},
|
||||
"source_created_at": int(row.get("create_time") or 0),
|
||||
"source_updated_at": int(row.get("update_time") or row.get("create_time") or 0),
|
||||
}
|
||||
|
||||
async def build_lottery_ai_history_payload(self, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
row = await self.fetchone(f"SELECT * FROM {self.table('lottery_ai_analysis')} WHERE id=%s LIMIT 1", (source_id,))
|
||||
if not row:
|
||||
return None
|
||||
result = row.get("result") if row.get("result") is not None else ""
|
||||
result_text = result if isinstance(result, str) else json_text(result)
|
||||
canonical = normalize_plain_text(result_text)
|
||||
title = f"{row.get('code') or '彩票'} 第{row.get('issue') or '-'}期 {row.get('module') or 'module'} 历史AI分析"
|
||||
return {
|
||||
"title": title,
|
||||
"summary": truncate_text(canonical, 220),
|
||||
"canonical_text": canonical,
|
||||
"keywords": ",".join(filter(None, [str(row.get("code") or ""), str(row.get("issue") or ""), str(row.get("module") or "")])),
|
||||
"metadata_json": {
|
||||
"code": str(row.get("code") or ""),
|
||||
"issue": str(row.get("issue") or ""),
|
||||
"module": str(row.get("module") or ""),
|
||||
},
|
||||
"source_created_at": int(row.get("create_time") or 0),
|
||||
"source_updated_at": int(row.get("update_time") or row.get("create_time") or 0),
|
||||
}
|
||||
|
||||
async def build_match_event_payload(self, source_id: int) -> Optional[Dict[str, Any]]:
|
||||
row = await self.fetchone(f"SELECT * FROM {self.table('match')} WHERE id=%s AND is_show=1 LIMIT 1", (source_id,))
|
||||
if not row:
|
||||
return None
|
||||
home = str(row.get("home_team") or "")
|
||||
away = str(row.get("away_team") or "")
|
||||
league = str(row.get("league_name") or "")
|
||||
match_time_value = int(row.get("match_time") or 0)
|
||||
match_time_text = time.strftime("%Y-%m-%d %H:%M", time.localtime(match_time_value)) if match_time_value > 0 else ""
|
||||
history = await self.fetchall(
|
||||
f"""
|
||||
SELECT * FROM {self.table('match_history')}
|
||||
WHERE (home_team=%s AND away_team=%s) OR (home_team=%s AND away_team=%s)
|
||||
ORDER BY match_time DESC LIMIT 5
|
||||
""",
|
||||
(home, away, away, home),
|
||||
)
|
||||
home_recent = await self.fetchall(
|
||||
f"SELECT * FROM {self.table('match_recent')} WHERE team_name=%s ORDER BY match_time DESC LIMIT 5",
|
||||
(home,),
|
||||
)
|
||||
away_recent = await self.fetchall(
|
||||
f"SELECT * FROM {self.table('match_recent')} WHERE team_name=%s ORDER BY match_time DESC LIMIT 5",
|
||||
(away,),
|
||||
)
|
||||
title = f"{home} vs {away}".strip()
|
||||
summary = (
|
||||
f"{league} {title}，{home} 对阵 {away}，比赛时间 {match_time_text}，"
|
||||
f"当前比分 {row.get('home_score') or 0}-{row.get('away_score') or 0}。"
|
||||
)
|
||||
canonical = "\n".join([
|
||||
summary,
|
||||
f"赔率：主胜{row.get('home_odds') or '-'}，平局{row.get('draw_odds') or '-'}，客胜{row.get('away_odds') or '-'}。",
|
||||
"历史交锋：" + normalize_plain_text(json_text(history)),
|
||||
f"{home}近期战绩：" + normalize_plain_text(json_text(home_recent)),
|
||||
f"{away}近期战绩：" + normalize_plain_text(json_text(away_recent)),
|
||||
])
|
||||
return {
|
||||
"title": title or f"赛事#{source_id}",
|
||||
"summary": normalize_plain_text(summary),
|
||||
"canonical_text": canonical.strip(),
|
||||
"keywords": ",".join(filter(None, [league, home, away])),
|
||||
"metadata_json": {
|
||||
"league": league,
|
||||
"teams": [team for team in [home, away] if team],
|
||||
"match_time": match_time_text,
|
||||
"status": int(row.get("status") or 0),
|
||||
},
|
||||
"source_created_at": int(row.get("create_time") or row.get("match_time") or 0),
|
||||
"source_updated_at": int(row.get("update_time") or row.get("match_time") or row.get("create_time") or 0),
|
||||
}
|
||||
|
||||
async def upsert_document(self, domain: str, subtype: str, source_id: int, embedder: "EmbeddingClient") -> Dict[str, Any]:
|
||||
payload = await self.build_payload(domain, subtype, source_id)
|
||||
if not payload:
|
||||
await self.mark_document_inactive(domain, subtype, source_id)
|
||||
return {"success": False, "error": "源内容不存在或不可用"}
|
||||
|
||||
content_hash = hashlib.md5(json_text([
|
||||
payload["title"],
|
||||
payload["summary"],
|
||||
payload["canonical_text"],
|
||||
payload["keywords"],
|
||||
payload["metadata_json"],
|
||||
]).encode("utf-8")).hexdigest()
|
||||
|
||||
existing = await self.fetchone(
|
||||
f"""
|
||||
SELECT * FROM {self.table('ai_kb_document')}
|
||||
WHERE domain=%s AND subtype=%s AND source_id=%s LIMIT 1
|
||||
""",
|
||||
(domain, subtype, source_id),
|
||||
)
|
||||
if existing and str(existing.get("content_hash") or "") == content_hash and int(existing.get("is_active") or 0) == 1:
|
||||
await self.update_document_metadata(int(existing["id"]), payload)
|
||||
chunk_count = await self.fetchone(
|
||||
f"SELECT COUNT(*) AS total FROM {self.table('ai_kb_chunk')} WHERE document_id=%s",
|
||||
(int(existing["id"]),),
|
||||
)
|
||||
return {"success": True, "document_id": int(existing["id"]), "chunk_count": int((chunk_count or {}).get("total") or 0), "skipped": True}
|
||||
|
||||
chunks = split_text(str(payload["canonical_text"] or ""))
|
||||
if not chunks:
|
||||
chunks = [payload["summary"] or payload["title"] or "empty"]
|
||||
embedding_result = await embedder.embed(chunks)
|
||||
embeddings = embedding_result.get("embeddings") or []
|
||||
if not embedding_result.get("success") or not embeddings:
|
||||
return {"success": False, "error": embedding_result.get("error") or "embedding失败"}
|
||||
|
||||
await self.connect()
|
||||
assert self.pool is not None
|
||||
async with self.pool.acquire() as conn:
|
||||
async with conn.cursor(aiomysql.DictCursor) as cur:
|
||||
await conn.begin()
|
||||
try:
|
||||
metadata_json = json_text(payload["metadata_json"])
|
||||
if existing:
|
||||
document_id = int(existing["id"])
|
||||
await cur.execute(
|
||||
f"""
|
||||
UPDATE {self.table('ai_kb_document')}
|
||||
SET title=%s, summary=%s, canonical_text=%s, keywords=%s, metadata_json=%s,
|
||||
source_created_at=%s, source_updated_at=%s, content_hash=%s, is_active=1, update_time=%s
|
||||
WHERE id=%s
|
||||
""",
|
||||
(
|
||||
payload["title"],
|
||||
payload["summary"],
|
||||
payload["canonical_text"],
|
||||
payload["keywords"],
|
||||
metadata_json,
|
||||
int(payload["source_created_at"]),
|
||||
int(payload["source_updated_at"]),
|
||||
content_hash,
|
||||
now_ts(),
|
||||
document_id,
|
||||
),
|
||||
)
|
||||
else:
|
||||
await cur.execute(
|
||||
f"""
|
||||
INSERT INTO {self.table('ai_kb_document')}
|
||||
(domain, subtype, source_id, title, summary, canonical_text, keywords, metadata_json,
|
||||
source_created_at, source_updated_at, content_hash, is_active, create_time, update_time)
|
||||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,1,%s,%s)
|
||||
""",
|
||||
(
|
||||
domain,
|
||||
subtype,
|
||||
source_id,
|
||||
payload["title"],
|
||||
payload["summary"],
|
||||
payload["canonical_text"],
|
||||
payload["keywords"],
|
||||
metadata_json,
|
||||
int(payload["source_created_at"]),
|
||||
int(payload["source_updated_at"]),
|
||||
content_hash,
|
||||
now_ts(),
|
||||
now_ts(),
|
||||
),
|
||||
)
|
||||
document_id = int(cur.lastrowid)
|
||||
|
||||
await cur.execute(f"DELETE FROM {self.table('ai_kb_chunk')} WHERE document_id=%s", (document_id,))
|
||||
dimension = int(embedding_result.get("dimension") or len(embeddings[0] or []))
|
||||
keyword_text = " ".join(filter(None, [
|
||||
str(payload.get("title") or ""),
|
||||
str(payload.get("summary") or ""),
|
||||
str(payload.get("keywords") or ""),
|
||||
json_text(payload.get("metadata_json") or {}),
|
||||
]))
|
||||
for index, chunk_text in enumerate(chunks):
|
||||
embedding = [float(value) for value in embeddings[index]]
|
||||
await cur.execute(
|
||||
f"""
|
||||
INSERT INTO {self.table('ai_kb_chunk')}
|
||||
(document_id, collection_name, chunk_index, chunk_text, embedding_json,
|
||||
embedding_model, embedding_dim, vector_norm, keyword_text, create_time, update_time)
|
||||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
||||
""",
|
||||
(
|
||||
document_id,
|
||||
COLLECTION_NAME,
|
||||
index,
|
||||
chunk_text,
|
||||
json_text(embedding),
|
||||
embedding_result.get("model") or embedder.model,
|
||||
dimension,
|
||||
vector_norm(embedding),
|
||||
keyword_text,
|
||||
now_ts(),
|
||||
now_ts(),
|
||||
),
|
||||
)
|
||||
await conn.commit()
|
||||
except Exception:
|
||||
await conn.rollback()
|
||||
raise
|
||||
return {"success": True, "document_id": document_id, "chunk_count": len(chunks)}
|
||||
|
||||
async def update_document_metadata(self, document_id: int, payload: Dict[str, Any]) -> None:
|
||||
await self.connect()
|
||||
assert self.pool is not None
|
||||
async with self.pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
f"""
|
||||
UPDATE {self.table('ai_kb_document')}
|
||||
SET title=%s, summary=%s, keywords=%s, metadata_json=%s,
|
||||
source_created_at=%s, source_updated_at=%s, update_time=%s
|
||||
WHERE id=%s
|
||||
""",
|
||||
(
|
||||
payload["title"],
|
||||
payload["summary"],
|
||||
payload["keywords"],
|
||||
json_text(payload["metadata_json"]),
|
||||
int(payload["source_created_at"]),
|
||||
int(payload["source_updated_at"]),
|
||||
now_ts(),
|
||||
document_id,
|
||||
),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
async def mark_document_inactive(self, domain: str, subtype: str, source_id: int) -> None:
|
||||
await self.connect()
|
||||
assert self.pool is not None
|
||||
async with self.pool.acquire() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
f"""
|
||||
UPDATE {self.table('ai_kb_document')}
|
||||
SET is_active=0, update_time=%s
|
||||
WHERE domain=%s AND subtype=%s AND source_id=%s
|
||||
""",
|
||||
(now_ts(), domain, subtype, source_id),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
|
||||
class EmbeddingClient:
|
||||
def __init__(self):
|
||||
self.api_key = os.environ.get("EMBEDDING_API_KEY") or os.environ.get("DASHSCOPE_API_KEY") or ""
|
||||
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)
|
||||
|
||||
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()
|
||||
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),
|
||||
}
|
||||
except Exception as exc:
|
||||
return {"success": False, "error": f"Embedding请求失败: {exc}", "embeddings": []}
|
||||
|
||||
|
||||
class KbWorker:
|
||||
def __init__(self, batch: int, block_ms: int, max_retries: int, worker_id: str):
|
||||
self.batch = max(1, batch)
|
||||
self.block_ms = max(100, block_ms)
|
||||
self.max_retries = max(0, max_retries)
|
||||
self.worker_id = worker_id
|
||||
redis_cfg = get_config().redis
|
||||
self.redis = aioredis.Redis(
|
||||
host=redis_cfg.host,
|
||||
port=redis_cfg.port,
|
||||
password=redis_cfg.password or None,
|
||||
db=redis_cfg.db,
|
||||
decode_responses=True,
|
||||
)
|
||||
self.db = KbDatabase()
|
||||
self.embedder = EmbeddingClient()
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.redis.close()
|
||||
await self.db.close()
|
||||
|
||||
async def ensure_group(self) -> None:
|
||||
try:
|
||||
await self.redis.xgroup_create(STREAM_KEY, GROUP_NAME, id="0", mkstream=True)
|
||||
logger.info("created redis stream group {} on {}", GROUP_NAME, STREAM_KEY)
|
||||
except Exception as exc:
|
||||
if "BUSYGROUP" not in str(exc):
|
||||
raise
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
await self.ensure_group()
|
||||
logger.info("KB worker started id={} batch={} block_ms={}", self.worker_id, self.batch, self.block_ms)
|
||||
while True:
|
||||
await self.consume_once()
|
||||
|
||||
async def consume_once(self) -> int:
|
||||
rows = await self.redis.xreadgroup(
|
||||
GROUP_NAME,
|
||||
self.worker_id,
|
||||
streams={STREAM_KEY: ">"},
|
||||
count=self.batch,
|
||||
block=self.block_ms,
|
||||
)
|
||||
processed = 0
|
||||
for _, messages in rows or []:
|
||||
for message_id, fields in messages:
|
||||
await self.handle_message(message_id, fields)
|
||||
processed += 1
|
||||
return processed
|
||||
|
||||
async def handle_message(self, message_id: str, fields: Dict[str, Any]) -> None:
|
||||
domain = str(fields.get("domain") or "")
|
||||
subtype = str(fields.get("subtype") or "")
|
||||
source_id = int(fields.get("source_id") or 0)
|
||||
attempts = int(fields.get("attempts") or 0)
|
||||
started = time.monotonic()
|
||||
try:
|
||||
if not domain or not subtype or source_id <= 0:
|
||||
raise ValueError("消息缺少 domain/subtype/source_id")
|
||||
result = await self.db.upsert_document(domain, subtype, source_id, self.embedder)
|
||||
if not result.get("success"):
|
||||
raise RuntimeError(str(result.get("error") or "upsert failed"))
|
||||
await self.redis.xack(STREAM_KEY, GROUP_NAME, message_id)
|
||||
logger.info(
|
||||
"kb upsert ok message={} source={}/{}/{} result={} elapsed={:.2f}s",
|
||||
message_id,
|
||||
domain,
|
||||
subtype,
|
||||
source_id,
|
||||
result,
|
||||
time.monotonic() - started,
|
||||
)
|
||||
except Exception as exc:
|
||||
await self.handle_failure(message_id, fields, attempts, str(exc))
|
||||
|
||||
async def handle_failure(self, message_id: str, fields: Dict[str, Any], attempts: int, error: str) -> None:
|
||||
next_attempt = attempts + 1
|
||||
payload = dict(fields)
|
||||
payload["attempts"] = str(next_attempt)
|
||||
payload["last_error"] = error[:500]
|
||||
payload["failed_at"] = str(now_ts())
|
||||
if next_attempt > self.max_retries:
|
||||
await self.redis.xadd(DEAD_STREAM_KEY, payload)
|
||||
await self.redis.xack(STREAM_KEY, GROUP_NAME, message_id)
|
||||
logger.error("kb message={} moved to dead stream after {} attempts: {}", message_id, next_attempt, error)
|
||||
return
|
||||
await self.redis.xadd(STREAM_KEY, payload)
|
||||
await self.redis.xack(STREAM_KEY, GROUP_NAME, message_id)
|
||||
logger.warning("kb message={} retry {}/{}: {}", message_id, next_attempt, self.max_retries, error)
|
||||
|
||||
|
||||
def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Sport Era AI KB Redis Stream worker")
|
||||
parser.add_argument("--batch", type=int, default=50)
|
||||
parser.add_argument("--block-ms", type=int, default=5000)
|
||||
parser.add_argument("--max-retries", type=int, default=3)
|
||||
parser.add_argument("--worker-id", default="")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
async def async_main(argv: Optional[List[str]] = None) -> int:
|
||||
args = parse_args(argv)
|
||||
load_config()
|
||||
worker_id = args.worker_id or f"{socket.gethostname()}-{os.getpid()}"
|
||||
worker = KbWorker(args.batch, args.block_ms, args.max_retries, worker_id)
|
||||
try:
|
||||
await worker.run_forever()
|
||||
finally:
|
||||
await worker.close()
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> int:
|
||||
return asyncio.run(async_main(argv))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user