695 lines
30 KiB
Python
695 lines
30 KiB
Python
#!/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())
|