freedakgmail
c949204662
借鉴 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 应用迁移并烟测。
98 lines
3.8 KiB
Python
98 lines
3.8 KiB
Python
"""任务消费者 — 从 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()
|