import os

DASHSCOPE_API_KEY = 'xxx'
os.environ["DASHSCOPE_API_KEY"] = DASHSCOPE_API_KEY

LANGCHAIN_API_KEY = 'xxx'
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY

TAVILY_API_KEY = 'xxx'
os.environ["TAVILY_API_KEY"] = TAVILY_API_KEY

from langchain_community.chat_models import ChatTongyi
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, MessagesState, START, END

model = ChatTongyi(model_name="qwen-max", temperature=0.1)

# 定义两个智能体——研究助手和摘要专家
def research_agent(state: MessagesState):
    """Call research agent"""
    messages = [SystemMessage(content="You are a research assistant. Given a topic, provide key facts and information.")] + state["messages"]
    response = model.invoke(messages)
    return {"messages": [response]}

def summarize_agent(state: MessagesState):
    """Call summarization agent"""
    messages = [SystemMessage(content="You are a summarization expert. Condense the given information into a brief summary.")] + state["messages"]
    response = model.invoke(messages)
    return {"messages": [response]}

workflow = StateGraph(MessagesState)
workflow.add_node("research", research_agent)
workflow.add_node("summarize", summarize_agent)

# 指定显式控制流，先由研究助手提供信息，再由摘要专家生成摘要
workflow.add_edge(START, "research")
workflow.add_edge("research", "summarize")
workflow.add_edge("summarize", END)

graph = workflow.compile()

for s in graph.stream(
    {"messages": [HumanMessage(content="强化学习")]},
    {"recursion_limit": 100},
):
    if "__end__" not in s:
        print(s)
        print("----")

import os

DASHSCOPE_API_KEY = 'xxx'
os.environ["DASHSCOPE_API_KEY"] = DASHSCOPE_API_KEY

LANGCHAIN_API_KEY = 'xxx'
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = LANGCHAIN_API_KEY

TAVILY_API_KEY = 'xxx'
os.environ["TAVILY_API_KEY"] = TAVILY_API_KEY

from langchain_community.chat_models import ChatTongyi
from typing import Annotated
from langchain_core.messages import SystemMessage, ToolMessage
from langgraph.prebuilt import ToolNode, InjectedState, create_react_agent

model = ChatTongyi(model_name="qwen-max", temperature=0.1)

# 定义两个智能体——研究助手和摘要专家
def research_agent(state: Annotated[dict, InjectedState]):
    """Call research agent"""
    messages = [SystemMessage(content="You are a research assistant. Given a topic, provide key facts and information.")] + state["messages"][:-1]
    response = model.invoke(messages)
    tool_call = state["messages"][-1].tool_calls[0]
    return {"messages": [ToolMessage(response.content, tool_call_id=tool_call["id"])]}

def summarize_agent(state: Annotated[dict, InjectedState]):
    """Call summarization agent"""
    messages = [SystemMessage(content="You are a summarization expert. Condense the given information into a brief summary.")] + state["messages"][:-1]
    response = model.invoke(messages)
    tool_call = state["messages"][-1].tool_calls[0]
    return {"messages": [ToolMessage(response.content, tool_call_id=tool_call["id"])]}

# 将研究助手和摘要专家作为工具
tool_node = ToolNode([research_agent, summarize_agent])
# 定义一个负责总体协调的智能体，其将研究助手和摘要专家作为工具，以ReAct的方式进行思考，分步调用工具
graph = create_react_agent(model, [research_agent, summarize_agent], state_modifier="First, research information on a given topic, Second, summarize based on information above.")

inputs = {"messages": [("user", "强化学习")]}

for s in graph.stream(inputs, stream_mode="values"):
    message = s["messages"][-1]
    if isinstance(message, tuple):
        print(message)
    else:
        message.pretty_print()