nanocode/TradingAgents/web_server.py

"""
TradingAgents Web API Server
FastAPI wrapper for the TradingAgents multi-agent trading framework.
"""
import os
import json
import datetime
import threading
from pathlib import Path

from dotenv import load_dotenv
load_dotenv()

from fastapi import FastAPI, HTTPException, BackgroundTasks, Request, Response, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import uvicorn

from jose import jwt, JWTError

from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG

# ── SSO 配置（与 SSO 中心共用）──────────────────
SECRET_KEY = "your-secret-key-change-in-production"
ALGORITHM = "HS256"
SSO_CENTER_URL = "http://129.211.173.87:9000"
COOKIE_DOMAIN = "129.211.173.87"

app = FastAPI(
    title="TradingAgents API",
    description="Multi-Agent LLM Financial Trading Framework",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── 全局任务存储 ──────────────────────────────────
tasks_lock = threading.Lock()
tasks_store = {}  # task_id -> {status, result, error, ...}
TASKS_DIR = Path(__file__).parent / "logs" / "tasks"
TASKS_DIR.mkdir(parents=True, exist_ok=True)

# 启动时从磁盘加载已有任务
def _load_tasks_from_disk():
    for f in sorted(TASKS_DIR.glob("*.json")):
        try:
            data = json.loads(f.read_text())
            tid = f.stem
            # 服务重启后，running/pending 任务已丢失后台线程，标记为失败
            if data.get("status") in ("running", "pending"):
                data["status"] = "failed"
                data["error"] = "服务重启，分析中断，请重新提交"
                f.write_text(json.dumps(data, ensure_ascii=False, indent=2))
            tasks_store[tid] = data
        except Exception:
            pass

_load_tasks_from_disk()


def _save_task_to_disk(task_id: str):
    """将任务写入磁盘 JSON 文件"""
    with tasks_lock:
        data = tasks_store.get(task_id)
    if data:
        (TASKS_DIR / f"{task_id}.json").write_text(json.dumps(data, ensure_ascii=False, indent=2))


# ── SSO 认证 ─────────────────────────────────────
async def require_sso(request: Request) -> str:
    """从 cookie 中验证 JWT，返回用户名"""
    token = request.cookies.get("access_token")
    if not token:
        raise HTTPException(401, detail="未登录")
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username = payload.get("sub")
        if not username:
            raise HTTPException(401, detail="无效令牌")
        return username
    except JWTError:
        raise HTTPException(401, detail="令牌已过期或无效")


# ── 请求/响应模型 ──────────────────────────────────
class AnalysisRequest(BaseModel):
    ticker: str
    date: Optional[str] = None  # YYYY-MM-DD, defaults to today
    analysts: Optional[list[str]] = ["market", "news"]
    max_debate_rounds: Optional[int] = 1
    output_language: Optional[str] = "Chinese"  # "Chinese" or "English"


# ── 后台分析任务 ──────────────────────────────────
def run_analysis_task(task_id: str, ticker: str, date: str, analysts: list, max_debate_rounds: int, output_language: str = "English"):
    config = DEFAULT_CONFIG.copy()
    config["max_debate_rounds"] = max_debate_rounds
    config["output_language"] = output_language

    try:
        with tasks_lock:
            tasks_store[task_id]["status"] = "running"
        _save_task_to_disk(task_id)

        graph = TradingAgentsGraph(selected_analysts=analysts, debug=False, config=config)
        state, decision = graph.propagate(ticker, date)

        with tasks_lock:
            tasks_store[task_id]["status"] = "completed"
            tasks_store[task_id]["result"] = {
                "ticker": ticker,
                "date": date,
                "decision": decision,
                "reports": {
                    k: v for k, v in state.items()
                    if k.endswith("_report") or k in ("investment_plan", "trader_investment_plan", "final_trade_decision")
                },
            }
    except Exception as e:
        with tasks_lock:
            tasks_store[task_id]["status"] = "failed"
            tasks_store[task_id]["error"] = str(e)
    finally:
        _save_task_to_disk(task_id)


# ── Web 页面 ──────────────────────────────────────
WEB_HTML = Path(__file__).parent / "web.html"


@app.get("/")
def index():
    from fastapi.responses import HTMLResponse
    return HTMLResponse(WEB_HTML.read_text("utf-8"))


@app.get("/api/health")
def health():
    return {"status": "ok", "time": datetime.datetime.now().isoformat()}


@app.post("/api/analyze")
def analyze(req: AnalysisRequest, background_tasks: BackgroundTasks, username: str = Depends(require_sso)):
    date = req.date or datetime.date.today().isoformat()
    task_id = f"{req.ticker}_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}"

    with tasks_lock:
        tasks_store[task_id] = {
            "task_id": task_id,
            "status": "pending",
            "ticker": req.ticker,
            "date": date,
            "result": None,
            "error": None,
            "created_at": datetime.datetime.now().isoformat(),
        }
    _save_task_to_disk(task_id)

    background_tasks.add_task(
        run_analysis_task, task_id, req.ticker, date, req.analysts, req.max_debate_rounds, req.output_language
    )
    return {"task_id": task_id, "status": "pending", "message": "Analysis started"}


@app.get("/api/task/{task_id}")
def get_task(task_id: str):
    with tasks_lock:
        task = tasks_store.get(task_id)
    if task:
        return task
    # 回退到磁盘（服务重启后内存丢失，但磁盘还在）
    disk_file = TASKS_DIR / f"{task_id}.json"
    if disk_file.exists():
        try:
            return json.loads(disk_file.read_text())
        except Exception:
            pass
    raise HTTPException(status_code=404, detail="Task not found")


# ── SSO 用户接口 ──────────────────────────────────
@app.get("/api/user")
async def api_user(username: str = Depends(require_sso)):
    """获取当前登录用户信息"""
    return {"username": username, "logged_in": True}


@app.post("/api/logout")
def api_logout(response: Response):
    """清除登录 cookie"""
    response.delete_cookie("access_token", domain=COOKIE_DOMAIN, path="/")
    return {"success": True}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8500)