feat(govai): 0617 优化首批 — 安全/私有化/深度研究/服务层/可观测性
借鉴 odysseus 的能力设计,全程净室实现、零 AGPL 代码、不引入 AGPL 依赖。
T1 提示注入防护: pkg/promptguard 包裹外部/知识库内容为不可信数据,buildMessages 移出 system 指令区。 T2 安全 CI: .github/workflows(ci+security: govulncheck/gitleaks/actionlint/hadolint/trivy)+dependabot+.hadolint.yaml;go.mod 加 toolchain go1.25.11 修复 20 个 stdlib CVE。 T3 管理员 2FA: 迁移 000016 + RFC6238 TOTP/备份码(pkg/auth, 零依赖) + 登录流程集成(后端)。 T4 本地模型: LLM/embedding 支持本地 vLLM/Ollama(OpenAI 兼容, 鉴权头条件发送, NoAuth) + docs/local-deploy.md。 T6 深度研究: 迁移 000017 + Python research-worker(净室多步流水线, 检索避开 SearXNG) + Go research 服务/handler/路由。 T7 service 层: 新增 internal/service/{research,twofa}, 2FA 业务逻辑从胖 handler 下沉, 接口注入可单测。 T10 缓存/可观测性: internal/cache(Redis+内存, 优雅降级) 接入 store 热点列表; Prometheus 指标+/metrics; docs/openapi.yaml。 验证: go build/vet/test ./... 全绿(8 包); research-worker 12 单测过; 真实 PG 应用迁移并烟测。
This commit is contained in:
@@ -0,0 +1,97 @@
|
||||
"""任务消费者 — 从 Redis 队列取研究任务并执行流水线。仿 ppt-worker 模式。"""
|
||||
|
||||
import json
|
||||
import time
|
||||
import signal
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import redis
|
||||
|
||||
from config import config
|
||||
from db import get_task, update_task_status, is_canceled
|
||||
import llm_client
|
||||
import search
|
||||
from pipeline import ResearchPipeline, CanceledError
|
||||
|
||||
|
||||
class ResearchWorker:
|
||||
def __init__(self):
|
||||
self.redis = redis.from_url(config.REDIS_URL, decode_responses=True)
|
||||
self.executor = ThreadPoolExecutor(max_workers=config.CONCURRENCY)
|
||||
self.running = True
|
||||
try:
|
||||
signal.signal(signal.SIGINT, self._shutdown)
|
||||
signal.signal(signal.SIGTERM, self._shutdown)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
def _shutdown(self, signum, frame):
|
||||
print(f"\n[Research] 收到信号 {signum},正在优雅关闭…")
|
||||
self.running = False
|
||||
|
||||
def start(self):
|
||||
print(f"[Research] 启动,并发数: {config.CONCURRENCY},监听队列: {config.TASK_QUEUE}")
|
||||
while self.running:
|
||||
try:
|
||||
result = self.redis.brpop(config.TASK_QUEUE, timeout=5)
|
||||
if result is None:
|
||||
continue
|
||||
_, raw = result
|
||||
task_id = json.loads(raw).get("task_id")
|
||||
if task_id:
|
||||
self.executor.submit(self._process, task_id)
|
||||
except redis.ConnectionError as e:
|
||||
print(f"[Research] Redis 连接失败: {e},5 秒后重试…")
|
||||
time.sleep(5)
|
||||
except Exception as e:
|
||||
print(f"[Research] 未知错误: {e}")
|
||||
time.sleep(1)
|
||||
self.executor.shutdown(wait=True)
|
||||
|
||||
def _redis_status(self, task_id, status, progress, message):
|
||||
key = f"{config.TASK_STATUS_PREFIX}{task_id}"
|
||||
self.redis.hset(key, mapping={"status": status, "progress": str(progress), "message": message})
|
||||
self.redis.expire(key, 3600)
|
||||
|
||||
def _process(self, task_id: str):
|
||||
try:
|
||||
task = get_task(task_id)
|
||||
if not task:
|
||||
print(f"[Research] 任务不存在: {task_id}")
|
||||
return
|
||||
if task["status"] != "pending":
|
||||
print(f"[Research] 跳过非 pending 任务: {task_id} ({task['status']})")
|
||||
return
|
||||
|
||||
def progress_cb(status, progress, message):
|
||||
# 同时更新 Redis(快速轮询)与 DB(Go 端读取)
|
||||
self._redis_status(task_id, status, progress, message)
|
||||
update_task_status(task_id, status, progress=progress, status_message=message)
|
||||
|
||||
pipeline = ResearchPipeline(
|
||||
task_id, task,
|
||||
llm_chat=llm_client.chat,
|
||||
search_fn=search.search,
|
||||
fetch_fn=search.fetch_text,
|
||||
update_fn=lambda **kw: update_task_status(task_id, **kw),
|
||||
progress_cb=progress_cb,
|
||||
is_canceled=lambda: is_canceled(task_id),
|
||||
max_subqueries=config.MAX_SUBQUERIES,
|
||||
max_sources=config.MAX_SOURCES,
|
||||
)
|
||||
pipeline.run()
|
||||
self._redis_status(task_id, "completed", 100, "完成")
|
||||
print(f"[Research] 任务完成: {task_id}")
|
||||
|
||||
except CanceledError:
|
||||
update_task_status(task_id, "canceled", status_message="任务已取消")
|
||||
self._redis_status(task_id, "canceled", 0, "已取消")
|
||||
print(f"[Research] 任务取消: {task_id}")
|
||||
except Exception as e:
|
||||
update_task_status(task_id, "failed", error_message=str(e))
|
||||
self._redis_status(task_id, "failed", 0, f"失败: {str(e)[:200]}")
|
||||
print(f"[Research] 任务失败: {task_id} - {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
ResearchWorker().start()
|
||||
Reference in New Issue
Block a user