""" app.py — LangGraph 流程編排 Space 的 FastAPI 入口 提供 SSE streaming，讓 Open WebUI 即時看到工作進度 """ import uuid import time import httpx import os import json import asyncio from typing import AsyncGenerator from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel from graph import company_graph, CompanyState app = FastAPI( title="AI 公司 — 流程編排器", description="LangGraph 驅動的任務編排 API", version="1.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) class TaskRequest(BaseModel): message: str session_id: str = "default" class TaskResponse(BaseModel): session_id: str response: str intent: str department: str @app.get("/") def root(): return { "service": "LangGraph Orchestrator", "endpoints": ["/run", "/stream", "/health", "/graph-schema"], } @app.get("/health") def health(): return {"status": "ok"} @app.get("/graph-schema") def graph_schema(): """回傳圖的結構（用於除錯）""" return {"nodes": list(company_graph.nodes.keys())} @app.post("/run", response_model=TaskResponse) async def run_task(req: TaskRequest): """同步執行任務（適合短任務）""" initial_state: CompanyState = { "user_input": req.message, "session_id": req.session_id, "intent": None, "department": None, "job_id": None, "job_status": None, "retry_count": 0, "crew_result": None, "openhands_result": None, "final_response": "", "error": None, "messages": [{"role": "user", "content": req.message}], } try: final_state = await asyncio.to_thread(company_graph.invoke, initial_state) return TaskResponse( session_id=req.session_id, response=final_state.get("final_response", ""), intent=final_state.get("intent", "unknown"), department=final_state.get("department", "unknown"), ) except Exception as e: raise HTTPException(500, str(e)) @app.post("/stream") async def stream_task(req: TaskRequest): """ SSE 串流執行（適合長任務） Open WebUI 的 pipeline 會訂閱這個端點 """ async def event_generator() -> AsyncGenerator[str, None]: initial_state: CompanyState = { "user_input": req.message, "session_id": req.session_id, "intent": None, "department": None, "job_id": None, "job_status": None, "retry_count": 0, "crew_result": None, "openhands_result": None, "final_response": "", "error": None, "messages": [{"role": "user", "content": req.message}], } try: # 用 stream_mode="updates" 逐步推送每個節點的更新 for chunk in company_graph.stream( initial_state, stream_mode="updates", ): for node_name, node_output in chunk.items(): # 過濾不重要的中間狀態 progress_msg = None if node_name == "classify": intent = node_output.get("intent", "?") dept = node_output.get("department", "?") progress_msg = f"🔍 任務分類完成：{intent}（部門：{dept}）" elif node_name == "route": progress_msg = f"📋 任務路由中..." elif node_name in ("crew_strategy", "crew_prd", "crew_marketing", "crew_research"): labels = { "crew_strategy": "CEO 制定策略中", "crew_prd": "PM 撰寫 PRD 中", "crew_marketing": "市場部創作中", "crew_research": "分析師研究中", } progress_msg = f"🤝 {labels[node_name]}..." elif node_name in ("openhands_code", "openhands_review"): labels = { "openhands_code": "🛠️ 工程團隊撰寫程式碼中", "openhands_review": "🔎 工程團隊執行 Code Review 中", } progress_msg = f"{labels[node_name]}..." elif node_name == "wait_job": retry = node_output.get("retry_count", 0) status = node_output.get("job_status", "running") if status != "done": progress_msg = f"⏳ 等待任務完成... ({retry * 5}s)" elif node_name == "format": result = node_output.get("final_response", "") data = json.dumps({"type": "result", "content": result}) yield f"data: {data}\n\n" return elif node_name == "error_handler": error = node_output.get("final_response", "Error") data = json.dumps({"type": "error", "content": error}) yield f"data: {data}\n\n" return if progress_msg: data = json.dumps({"type": "progress", "content": progress_msg}) yield f"data: {data}\n\n" await asyncio.sleep(0.1) except Exception as e: data = json.dumps({"type": "error", "content": f"Orchestration error: {e}"}) yield f"data: {data}\n\n" return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, ) # ... 原有的 /run, /stream, /health, /graph-schema 端點 ... @app.get("/v1/models") def list_models(): return { "object": "list", "data": [ {"id": "ai-company/auto", "object": "model", "owned_by": "ai-company"}, {"id": "ai-company/strategy", "object": "model", "owned_by": "ai-company"}, {"id": "ai-company/engineering", "object": "model", "owned_by": "ai-company"}, {"id": "ai-company/marketing", "object": "model", "owned_by": "ai-company"}, ] } @app.post("/v1/chat/completions") async def openai_compatible(req: dict): message = req["messages"][-1]["content"] session_id = req.get("user", "default") model = req.get("model", "ai-company") async def generate(): buffer = "" async with httpx.AsyncClient(timeout=300) as client: async with client.stream( "POST", "http://localhost:7860/stream", json={"message": message, "session_id": session_id}, ) as resp: async for chunk in resp.aiter_text(): buffer += chunk while "\n" in buffer: line, buffer = buffer.split("\n", 1) line = line.strip() if not line.startswith("data: "): continue try: data = json.loads(line[6:]) except Exception: continue content = data.get("content", "") delta = { "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion.chunk", "created": int(time.time()), "model": model, "choices": [{"delta": {"content": content}, "index": 0, "finish_reason": None}], } yield f"data: {json.dumps(delta)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(generate(), media_type="text/event-stream") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)