from typing import TypedDict, List
from langgraph.graph import StateGraph, END
from anthropic import Anthropic

# 1. 定义智能体状态
class AgentState(TypedDict):
    messages: List[dict]
    screenshot_base64: str | None
    retry_count: int

# 2. 主管节点 (Supervisor - Claude Opus 4.5)
async def supervisor_node(state: AgentState):
    client = Anthropic()
    
    # 【SOTA优化】启用 Prompt Caching
    # 将包含几千行 MCP 工具定义的 System Prompt 缓存，大幅降低成本
    system_prompt = {
        "type": "text", 
        "text": "你是全能智能体主管。利用 MCP 工具完成任务...",
        "cache_control": {"type": "ephemeral"} 
    }

    response = client.beta.messages.create(
        model="claude-3-opus-20250601",
        max_tokens=4096,
        system=[system_prompt],
        messages=state["messages"],
        tools=mcp_tools, # 动态加载 E2B, Tavily 等 MCP 工具
        betas=["computer-use-2025", "prompt-caching-20240731"]
    )
    return {"messages": [response]}

# 3. 视觉监控节点 (Monitor - Gemini 3 Flash)
async def monitor_node(state: AgentState):
    # 使用极低成本模型检查页面加载状态 (Loading / Error / Done)
    status = await gemini_flash.check_screen(state["screenshot_base64"])
    
    if status == "loading": 
        return {"action": "wait"}
    if status == "error": 
        return {"retry_count": state["retry_count"] + 1}
        
    return {"action": "continue"}

# 4. 工作流编排
workflow = StateGraph(AgentState)
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("tool_executor", tool_node) 
workflow.add_node("monitor", monitor_node)

workflow.set_entry_point("supervisor")
workflow.add_edge("tool_executor", "monitor") # 执行后先监控
workflow.add_conditional_edges("monitor", monitor_router) 

app = workflow.compile()