vent_agent_2.0/mock_result_platform/app.py

from contextlib import contextmanager
from datetime import datetime
import hmac
import json
from pathlib import Path
import sqlite3
import threading
import uuid
from zoneinfo import ZoneInfo

from fastapi import Depends, FastAPI, Header, HTTPException, Query
from fastapi.responses import RedirectResponse
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

from runtime.result_models import PlatformResultPayload

LOCAL_TOKEN = "local-result-platform-token"


class ResultStore:
    def __init__(self, path):
        self.path = Path(path).resolve()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self.connect() as conn:
            conn.execute("""CREATE TABLE IF NOT EXISTS received_judgments (
                agent_code TEXT NOT NULL, task_id TEXT NOT NULL, payload TEXT NOT NULL,
                receipt TEXT NOT NULL, received_at TEXT NOT NULL, PRIMARY KEY (agent_code, task_id))""")

    @contextmanager
    def connect(self):
        conn = sqlite3.connect(str(self.path), timeout=30)
        conn.row_factory = sqlite3.Row
        try:
            with conn:
                yield conn
        finally:
            conn.close()

    def accept(self, payload):
        text = json.dumps(payload, ensure_ascii=False, sort_keys=True, allow_nan=False)
        with self.connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            prior = conn.execute("SELECT payload, receipt FROM received_judgments WHERE agent_code = ? AND task_id = ?",
                                 (payload["agentCode"], payload["taskId"])).fetchone()
            if prior:
                if prior["payload"] != text:
                    raise ValueError("同一个 agentCode/taskId 的报文内容不一致")
                return json.loads(prior["receipt"])
            now = datetime.now(ZoneInfo("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
            receipt = {"code": 200, "success": True, "data": {
                "taskId": payload["taskId"], "resultId": uuid.uuid4().hex, "receivedAt": now}}
            conn.execute("INSERT INTO received_judgments VALUES (?, ?, ?, ?, ?)",
                         (payload["agentCode"], payload["taskId"], text, json.dumps(receipt, ensure_ascii=False), now))
            return receipt

    def rows(self, limit=50):
        with self.connect() as conn:
            rows = conn.execute("SELECT payload, receipt, received_at FROM received_judgments ORDER BY received_at DESC, rowid DESC LIMIT ?", (limit,)).fetchall()
            return [{"payload": json.loads(row["payload"]), "receipt": json.loads(row["receipt"]),
                     "receivedAt": row["received_at"]} for row in rows]

    def count(self):
        with self.connect() as conn:
            return conn.execute("SELECT COUNT(*) FROM received_judgments").fetchone()[0]


def create_app(store, *, token=LOCAL_TOKEN, fail_first=0):
    if not token:
        raise ValueError("本地结果平台必须设置访问凭证")
    app = FastAPI(title="主通风机结果接收平台", description="独立接收智能体研判 JSON，与取数网关分开运行。")
    app.state.store = store
    remaining, lock = [fail_first], threading.Lock()
    scheme = HTTPBearer(auto_error=False)

    def auth(credentials: HTTPAuthorizationCredentials | None = Depends(scheme)):
        if (credentials is None or credentials.scheme.lower() != "bearer"
                or not hmac.compare_digest(credentials.credentials.encode(), token.encode())):
            raise HTTPException(401, "结果平台认证失败")

    @app.get("/", include_in_schema=False)
    def home():
        return RedirectResponse("/docs")

    @app.get("/health")
    def health():
        return {"status": "ok", "received": store.count()}

    @app.post("/api/judgments", dependencies=[Depends(auth)])
    def receive(payload: PlatformResultPayload, idempotency_key: str = Header(default="")):
        if idempotency_key and idempotency_key != payload.taskId:
            raise HTTPException(409, "幂等键与 taskId 不一致")
        with lock:
            if remaining[0] > 0:
                remaining[0] -= 1
                raise HTTPException(503, "接收服务暂不可用", headers={"Retry-After": "2"})
        try:
            return store.accept(payload.model_dump(mode="json"))
        except ValueError as exc:
            raise HTTPException(409, str(exc)) from None

    @app.get("/api/judgments", dependencies=[Depends(auth)])
    def received(limit: int = Query(default=50, ge=1, le=200)):
        return {"items": store.rows(limit)}

    return app