AI Agents-8 | 构建Muti-Agent系统 原创

from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
from IPython.display import display, Image

# 定义一个简单的工具
@tool
def get_weather(location: str) -> str:
    """获取指定地点的天气信息"""
    returnf"{location} 的天气是晴天"

# 创建智能体
agent = create_react_agent(
    model="gpt-4o-mini",
    tools=[get_weather],
    name="weather_agent",
    prompt=(
        "You are a weather agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with weather-related tasks\n"
        "- After you're done with your tasks, respond directly to the user\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    )
)

pip install langgraph-supervisor langchain-openai

import os
os.environ["OPENAI_API_KEY"] = "<your_api_key>"

from langchain_openai import ChatOpenAI
from langgraph.graph import START, END
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor

# 选择 LLM
model = ChatOpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))

# 定义工具
def add(a: float, b: float) -> float:
    """加法"""
    return a + b

def multiply(a: float, b: float) -> float:
    """乘法"""
    return a * b

def divide(a: float, b: float) -> float:
    """除法"""
    return a / b

# 创建研究智能体
research_agent = create_react_agent(
    model=model,
    tools=[TavilySearchResults(max_results=3, tavily_api_key=os.getenv("TAVILY_API_KEY"))],
    name="research_agent",
    prompt=(
        "You are a research agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with research-related tasks, DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    )
)

# 创建数学智能体
math_agent = create_react_agent(
    model=model,
    tools=[add, multiply, divide],
    name="math_agent",
    prompt=(
        "You are a math agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with math-related tasks\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    )
)

# 创建监督者智能体
supervisor_agent = create_supervisor(
    model=model,
    agents=[research_agent, math_agent],
    prompt=(
        "You are a supervisor managing two agents:\n"
        "- a research agent. Assign research-related tasks to this agent\n"
        "- a math agent. Assign math-related tasks to this agent\n"
        "Assign work to one agent at a time, do not call agents in parallel.\n"
        "Do not do any work yourself."
    ),
    add_handoff_back_messages=True,
    output_mode="full_history"
).compile()

def test_supervisor_functionality():
    """测试监督者模式的功能"""
    print("Query: find US and New York state GDP in 2024. what % of US GDP was New York state?")
    print("-" * 80)
    try:
        for chunk in supervisor_agent.stream(
            {
                "messages": [
                    {
                        "role": "user",
                        "content": "find US and New York state GDP in 2024. what % of US GDP was New York state?",
                    }
                ]
            },
            subgraphs=False
        ):
            print(chunk)
        print("Test completed successfully")
    except Exception as e:
        print(f"Test failed with error: {str(e)}")
    print("=" * 80)

if __name__ == "__main__":
    test_supervisor_functionality()

supervisor = create_supervisor(
    agents=[agent1, agent2],
    output_mode="last_message"
)

research_team = create_supervisor(
    [research_agent, math_agent],
    model=model,
    supervisor_name="research_supervisor"
).compile(name="research_team")

writing_team = create_supervisor(
    [writing_agent, publishing_agent],
    model=model,
    supervisor_name="writing_supervisor"
).compile(name="writing_team")

top_level_supervisor = create_supervisor(
    [research_team, writing_team],
    model=model,
    supervisor_name="top_level_supervisor"
).compile(name="top_level_supervisor")

from langgraph_supervisor import create_handoff_tool

# 创建一个自定义的交接工具
assign_to_research_agent = create_handoff_tool(
    agent_name="research_agent",
    descriptinotallow="Assign task to a researcher agent.",
)

# 将自定义工具传递给监督者
supervisor_agent = create_react_agent(
    model="openai:gpt-4o-mini",
    tools=[assign_to_research_agent],
    prompt=(
        "You are a supervisor managing two agents:\n"
        "- a research agent. Assign research-related tasks to this agent\n"
        "- a math agent. Assign math-related tasks to this agent\n"
        "Assign work to one agent at a time, do not call agents in parallel.\n"
        "Do not do any work yourself."
    ),
    name="supervisor",
)

from langgraph_swarm import create_swarm, create_handoff_tool

# 创建蜂群模式的智能体
swarm_agent = create_swarm(
    agents=[research_agent, math_agent],
    default_active_agent="math_agent",
).compile()

# 测试蜂群模式的功能
def test_swarm_functionality():
    print("Query: find US and New York state GDP in 2024. what % of US GDP was New York state?")
    print("-" * 80)
    try:
        for chunk in swarm_agent.stream(
            {
                "messages": [
                    {
                        "role": "user",
                        "content": "find US and New York state GDP in 2024. what % of US GDP was New York state?",
                    }
                ]
            },
            subgraphs=False
        ):
            print(chunk)
        print("Test completed successfully")
    except Exception as e:
        print(f"Test failed with error: {str(e)}")
    print("=" * 80)

if __name__ == "__main__":
    test_swarm_functionality()

response = agent.invoke({"messages": [{"role": "user", "content": "what is the weather in sf"}]})

response = await agent.ainvoke({"messages": [{"role": "user", "content": "what is the weather in sf"}]})