app.py

"""
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)