Files
GovAI/research-worker/app.py
T
freedakgmail c949204662 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 应用迁移并烟测。
2026-06-17 17:52:47 +08:00

90 lines
2.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Research Worker HTTP APIFastAPI)— 供 Go 后端调用 / 调试。
任务正常由 Go 后端写入 research_tasks 并 LPush 到 Redis 队列;本服务的后台线程消费队列。
这里另外提供状态查询与健康检查端点。
"""
import json
import uuid
import threading
from typing import Optional
import psycopg
import redis
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from config import config
from db import get_task
from worker import ResearchWorker
app = FastAPI(title="Research Worker API", version="1.0.0")
rdb = redis.from_url(config.REDIS_URL, decode_responses=True)
class CreateTaskRequest(BaseModel):
user_id: str
topic: str
app_id: Optional[str] = None
config: dict = {}
class TaskStatusResponse(BaseModel):
task_id: str
status: str
progress: int
status_message: Optional[str] = None
error_message: Optional[str] = None
@app.post("/api/tasks")
def create_task(req: CreateTaskRequest):
task_id = str(uuid.uuid4())
with psycopg.connect(config.DATABASE_URL) as conn:
with conn.cursor() as cur:
cur.execute(
"INSERT INTO research_tasks (id, user_id, app_id, topic, config) "
"VALUES (%s, %s, %s, %s, %s)",
(task_id, req.user_id, req.app_id, req.topic, json.dumps(req.config)),
)
conn.commit()
rdb.lpush(config.TASK_QUEUE, json.dumps({"task_id": task_id}))
return {"task_id": task_id, "status": "pending"}
@app.get("/api/tasks/{task_id}", response_model=TaskStatusResponse)
def get_task_status(task_id: str):
cached = rdb.hgetall(f"{config.TASK_STATUS_PREFIX}{task_id}")
if cached:
return TaskStatusResponse(
task_id=task_id,
status=cached.get("status", "unknown"),
progress=int(cached.get("progress", 0)),
status_message=cached.get("message"),
)
task = get_task(task_id)
if not task:
raise HTTPException(status_code=404, detail="任务不存在")
return TaskStatusResponse(
task_id=task_id,
status=task["status"],
progress=task["progress"],
status_message=task.get("status_message"),
error_message=task.get("error_message"),
)
@app.get("/health")
def health():
return {"status": "ok", "service": "research-worker"}
def _start_worker():
ResearchWorker().start()
if __name__ == "__main__":
threading.Thread(target=_start_worker, daemon=True).start()
uvicorn.run(app, host=config.HOST, port=config.PORT)