LangGraph构建Ai智能体-14-多智能体架构前言在AI Agent应用中，多智能体架构涉及多个智能体协同工作，

前言

在AI Agent应用中，多智能体架构涉及多个智能体协同工作，以实现复杂的目标。

智能体监督器（Agent Supervisor） 在管理这些智能体，通过协调任务、分配责任和监控结果，确保整个多智能体架构的一致性和质量。监督器既负责协调，也负责评估，作为中心节点，确保智能体之间协作顺畅，避免目标冲突或重复工作。

智能体监督器的特点如下:

任务分配与协调：监督器根据请求的类型和每个智能体的能力分配任务。可以避免重复工作，让智能体保持专注，并通过确保每个智能体的特定技能得到利用来优化工作流程。

监控与错误纠正：通过观察每个智能体的表现和输出，监督器可以检测到某个智能体是否需要帮助，提供纠正措施，或者在某个智能体的输出不准确或不完整时触发额外的动作。

通信管理：有效的多智能体架构需要高效的通信。监督器帮助管理消息交换，决定哪个智能体应该在多轮对话中回应，尤其是在智能体的回应会影响工作流程的情况下。

一个典型例子是研究助手智能体。这个架构由多个智能体组成，每个智能体都有特定的角色和能力，例如：

文献检索智能体：负责从数据库中检索相关文献。
数据分析智能体：负责对检索到的文献进行分析，提取关键信息。
报告生成智能体：负责根据分析结果生成研究报告。

在这个网络中，智能体监督器的作用如下：

任务分配：根据用户的研究需求，将任务分配给不同的智能体。例如，如果用户需要最新的研究成果，监督器会将任务分配给文献检索智能体；如果需要对数据进行深入分析，则分配给数据分析智能体。
监控与纠正：监督器会监控每个智能体的输出。如果某个智能体检索到的文献不相关，或者分析结果不准确，监督器会介入，提供纠正措施，或者重新分配任务。
通信协调：在多轮对话中，监督器会管理智能体之间的通信。例如，当用户询问某个研究领域的最新进展时，监督器会协调文献检索智能体和数据分析智能体的输出，确保最终生成的报告既全面又准确。

多智能体系统的一个显著特点是“对话编程”的整合。在这种模式下，智能体通过结构化的对话来协调和分配任务。AutoGen提供了“以对话为中心的计算”工具，允许多智能体通过消息交换动态地塑造工作流程，每个智能体都根据其专业角色发挥作用。

这种架构不仅提高了系统的灵活性和适应性，还通过智能体之间的协作，解决了复杂任务，同时避免了单个智能体可能遇到的局限性。

实现代码-研究助手

在研究助手中，多个智能体可以处理信息收集、分析和报告的不同方面。例如：

研究智能体（Research Agent）：进行文献综述并从科学数据库中收集数据。
数据分析智能体（Data Analysis Agent）：使用统计工具处理定量数据。
报告智能体（Reporting Agent）：将研究发现整理成结构化的文档。

监督器通过启动任务给研究智能体，将分析任务分配给数据分析智能体，并以报告智能体完成任务来协调这些智能体。如果需要额外的数据分析，例如，监督器会介入，指导数据分析智能体用特定参数重新处理数据。

具体流程示例

任务启动：
- 用户（研究人员）提出研究请求：“请帮我分析最近五年内关于人工智能在医疗领域的研究进展。”
- 智能体监督器接收请求，并将任务分配给研究智能体。
研究智能体执行任务：
- 研究智能体从PubMed、IEEE Xplore等科学数据库中检索相关文献。
- 研究智能体收集到100篇相关文献，并将结果反馈给智能体监督器：“已收集到100篇相关文献，下一步需要进行数据分析。”
数据分析智能体执行任务：
- 智能体监督器将数据分析任务分配给数据分析智能体。
- 数据分析智能体使用统计工具（如R或Python的Pandas库）对文献中的数据进行处理，提取关键指标（如研究数量、引用次数、主要研究方向等）。
- 数据分析智能体完成初步分析后，反馈结果给智能体监督器：“初步分析完成，发现过去五年内相关研究数量逐年增长，主要集中在机器学习和图像识别领域。”
报告智能体执行任务：
- 智能体监督器将报告生成任务分配给报告智能体。
- 报告智能体将研究发现整理成结构化的文档，包括研究背景、数据分析结果和结论。
- 报告智能体生成报告后，反馈给智能体监督器：“报告已生成，包含研究背景、数据分析结果和结论。”
监督器介入调整：
- 智能体监督器在审查报告时发现，数据分析部分对某些关键指标的解释不够详细。
- 智能体监督器介入，指示数据分析智能体重新处理数据，重点关注某些特定参数（如不同研究方法的效果对比）。
- 数据分析智能体重新分析后，将更新后的结果反馈给报告智能体。
- 报告智能体根据新的分析结果更新报告，并再次反馈给智能体监督器。
任务完成：
- 智能体监督器确认报告内容完整且准确，通知用户任务完成：“研究助理网络已完成您的请求，报告已生成并发送至您的邮箱。”

import os
from typing import List, Literal, TypedDict

from dotenv import load_dotenv
from langchain.schema import Document  # Import the Document class
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
# pip install pypdf
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
# pip install beautifulsoup4
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import BaseMessage
from langgraph.prebuilt import create_react_agent

from langchain_experimental.tools import PythonREPLTool
import functools
import operator
from typing import Sequence
from typing_extensions import TypedDict
from typing import Annotated

load_dotenv()
embeddings = DashScopeEmbeddings(
    dashscope_api_key=os.getenv("OPENAI_API_KEY"),
    model="text-embedding-v4",
)
model = ChatOpenAI(model="qwen-plus",
                   base_url=os.getenv("BASE_URL"),
                   api_key=os.getenv("OPENAI_API_KEY"),
                   temperature=0,
                   streaming=True)

tavily_tool = TavilySearchResults(max_results=5)
python_repl_tool = PythonREPLTool()

system_prompt = (
    """
     你是负责协调以下工作人员对话的监督者：Researcher,Coder
     根据以下用户请求，回复下一步该由哪位工作人员行动。每位工作人员将执行一项任务，并回复他们的结果和状态。
     任务完成后，回复“FINISH”。
        
    You must respond using JSON format with the following structure:
    {{"next": "next worker name"}}
    """
)


class routeResponse(BaseModel):
    next: Literal["Researcher", "Coder", "FINISH"]


prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="messages"),
        (
            "system",
            """根据上面的对话，下一步该由谁行动？还是应该结束对话？请选择以下选项之一：
            Researcher,Coder,FINISH
            """,
        ),
    ]
)


# 监管节点
def supervisor_agent(state):
    supervisor_chain = prompt | model.with_structured_output(routeResponse)
    return supervisor_chain.invoke(state)


class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str


research_agent = create_react_agent(model, tools=[tavily_tool])
code_agent = create_react_agent(model, tools=[python_repl_tool])


def research_node(state: AgentState):
    result = research_agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name="Researcher")]
    }


def code_node(state: AgentState):
    result = code_agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name="Coder")]
    }


workflow = StateGraph(AgentState)
workflow.add_node("Researcher", research_node)
workflow.add_node("Coder", code_node)
workflow.add_node("supervisor", supervisor_agent)

workflow.add_edge(START, "supervisor")
workflow.add_conditional_edges(
    "supervisor",
    lambda x: x["next"],
    {
        "Researcher": "Researcher",
        "Coder": "Coder",
        "FINISH": END
    }
)
workflow.add_edge("Researcher", "supervisor")
workflow.add_edge("Coder", "supervisor")

graph = workflow.compile()

inputs = {"messages": [HumanMessage(content="研究人工智能的最新进展并总结。")]}
for s in graph.stream(input=inputs):
    if "__end__" not in s:
        print(s)
        print("----")
#
inputs = {"messages": [HumanMessage(content="撰写一份关于pikas的简短text研究报告,保存到当前文件夹下")]}
for s in graph.stream(input=inputs, config={"recursion_limit": 100}):
    if "__end__" not in s:
        print(s)
        print("----")

执行结果

{'supervisor': {'next': 'Researcher'}}
----
{'Researcher': {'messages': [HumanMessage(content='人工智能（AI）在2023年取得了显著的进展，涵盖技术、应用、治理等多个方面。以下是对这些进展的总结：\n\n### 1. **生成式AI和大模型的爆发**\n   - 生成式AI（AIGC）成为2023年的焦点，以**GPT-4**、**PaLM-E**、**Stable Diffusion**等为代表的大型语言模型和多模态模型，展示了强大的内容生成能力。这些模型能够生成高质量文本、图像、代码，甚至推动机器人控制等跨领域应用。\n   - 大模型参数量持续增加，例如**PaLM**模型的规模是2019年GPT-2的360倍，成本也提高了160倍，推动了深度学习技术的突破。\n\n### 2. **可信AI和安全治理的加强**\n   - 随着AI的快速发展，其**安全性和可信性**成为关注焦点。各国政府和企业纷纷出台政策和技术手段来应对AI可能带来的风险：\n     - 中国发布《生成式人工智能服务管理办法》征求意见稿。\n     - 欧盟通过《人工智能法案》，要求生成式AI披露内容来源并保护版权。\n     - 英伟达、谷歌、OpenAI等公司推出工具来限制AI的不良输出，如**NeMo Guardrails**和**ChatGPT安全方法**。\n   - 与此同时，研究人员开始探索AI的伦理问题，如生成内容的偏见、错误信息传播等。\n\n### 3. **算力和硬件的突破**\n   - 算力需求呈指数级增长，大模型训练需要更强大的计算支持。GPU、类脑芯片、量子计算等新技术不断演进，为AI提供更高效的算力基础。\n   - **硬件创新**包括：\n     - 新架构芯片的开发，如**存内计算**和**类脑芯片**。\n     - 企业级生成式AI解决方案，如英特尔和BCG的合作。\n\n### 4. **AI在多个领域的应用深化**\n   - **具身智能**：AI与机器人技术结合，实现更复杂的任务控制。例如：\n     - 微软提出ChatGPT应用于机器人的设计原则。\n     - 谷歌的**PaLM-E**模型可直接执行机器人指令。\n     - 李飞飞团队的**VoxPoser**项目，通过自然语言指令控制机器人。\n   - **脑机接口**：日本研究团队成功实现通过非侵入式脑机接口控制邮件发送，展示了AI与神经科学的结合潜力。\n   - **科学研究**：AI在多个领域加速科学发现，如氢聚变控制、抗体生成等。\n\n### 5. **AI与产业融合**\n   - 各国将AI作为核心竞争力的重要领域，推动**数实融合**，促进经济发展：\n     - 全球AI市场收入在2022年达到4497亿美元，预计未来几年将继续增长。\n     - 金融、医疗、制造、交通等行业广泛应用AI技术，例如：\n       - 医疗领域：提高诊断准确性、个性化医疗。\n       - 金融领域：风险评估、欺诈检测。\n       - 零售领域：需求预测、个性化营销。\n\n### 6. **AI的未来展望**\n   - **通用人工智能（AGI）**的研究正在加速。深度学习奠基人杰弗里·辛顿预测，AGI可能在5到20年内实现，可能带来新一轮生产力变革。\n   - **AI与人类协作**：研究表明，AI不仅能替代部分工作，还能创造新的岗位。预计到2030年，AI将净增5000万个就业岗位。\n   - **AI伦理与监管**：随着AI技术的普及，全球对AI的治理和监管需求日益迫切，未来需要更完善的法律框架和国际合作。\n\n### 总结\n2023年，人工智能在技术、应用和治理层面都取得了显著进展。生成式AI的爆发推动了AI的广泛应用，同时也引发了对安全性和伦理问题的讨论。随着算力提升和硬件创新，AI的能力不断增强，未来有望在更多领域实现突破。然而，如何平衡AI的发展与治理，将是未来需要重点关注的方向。', additional_kwargs={}, response_metadata={}, name='Researcher')]}}
----
{'supervisor': {'next': 'FINISH'}}

另一个执行的结果

{'supervisor': {'next': 'Researcher'}}
----
{'Researcher': {'messages': [HumanMessage(content='# 关于Pikas的研究报告\n\n## 引言\nPikas（鼠兔）是一种小型草食性哺乳动物，属于兔科。它们主要分布于北美西部的山区，栖息在岩石堆（talus）中。由于其对环境变化的高度敏感性，Pikas 成为了研究气候变化影响的重要物种。\n\n## 栖息地\nPikas 偏好岩石大小在8英寸到3英尺之间的岩石堆，较大的或较少孤立的岩石场通常更有可能支持Pikas的生存。它们利用岩石堆作为避难所，以躲避高温和捕食者。然而，研究表明，尽管岩石堆可以提供一定的温度调节，但在极端高温下，Pikas 仍然需要在地表进行觅食和繁殖，这使得它们面临更大的生存压力。\n\n## 气候变化的影响\nPikas 对温度变化极为敏感。研究发现，夏季的高温会抑制它们的觅食行为和生存能力，尤其是幼年Pikas的存活和扩散能力也会受到影响。在加利福尼亚州和内华达州，Pikas 的数量在较温暖的地区已经出现下降。尽管一些科学家认为地下温度避难所可以保护Pikas免受全球变暖的影响，但越来越多的证据表明，气候变化将对Pikas产生不利影响。\n\n### 雪覆盖的重要性\n最近的研究还表明，雪覆盖可能是影响Pikas分布的关键因素之一。较长的雪覆盖期可以为Pikas提供更好的绝缘层和水源，从而提高它们的生存率。一项研究发现，雪覆盖的持续时间比夏季平均温度更能预测Pikas的分布情况。\n\n## 保护现状\nPikas的保护现状引起了广泛关注。美国国家公园管理局（NPS）在2010年代启动了“Pikas in Peril”项目，旨在研究气候变化对Pikas的影响。研究结果显示，在科罗拉多州的落基山国家公园，如果气候变化趋势持续下去，Pikas可能在2100年前后完全消失。\n\n## 结论\nPikas 是研究气候变化对生态系统影响的重要指标物种。它们的生存依赖于适宜的栖息地和稳定的气候条件。随着全球气温的上升，Pikas面临着越来越大的威胁。保护Pikas不仅有助于维持生态平衡，还可以为其他物种的保护提供借鉴。\n\n## 参考文献\n1. California Department of Fish and Wildlife. (2021). Habitat Study Shows How American Pikas Fare in Warming Temperatures.\n2. Adventure Scientists. (2021). What Can Pikas Show Us About the Planet?\n3. Wyoming Untrapped. (2024). New Pika Research Suggests Snowpack Could Be Top Predictor for Occurrence.\n4. Smith, A. T. (2020). Conservation status of American pikas (Ochotona princeps).\n5. National Park Service. (2010s). Pikas in Peril.\n\n---\n\n这份报告总结了Pikas的基本生态特征、栖息地需求、气候变化对其的影响以及当前的保护措施。希望这份报告能够帮助读者更好地了解Pikas及其面临的挑战。', additional_kwargs={}, response_metadata={}, name='Researcher')]}}
----
{'supervisor': {'next': 'Coder'}}
----
Python REPL can execute arbitrary code. Use with caution.
{'Coder': {'messages': [HumanMessage(content='报告已成功保存为名为 `pikas_report.txt` 的文本文件。如果您需要进一步操作或查看文件内容，请告诉我！', additional_kwargs={}, response_metadata={}, name='Coder')]}}
----
{'supervisor': {'next': 'Coder'}}
----
{'Coder': {'messages': [HumanMessage(content='报告已成功保存为名为 `pikas_report.txt` 的文本文件。如果您需要进一步操作或查看文件内容，请告诉我！', additional_kwargs={}, response_metadata={}, name='Coder')]}}
----
{'supervisor': {'next': 'FINISH'}}
----

并且把文件成功保存了下来

实现代码2-自动化客服

在自动化客服中，可以部署一个多智能体系统，每个智能体专注于不同的客户需求，例如：

查询智能体（Query Agent）：处理客户的初始咨询。
解决方案智能体（Resolution Agent）：解决技术或物流问题。
升级智能体（Escalation Agent）：处理复杂查询或将案件升级到人工代表。

监督器根据客户查询的复杂性分配任务，将简单的咨询分配给查询智能体，将复杂问题分配给解决方案智能体。如果需要人工干预，监督器会将任务委托给升级智能体，确保流程顺畅，高效处理多样化的客户需求。

具体流程示例

客户发起咨询：
- 客户通过在线客服平台提交问题：“我的订单在哪里？”
- 智能体监督器接收问题并评估其复杂性。
任务分配：
- 智能体监督器判断这是一个简单的查询，将其分配给查询智能体。
- 查询智能体回复：“您的订单正在运输中，预计明天到达。”
复杂问题处理：
- 客户提交更复杂的问题：“我的订单显示已发货，但我没有收到任何物流信息。”
- 智能体监督器将问题分配给解决方案智能体。
- 解决方案智能体通过系统查询物流信息并回复：“物流信息已更新，您的包裹正在路上。”
升级到人工：
- 客户提出一个非常复杂的问题：“我在下单时选择了错误的地址，现在订单已经发货了，我该怎么办？”
- 智能体监督器判断需要人工干预，将问题分配给升级智能体。
- 升级智能体联系人工客服代表，并将问题详细信息传递过去。
- 人工客服代表与客户沟通并解决问题。

import os
from typing import List, Literal, TypedDict

from dotenv import load_dotenv
from langchain.schema import Document  # Import the Document class
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
# pip install pypdf
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
# pip install beautifulsoup4
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import BaseMessage
from langgraph.prebuilt import create_react_agent

from langchain_experimental.tools import PythonREPLTool
import functools
import operator
from typing import Sequence
from typing_extensions import TypedDict
from typing import Annotated
from functools import partial
from langchain_core.tools import tool

load_dotenv()
embeddings = DashScopeEmbeddings(
    dashscope_api_key=os.getenv("OPENAI_API_KEY"),
    model="text-embedding-v4",
)
model = ChatOpenAI(model="qwen-plus",
                   base_url=os.getenv("BASE_URL"),
                   api_key=os.getenv("OPENAI_API_KEY"),
                   temperature=0,
                   streaming=True)


# 客服的路由
class RouteResponseCS(BaseModel):
    next: Literal["Query_Agent", "Resolution_Agent", "Escalation_Agent", "FINISH"]


members_cs = ["Query_Agent", "Resolution_Agent", "Escalation_Agent"]

prompt_cs = ChatPromptTemplate.from_messages([
    ("system",
     """You are a Customer Service Supervisor managing agents: Query_Agent, Resolution_Agent, Escalation_Agent.
        Query_Agent：处理客户的初始查询
        Resolution_Agent：解决技术或后勤问题
        Escalation_Agent：处理复杂的查询或将案件上报给人工代表
     """),
    MessagesPlaceholder(variable_name="messages"),
    ("system", """Choose the next agent to act from Query_Agent, Resolution_Agent, Escalation_Agent, FINISH
    
    You must respond using JSON format with the following structure:
    {{"next": "next agent"}}
    """),
])


# 监管节点
def supervisor_agent_cs(state):
    supervisor_chain_cs = prompt_cs | model.with_structured_output(RouteResponseCS)
    return supervisor_chain_cs.invoke(state)


def agent_node(state, agent, name):
    result = agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name=name)]
    }


@tool
def escalation(question: str):
    """人工介入服务的工具"""

    print("人工处理节点", question)
    return "您的问题我们已经收到，请你耐心等待回复"


# 查询节点
query_agent = create_react_agent(model, tools=[TavilySearchResults(max_results=5)])
# 解决方案节点
resolution_agent = create_react_agent(model, tools=[PythonREPLTool()])
# 人工节点
escalation_agent = create_react_agent(model, tools=[escalation])


def query_node(state):
    result = query_agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name="Query_Agent")]
    }


def resolution_node(state):
    result = resolution_agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name="Resolution_Agent")]
    }


def escalation_node(state):
    result = escalation_agent.invoke(state)
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name="Escalation_Agent")]
    }


# Define Customer Service graph state and workflow
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str


# Initialize StateGraph and add nodes
workflow_cs = StateGraph(AgentState)
workflow_cs.add_node("Query_Agent", query_node)
workflow_cs.add_node("Resolution_Agent", resolution_node)
workflow_cs.add_node("Escalation_Agent", escalation_node)
workflow_cs.add_node("supervisor", supervisor_agent_cs)

workflow_cs.add_edge(START, "supervisor")
# Define conditional map for routing
workflow_cs.add_conditional_edges(
    "supervisor",
    lambda x: x["next"],
    {
        "Query_Agent": "Query_Agent",
        "Resolution_Agent": "Resolution_Agent",
        "Escalation_Agent": "Escalation_Agent",
        "FINISH": END
    }
)
workflow_cs.add_edge("Query_Agent", "supervisor")
workflow_cs.add_edge("Resolution_Agent", "supervisor")
workflow_cs.add_edge("Escalation_Agent", "supervisor")

# Compile and test the graph
graph_cs = workflow_cs.compile()

#
# Example input for testing
inputs_cs = {"messages": [HumanMessage(content="帮我重置下密码")]}
# inputs_cs = {"messages": [HumanMessage(content="手机怎么开机")]}
# inputs_cs = {"messages": [HumanMessage(content="手机怎么开不了机，我要退货")]}

# Run the graph
for output in graph_cs.stream(inputs_cs):
    if "__end__" not in output:
        print(output)

股票分析的例子

在金融咨询系统中，多个智能体协作收集和分析市场数据、进行财务分析，并根据用户查询检索新闻更新。以下是每个智能体的具体职责：

市场数据智能体（Market Data Agent）：使用yfinance库检索实时股票价格和相关财务指标。该智能体为客户提供指定股票或资产的当前市场状况的直接洞察。
财务分析智能体（Financial Analysis Agent）：处理财务数据，计算投资回报率（ROI）和增长率等指标，为客户提供关于潜在或现有投资的评估性见解。
财经新闻智能体（Financial News Agent）：检索相关的财经新闻文章和更新，提供市场趋势、公司公告和经济发展的总结性见解。

监督智能体（Supervisor Agent） 作为协调者，根据客户的查询内容，将请求首先分配给市场数据智能体（如果请求与股票价格或基本市场信息相关）。如果需要更深入的分析，例如计算财务表现或查找相关市场新闻，监督智能体会将查询路由到分析智能体或新闻智能体。这种协作方式提供了一个高效、连贯的工作流程，每个智能体都处理金融查询的不同方面，最终生成一个全面且结构化的回答，能够满足不同客户的需求。

import os
from typing import List, Literal, TypedDict

from dotenv import load_dotenv
from langchain.schema import Document  # Import the Document class
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
# pip install pypdf
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
# pip install beautifulsoup4
from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import BaseMessage
from langgraph.prebuilt import create_react_agent

from langchain_experimental.tools import PythonREPLTool
import functools
import operator
from typing import Sequence
from typing_extensions import TypedDict
from typing import Annotated
from functools import partial
from langchain_core.tools import tool
# import yfinance as yf
import baostock as bs

load_dotenv()
embeddings = DashScopeEmbeddings(
    dashscope_api_key=os.getenv("OPENAI_API_KEY"),
    model="text-embedding-v4",
)
model = ChatOpenAI(model="qwen-plus",
                   base_url=os.getenv("BASE_URL"),
                   api_key=os.getenv("OPENAI_API_KEY"),
                   temperature=0,
                   streaming=True)


# Route Response structure for supervisor's decision
class RouteResponseFin(BaseModel):
    next: Literal["Market_Data_Agent", "Analysis_Agent", "News_Agent", "FINISH"]


# Define agent members
members_fin = ["Market_Data_Agent", "Analysis_Agent", "News_Agent"]

# Supervisor prompt setup
system_prompt_fin = (
    "You are a Financial Services Supervisor managing the following agents: "
    """{agents}. Select the next agent to handle the current query. 
    
    You must respond using JSON format with the following structure:
    {{"next": "next agent"}}
    """
)

prompt_fin = ChatPromptTemplate.from_messages([
    ("system", system_prompt_fin),
    MessagesPlaceholder(variable_name="messages"),
    ("system", "Choose the next agent from: {options}."),
]).partial(agents=", ".join(members_fin), options=str(members_fin))


# Supervisor Agent
def supervisor_agent_fin(state):
    supervisor_chain_fin = prompt_fin | model.with_structured_output(RouteResponseFin)
    return supervisor_chain_fin.invoke(state)


def fetch_stock_price(query):
    """获取公司的股票价格"""
    resp = model.invoke(input=f"""
    查询{query}的股票代码,仅仅返回股票代码，不附加任何额外信息
    
    上证代码返回格式为: sh.600406
    深证代码返回格式为: sz.399001
    """)
    print("resp", resp)
    code = resp.content
    try:
        lg = bs.login()
        rs = bs.query_history_k_data_plus(code, "close",
                                          start_date='2025-08-14', end_date='2025-08-16',
                                          frequency="d", adjustflag="3")
        current_price = rs.data[0][0]
        bs.logout()
        return f"{query}的股票价格是RMB{current_price}."
    except Exception as e:
        return f"Error retrieving stock data for {query}: {str(e)}"


def agent_node(state, agent, name):
    result = agent.invoke(state)
    print(f"{name} Output: {result['messages'][-1].content}")
    return {
        "messages": [HumanMessage(content=result["messages"][-1].content, name=name)]
    }


market_data_prompt = (
    """您是市场数据代理。您的角色是检索最新的股票价格或基于用户查询的市场信息。确保您的回复包括当前价格以及任何相关的市场细节（如有）。"""
)
market_data_agent = create_react_agent(model, tools=[fetch_stock_price], state_modifier=market_data_prompt)
market_data_node = functools.partial(agent_node, agent=market_data_agent, name="Market_Data_Agent")


# 2. Financial Analysis Tool and Agent Prompt
def perform_financial_analysis(query):
    """Perform financial analysis based on user query."""
    if "ROI" in query:
        initial_investment = 1000
        final_value = 1200
        roi = ((final_value - initial_investment) / initial_investment) * 100
        return f"For an initial investment of ${initial_investment} yielding ${final_value}, the ROI is {roi}%."
    return "No relevant financial analysis found."


analysis_prompt = (
    "You are the Financial Analysis Agent. Analyze the financial data provided in the query. "
    "Perform calculations like ROI, growth rates, or other financial metrics as required. "
    "Provide a clear and concise response."
    "Only use the following tools:"
    "perform_financial_analysis"

)

analysis_agent = create_react_agent(model, tools=[perform_financial_analysis], state_modifier=analysis_prompt)
analysis_node = functools.partial(agent_node, agent=analysis_agent, name="Analysis_Agent")

# 3. Financial News Tool and Agent Prompt
financial_news_tool = TavilySearchResults(max_results=5)
news_prompt = (
    "You are the Financial News Agent. Retrieve the latest financial news articles relevant to the user's query. "
    "Use search tools to gather up-to-date news information and summarize key points."
    "Do not quote sources, just give a summary."
)
financial_news_agent = create_react_agent(model, tools=[financial_news_tool], state_modifier=news_prompt)
financial_news_node = functools.partial(agent_node, agent=financial_news_agent, name="News_Agent")


# Define Workflow State
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next: str


workflow_fin = StateGraph(AgentState)
workflow_fin.add_node("Market_Data_Agent", market_data_node)
workflow_fin.add_node("Analysis_Agent", analysis_node)
workflow_fin.add_node("News_Agent", financial_news_node)
workflow_fin.add_node("supervisor", supervisor_agent_fin)

# Conditional map for routing based on supervisor's decision
conditional_map_fin = {
    "Market_Data_Agent": "Market_Data_Agent",
    "Analysis_Agent": "Analysis_Agent",
    "News_Agent": "News_Agent",
    "FINISH": END  # This will end the workflow when supervisor decides
}
workflow_fin.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map_fin)
workflow_fin.add_edge(START, "supervisor")

# Compile the workflow
graph_fin = workflow_fin.compile()

# Testing the workflow with an example input
inputs_fin = {"messages": [HumanMessage(content="分析下国电南瑞的股票")]}
for output in graph_fin.stream(inputs_fin):
    if "__end__" not in output:
        print(output)