from typing import TypedDict
from langgraph.constants import START, END
from langgraph.graph import StateGraph

# 1. 定义图状态 (TypedDict)
class AgentState(TypedDict):
    query: str
    category: str
    response: str

# 2. 定义节点函数 (仅返回增量更新字典)
def classify_node(state: AgentState) -> dict:
    category = "technical" if "代码" in state["query"] else "billing"
    return {"category": category}

def tech_node(state: AgentState) -> dict:
    return {"response": f"【技术支持】已处理关于“{state['query']}”的技术咨询。"}

def billing_node(state: AgentState) -> dict:
    return {"response": f"【账单服务】已查询关于“{state['query']}”的账单明细。"}

# 3. 条件路由函数
def route_intent(state: AgentState) -> str:
    return state["category"]

# 4. 构建图与配置动态条件边 (Conditional Edges)
builder = StateGraph(AgentState)

builder.add_node("classify", classify_node)
builder.add_node("tech_support", tech_node)
builder.add_node("billing_support", billing_node)

builder.add_edge(START, "classify")
builder.add_conditional_edges(
    "classify",
    route_intent,
    {
        "technical": "tech_support",
        "billing": "billing_support",
    },
)
builder.add_edge("tech_support", END)
builder.add_edge("billing_support", END)

# 5. 编译与执行
graph = builder.compile()
final_state = graph.invoke({"query": "请问这段 Python 代码报错怎么解决？"})
print(final_state["response"])

from typing import TypedDict, List
from typing_extensions import NotRequired

# 可设置 total=False 允许实例化时只传入部分初始字段
class AgentState(TypedDict, total=False):
    user_query: str
    conversation_history: List[str]
    search_results: NotRequired[List[str]]  # 可选增量字段

from pydantic import BaseModel, Field
from typing import Optional

class AgentStateModel(BaseModel):
    user_id: str = Field(description="用户标识")
    query: str = Field(min_length=1, description="用户问题")
    final_answer: Optional[str] = None

from typing_extensions import TypedDict, Annotated
import operator
from langgraph.graph import add_messages

class StateWithReducer(TypedDict):
    current_step: str                        # 默认覆盖
    logs: Annotated[list, operator.add]      # 列表追加
    messages: Annotated[list, add_messages]  # 消息自动化合并

# 错误做法：直接原地修改并返回整个 state
def wrong_node(state: MyState) -> MyState:
    state["final_answer"] = "done"
    return state

# 正确做法：仅返回变更字段的增量字典
def correct_node(state: MyState) -> dict:
    return {"final_answer": "done"}

from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.runtime import Runtime

# 1. 定义 Context 依赖类型与 State 类型
class AppContext(TypedDict):
    db_service: object

class AppState(TypedDict):
    query: str
    result: str

# 2. 节点中通过 runtime.context 获取注入的服务
def process_node(state: AppState, runtime: Runtime[AppContext]) -> dict:
    db = runtime.context["db_service"]
    # 使用 db 服务进行查询...
    return {"result": "processed"}

# 3. 绑定 state_schema 与 context_schema
builder = StateGraph(state_schema=AppState, context_schema=AppContext)
builder.add_node("process", process_node)
builder.add_edge(START, "process")
builder.add_edge("process", END)

graph = builder.compile()

# 4. 执行时传入 context 实体
res = graph.invoke(
    {"query": "test"},
    context={"db_service": "DatabaseInstance"}
)

# 顺序执行：node_a -> node_b
builder.add_edge("node_a", "node_b")

# 并行并发：node_a 分发至 node_b 与 node_c
builder.add_edge("node_a", "node_b")
builder.add_edge("node_a", "node_c")

from typing import TypedDict
from langgraph.graph import StateGraph, START, END

class RouteState(TypedDict):
    query: str
    need_web_search: bool

def route_decision(state: RouteState):
    if state["need_web_search"]:
        # 返回列表触发并行处理
        return ["rag_node", "web_node"]
    return "rag_node"

builder = StateGraph(RouteState)
builder.add_node("rag_node", lambda s: {"result": "rag"})
builder.add_node("web_node", lambda s: {"result": "web"})

builder.add_conditional_edges(
    START,
    route_decision,
    {
        "rag_node": "rag_node",
        "web_node": "web_node"
    }
)

# 当校验节点失败时流转至纠错节点，纠错完成后重新指向校验节点形成回环
builder.add_conditional_edges(
    "validate_node",
    lambda state: "execute_node" if state.get("error") is None else "correct_node",
    {
        "correct_node": "correct_node",
        "execute_node": "execute_node"
    }
)
# 纠错后重新校验
builder.add_edge("correct_node", "validate_node")
builder.add_edge("execute_node", END)

# 设置最大执行步数为 10 步
result = graph.invoke(
    {"query": "input"},
    config={"recursion_limit": 10}
)

from langgraph.types import Command
from typing import Literal

def decision_node(state: dict) -> Command[Literal["node_b", "node_c"]]:
    if state.get("score", 0) > 80:
        return Command(update={"status": "passed"}, goto="node_c")
    return Command(update={"status": "failed"}, goto="node_b")

from langgraph.types import Send
from typing import Annotated
import operator

class MapReduceState(TypedDict):
    items: list[str]
    results: Annotated[list, operator.add]

def map_dispatcher(state: MapReduceState):
    # 动态为每个 item 生成一个 worker 节点任务
    return [Send("worker_node", {"item": item}) for item in state["items"]]

def worker_node(state: dict) -> dict:
    return {"results": [f"Processed {state['item']}"]}

# 同步调用
final_state = graph.invoke({"query": "hello"})

# 异步调用
final_state = await graph.ainvoke({"query": "hello"})

class KBQueryWorkflow:
    def __init__(self):
        # 初始化并编译图对象
        self.app = create_query_graph()
        
    async def run(self, initial_state: dict):
        return await self.app.ainvoke(initial_state)
        
    async def stream(self, initial_state: dict):
        async for output in self.app.astream(initial_state):
            yield output

# 1. updates 模式：接收每一步节点产出的状态增量
async for event in graph.astream(initial_state, stream_mode="updates"):
    print(f"Node Update: {event}")

# 2. values 模式：接收每一步执行结束后的全量 State 视图
async for state_snapshot in graph.astream(initial_state, stream_mode="values"):
    print(f"State Snapshot: {state_snapshot}")

# 3. messages 模式：捕获大模型节点的流式 Token 输出
async for message_chunk, metadata in graph.astream(initial_state, stream_mode="messages"):
    if message_chunk.content:
        print(message_chunk.content, end="", flush=True)

from typing import Annotated
from typing_extensions import TypedDict
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END, add_messages
from langgraph.prebuilt import ToolNode, tools_condition

# 1. 定义工具
@tool
def multiply(a: int, b: int) -> int:
    """计算两个整数的乘积。"""
    return a * b

tools = [multiply]

# 2. 定义状态模式
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 3. 绑定工具到模型
llm = ChatOpenAI(model="gpt-4o-mini").bind_tools(tools)

def chatbot_node(state: State) -> dict:
    return {"messages": [llm.invoke(state["messages"])]}

# 4. 构建图结构
builder = StateGraph(State)
builder.add_node("chatbot", chatbot_node)
builder.add_node("tools", ToolNode(tools))  # 实例化 ToolNode 节点

builder.add_edge(START, "chatbot")

# 结合 tools_condition 动态判断：若模型产生 tool_calls 则转向 "tools" 节点，否则终止转向 END
builder.add_conditional_edges(
    "chatbot",
    tools_condition,
    {
        "tools": "tools",
        "__end__": END
    }
)
builder.add_edge("tools", "chatbot")  # 工具执行完成后返回 chatbot 节点

graph = builder.compile()

from langgraph.types import RetryPolicy

builder.add_node(
    "search_node",
    search_node_func,
    retry_policy=RetryPolicy(
        max_attempts=3,
        initial_interval=1.0,
        retry_on=(ConnectionError, TimeoutError)
    )
)

from langgraph.prebuilt import ToolNode

# handle_tool_errors=True 会自动将工具执行报错转化为 ToolMessage 节点输出
tool_node = ToolNode(tools=[fetch_user_data], handle_tool_errors=True)
builder.add_node("tools", tool_node)