from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage
import operator

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    research_notes: str
    code_output: str
    review_feedback: str
    iteration_count: int

# 初始状态
state = {"messages": [msg1, msg2]}

# 节点 A 返回的更新
update_a = {"messages": [msg3]}

# 节点 B 返回的更新
update_b = {"messages": [msg4]}

# LangGraph 使用 operator.add 合并更新
state["messages"] = operator.add(state["messages"], update_a["messages"])
# 现在 state["messages"] = [msg1, msg2, msg3]

state["messages"] = operator.add(state["messages"], update_b["messages"])
# 最终 state["messages"] = [msg1, msg2, msg3, msg4]

def merge_dicts(existing: dict, new: dict) -> dict:
    """自定义 Reducer：深度合并两个字典"""
    result = existing.copy()
    result.update(new)
    return result

class ComplexState(TypedDict):
    config: Annotated[dict, merge_dicts]
    results: Annotated[list, operator.add]

from langgraph.graph import StateGraph, END

# 1. 定义状态
class SimpleState(TypedDict):
    input: str
    output: str
    counter: int

# 2. 定义节点函数
def node_1(state: SimpleState) -> SimpleState:
    """第一个节点：处理输入"""
    print("Running Node 1")
    return {"counter": state.get("counter", 0) + 1}

def node_2(state: SimpleState) -> SimpleState:
    """第二个节点：生成输出"""
    print("Running Node 2")
    processed = f"Processed: {state['input']}"
    return {"output": processed, "counter": state["counter"] + 1}

# 3. 创建图
graph = StateGraph(SimpleState)

# 4. 添加节点
graph.add_node("node1", node_1)
graph.add_node("node2", node_2)

# 5. 设置入口点
graph.set_entry_point("node1")

# 6. 添加边
graph.add_edge("node1", "node2")
graph.add_edge("node2", END)

# 7. 编译图
app = graph.compile()

result = app.invoke({"input": "Hello, LangGraph!"})
print(result)
# {'input': 'Hello, LangGraph!', 'counter': 2, 'output': 'Processed: Hello, LangGraph!'}

def should_continue(state: SimpleState) -> str:
    """路由函数：决定下一步去哪里"""
    if state["counter"] < 5:
        return "node1"  # 继续循环
    else:
        return "node2"  # 结束循环

# 修改图的构建
graph = StateGraph(SimpleState)
graph.add_node("node1", node_1)
graph.add_node("node2", node_2)
graph.set_entry_point("node1")

# 使用条件边
graph.add_conditional_edges(
    "node1",  # 源节点
    should_continue,  # 路由函数
    {
        "node1": "node1",  # 映射返回值到目标节点
        "node2": "node2"
    }
)
graph.add_edge("node2", END)

app = graph.compile()

from langgraph.graph import MessageGraph
from langchain_openai import ChatOpenAI

model = ChatOpenAI(model="gpt-4o")

def call_model(state):
    return [model.invoke(state)]

graph = MessageGraph()
graph.add_node("agent", call_model)
graph.set_entry_point("agent")
graph.add_edge("agent", END)

app = graph.compile()
result = app.invoke([("user", "你好，解释一下量子纠缠")])

pip install langgraph langchain langchain-openai tavily-python python-dotenv

import os
from dotenv import load_dotenv
load_dotenv()

os.environ["OPENAI_API_KEY"] = "your-api-key"
os.environ["TAVILY_API_KEY"] = "your-tavily-key"

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode
import operator
from tavily import TavilyClient

tavily = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])

@tool
def search(query: str) -> str:
    """搜索互联网获取最新信息"""
    result = tavily.search(query=query, max_results=3)
    return "\n\n".join([
        f"来源: {r['title']}\n{r['content']}"
        for r in result['results']
    ])

tools = [search]
tool_node = ToolNode(tools)

model = ChatOpenAI(model="gpt-4o", temperature=0)
model = model.bind_tools(tools)

class ResearchState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]

def should_continue(state: ResearchState) -> str:
    """决定是否需要继续搜索"""
    messages = state["messages"]
    last_message = messages[-1]
    
    # 如果 LLM 调用了工具，就继续执行工具
    if last_message.tool_calls:
        return "tools"
    # 否则结束
    return END

def call_model(state: ResearchState):
    """调用 LLM"""
    messages = state["messages"]
    response = model.invoke(messages)
    return {"messages": [response]}

from langgraph.graph import StateGraph, END

# 构建图
workflow = StateGraph(ResearchState)

workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

workflow.set_entry_point("agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {
        "tools": "tools",
        END: END
    }
)

workflow.add_edge("tools", "agent")

app = workflow.compile()

# 运行
inputs = {
    "messages": [
        HumanMessage(content="研究 2026 年最新的 AI Agent 框架发展趋势")
    ]
}

result = app.invoke(inputs)
print(result["messages"][-1].content)

class MultiAgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str  # 下一个要执行的 Agent
    research_result: str
    code_output: str
    review_feedback: str
    iteration: int

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# Supervisor Agent
supervisor_prompt = ChatPromptTemplate.from_messages([
    ("system", """你是团队的协调者。你的任务是：
    1. 分析用户的需求
    2. 决定下一步由哪个 Agent 执行
    3. 整合各个 Agent 的输出，判断任务是否完成
    
    可用的 Agent:
    - Researcher: 需要获取外部信息时调用
    - Coder: 需要编写或执行代码时调用
    - Reviewer: 需要审查代码质量时调用
    - FINISH: 任务已经完成
    
    只返回下一个要执行的 Agent 名称。"""),
    MessagesPlaceholder(variable_name="messages"),
])

supervisor_chain = supervisor_prompt | ChatOpenAI(model="gpt-4o")

# Research Agent
researcher_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是研究专家。使用搜索工具获取最新、最准确的信息。"),
    MessagesPlaceholder(variable_name="messages"),
])
researcher_chain = researcher_prompt | ChatOpenAI(model="gpt-4o").bind_tools(tools)

# Code Agent
coder_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是资深开发工程师。编写清晰、可运行、经过测试的代码。"),
    MessagesPlaceholder(variable_name="messages"),
])
coder_chain = coder_prompt | ChatOpenAI(model="gpt-4o")

# Review Agent
reviewer_prompt = ChatPromptTemplate.from_messages([
    ("system", """你是代码审查专家。检查代码的：
    1. 正确性：是否有 bug？
    2. 安全性：是否有安全漏洞？
    3. 可读性：代码是否清晰？
    4. 最佳实践：是否遵循编码规范？
    
    提供具体的改进建议。"""),
    MessagesPlaceholder(variable_name="messages"),
])
reviewer_chain = reviewer_prompt | ChatOpenAI(model="gpt-4o")

def supervisor_node(state: MultiAgentState):
    result = supervisor_chain.invoke(state)
    return {
        "messages": [result],
        "next": result.content.strip(),
        "iteration": state.get("iteration", 0) + 1
    }

def researcher_node(state: MultiAgentState):
    result = researcher_chain.invoke(state)
    return {
        "messages": [result],
        "research_result": result.content,
        "next": "supervisor"
    }

def coder_node(state: MultiAgentState):
    result = coder_chain.invoke(state)
    return {
        "messages": [result],
        "code_output": result.content,
        "next": "reviewer"
    }

def reviewer_node(state: MultiAgentState):
    result = reviewer_chain.invoke(state)
    return {
        "messages": [result],
        "review_feedback": result.content,
        "next": "supervisor"
    }

# 内存 Checkpointer（开发测试用）
from langgraph.checkpoint.memory import MemorySaver
memory = MemorySaver()

# SQLite Checkpointer（生产环境常用）
from langgraph.checkpoint.sqlite import SqliteSaver
sqlite_saver = SqliteSaver.from_conn_string("checkpoints.sqlite")

# Postgres Checkpointer（企业级）
from langgraph.checkpoint.postgres import PostgresSaver
# 需要 psycopg2

app = workflow.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "conversation-123"}}

# 第一次对话
result = app.invoke(
    {"messages": [HumanMessage(content="帮我写一个快速排序")]},
    config=config
)

# 可以中断，然后继续对话
result = app.invoke(
    {"messages": [HumanMessage(content="解释一下时间复杂度")]},
    config=config
)

workflow = StateGraph(AgentState)
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", should_continue)
workflow.add_edge("tools", "agent")

# 关键：在 tools 节点前添加中断点
app = workflow.compile(
    checkpointer=memory,
    interrupt_before=["tools"]  # 在执行 tools 前暂停
)

config = {"configurable": {"thread_id": "human-in-loop-1"}}

# 第一次调用，执行到 tools 前暂停
result = app.invoke(inputs, config=config)

# 检查当前状态
current_state = app.get_state(config)
print("下一个要执行的节点:", current_state.next)
# 下一个要执行的节点: ('tools',)

# 人工检查工具调用是否合理
last_message = current_state.values["messages"][-1]
print("准备调用工具:", last_message.tool_calls)

# 如果确认没问题，继续执行
result = app.invoke(None, config=config)  # 传入 None 表示继续

# 修改工具调用参数
updated_tool_calls = modify_tool_calls(last_message.tool_calls)

# 更新状态
app.update_state(
    config,
    {"messages": [AIMessage(tool_calls=updated_tool_calls)]}
)

# 然后继续执行
result = app.invoke(None, config=config)

# 流式输出消息
for chunk in app.stream(inputs, config):
    for node_name, output in chunk.items():
        if "messages" in output:
            for msg in output["messages"]:
                print(f"[{node_name}]: {msg.content}")

# Token 级别的流式输出
async for event in app.astream_events(inputs, version="v1"):
    if event["event"] == "on_chat_model_stream":
        token = event["data"]["chunk"].content
        print(token, end="", flush=True)

from tenacity import retry, stop_after_attempt, wait_exponential

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
def call_model_with_retry(state):
    return model.invoke(state["messages"])

from langgraph.constants import Send

def parallel_tasks(state):
    """并发执行多个任务"""
    return [
        Send("research", {"query": state["topic"]}),
        Send("code_generation", {"spec": state["spec"]}),
        Send("review", {"draft": state["draft"]})
    ]

workflow.add_conditional_edges("start", parallel_tasks)

# 查看当前状态
state = app.get_state(config)
print("当前状态:", state.values)
print("下一个节点:", state.next)

# 查看完整历史
for history in app.get_state_history(config):
    print(f"步骤 {history.config['configurable']['checkpoint_ns']}:")
    print(f"状态: {history.values}")

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "your-api-key"

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
import operator
import os

# 状态定义
class CodeReviewState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    code: str
    analysis_result: str
    security_findings: str
    style_findings: str
    final_report: str
    review_completed: bool

# 初始化模型
model = ChatOpenAI(model="gpt-4o", temperature=0)
checkpointer = SqliteSaver.from_conn_string("code_review.db")

# 各个 Agent 的 Prompt
def create_analyzer_chain():
    from langchain_core.prompts import ChatPromptTemplate
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """你是资深代码分析专家。请分析以下代码：
        1. 识别代码的主要功能
        2. 找出潜在的逻辑问题
        3. 标记复杂度较高的区域
        4. 建议可以优化的地方
        
        请提供详细、具体的分析。"""),
        ("human", "待分析代码:\n\n{code}")
    ])
    
    return prompt | model

def create_security_chain():
    from langchain_core.prompts import ChatPromptTemplate
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """你是安全审计专家。检查代码中的安全漏洞：
        1. SQL 注入风险
        2. XSS 漏洞
        3. 认证和授权问题
        4. 敏感信息泄露
        5. 不安全的依赖
        
        每个发现都要说明风险等级和修复建议。"""),
        ("human", "待检查代码:\n\n{code}\n\n分析结果:\n{analysis}")
    ])
    
    return prompt | model

def create_style_chain():
    from langchain_core.prompts import ChatPromptTemplate
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """你是代码风格审查专家。检查：
        1. 是否符合 PEP 8 规范
        2. 命名是否清晰
        3. 是否有适当的注释和文档
        4. 代码结构是否合理
        5. 可测试性
        
        提供具体的改进建议。"""),
        ("human", "待检查代码:\n\n{code}\n\n分析结果:\n{analysis}")
    ])
    
    return prompt | model

def create_summarizer_chain():
    from langchain_core.prompts import ChatPromptTemplate
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", """你是技术团队负责人。请整合所有审查结果，生成最终报告。
        报告结构：
        1. 总体评价
        2. 关键发现（按优先级排序）
        3. 具体修复建议
        4. 评分（0-100分）
        
        报告应该专业、清晰、可操作。"""),
        ("human", """代码分析:
        {analysis}
        
        安全检查:
        {security}
        
        风格检查:
        {style}""")
    ])
    
    return prompt | model

# 节点函数
def analyzer_node(state: CodeReviewState):
    chain = create_analyzer_chain()
    result = chain.invoke({"code": state["code"]})
    return {
        "messages": [result],
        "analysis_result": result.content
    }

def security_node(state: CodeReviewState):
    chain = create_security_chain()
    result = chain.invoke({
        "code": state["code"],
        "analysis": state["analysis_result"]
    })
    return {
        "messages": [result],
        "security_findings": result.content
    }

def style_node(state: CodeReviewState):
    chain = create_style_chain()
    result = chain.invoke({
        "code": state["code"],
        "analysis": state["analysis_result"]
    })
    return {
        "messages": [result],
        "style_findings": result.content
    }

def summarizer_node(state: CodeReviewState):
    chain = create_summarizer_chain()
    result = chain.invoke({
        "analysis": state["analysis_result"],
        "security": state["security_findings"],
        "style": state["style_findings"]
    })
    return {
        "messages": [result],
        "final_report": result.content,
        "review_completed": True
    }

# 构建图
def build_code_review_graph():
    workflow = StateGraph(CodeReviewState)
    
    # 添加节点
    workflow.add_node("analyzer", analyzer_node)
    workflow.add_node("security", security_node)
    workflow.add_node("style", style_node)
    workflow.add_node("summarizer", summarizer_node)
    
    # 设置入口
    workflow.set_entry_point("analyzer")
    
    # 添加边
    workflow.add_edge("analyzer", "security")
    workflow.add_edge("analyzer", "style")
    workflow.add_edge(["security", "style"], "summarizer")
    workflow.add_edge("summarizer", END)
    
    return workflow.compile(checkpointer=checkpointer)

# 使用示例
if __name__ == "__main__":
    app = build_code_review_graph()
    
    code_to_review = """
    def process_user_input(user_id, input_data):
        # 直接拼接 SQL，危险！
        query = f"SELECT * FROM users WHERE id = {user_id}"
        result = db.execute(query)
        
        # 没有输入验证
        user_input = input_data['content']
        
        # 执行用户提供的代码
        eval(user_input)
        
        return result
    """
    
    config = {"configurable": {"thread_id": "review-001"}}
    result = app.invoke({
        "code": code_to_review,
        "review_completed": False
    }, config=config)
    
    print("=" * 60)
    print("代码审查报告")
    print("=" * 60)
    print(result["final_report"])