LangChain 1.0 搭建 Agent 智能体应用实战

from langchain.agents import create_agent

agent = create_agent(
    model=model,                    # 模型
    tools=[order_query_tool],       # 工具
    system_prompt="你是一个订单查询助手，能够查询订单状态和明细。" , # 系统提示
    middlewares=[order_query_middleware],                    # 中间件
    checkpointer=checkpointer,      # 检查点短期记忆
    store=store,                    # 状态存储长期记忆
    state_schema=OrderQueryState,   # 扩展状态（如需要）
    context_schema=AgentContext,    # 上下文状态（如需要）
    response_format=ResponseModel   # 结构化输出（如需要）
)

# ============ 限制最大 3 次循环 ============
config = {
    "configurable": {"thread_id": "limit_demo"}, # 限制thread_id 线程ID
    "recursion_limit": 3  # 最多 3 次迭代，或使用中间件进行精确跟踪和终止循环
}

result = agent.invoke(
        {"messages": [{"role": "user", "content": "LangChain 1.0 发布日期"}]},
        config=config
    )

# 环境依赖版本
!pip list | grep langchain

#python 版本
!python --version

# 加载环境
from dotenv import load_dotenv

# 加载 .env 环境变量
load_dotenv(override=True)

# 1. 定义带速率限制的load_chat_model函数
from langchain.chat_models import init_chat_model
from langchain_core.rate_limiters import InMemoryRateLimiter

# 2. 配置速率限制器
rate_limiter = InMemoryRateLimiter(
    requests_per_second=5,       # 每秒最多5个请求
    check_every_n_seconds=1.0    # 每1秒检查一次是否超过速率限制
)

# 3. 对模型调用进行封装，后续直接调用传参数就行
def load_chat_model(
    model: str,
    provider: str,
    temperature: float = 0.7,
    max_tokens: int | None = None,
    base_url: str | None = None,
):
    return init_chat_model(
        model=model,               # 模型名称
        model_provider=provider,   # 模型供应商
        temperature=temperature,   # 温度参数，用于控制模型的随机性，值越小则随机性越小
        max_tokens=max_tokens,     # 最大生成token数
        base_url=base_url,         # 专用于自定义 API Server 或代理
        rate_limiter=rate_limiter  # 自动限速
    )

#!pip install langchain-tavily

# 1.导入相关库
from langchain.agents import create_agent
from langchain_community.tools.tavily_search import TavilySearchResults # 导入第三方社区集成 Tavily 搜索工具
from langchain_tavily import TavilySearch

# 2.导入模型和工具
web_search = TavilySearchResults(max_results=2)

# 3.创建模型
model = load_chat_model(model="deepseek-chat",provider="deepseek")

# 4.创建Agent
agent = create_agent(
    model=model,
    tools=[web_search],
    system_prompt="你是一名多才多艺的智能助手，可以调用工具帮助用户解决问题。"
)

# 5.运行Agent获得结果
result = agent.invoke(
    {"messages": [{"role": "user", "content": "请帮我查询2024年诺贝尔物理学奖得主是谁？"}]}
)

result['messages'][-1].content

#!pip install langchain-experimental

from langchain_core.tools import tool
from langchain.agents import create_agent

# 1. 定义一个简单的 Tool (Runnable)
@tool
def multiply(a: int, b: int) -> int:
    """Multiplies a and b."""
    return a * b

# 2.导入模型
model = load_chat_model(
    model="gpt-4o-mini",    # 指定OpenAI的gpt-4o-mini模型
    provider="openai",      # 指定模型提供商为openai
)

# 3.创建Agent
agent = create_agent(model=model,tools=[multiply])

# 4. 调用Agent
response = agent.invoke({
    "messages": [{
        "role": "user",
        "content": "帮我计算12乘以6等于多少？"
    }]
})

response["messages"]

response["messages"][-1].content

from pydantic import BaseModel, Field
from langchain_core.tools import StructuredTool

"""
1. 通过 Pydantic BaseModel 定义参数，提供：
- 参数描述（description）
- 必填/可选约束
- 更清晰的 Schema 文档
"""
class DivideInput(BaseModel):
    """除法工具输入参数"""
    dividend: float = Field(description="被除数")
    divisor: float = Field(description="除数，不能为零")

def divide(dividend: float, divisor: float) -> float:
    """执行除法运算，支持浮点数"""
    if divisor == 0:
        raise ValueError("除数不能为零")
    return dividend / divisor

# 2. 创建带参数校验的工具
division_tool = StructuredTool.from_function(
    func=divide,
    name="DivisionTool",
    description="安全执行除法运算，自动处理除零错误",
    args_schema=DivideInput,  # 显式指定参数模式
    return_direct=False,  # 是否直接返回工具结果（不经过 LLM 再次处理）
)

# 3. 测试参数校验（触发 Pydantic 验证）
try:
    division_tool.invoke({"a": 10, "b": 2})  # 错误：参数名不匹配
except Exception as e:
    print(f"参数校验失败：{e}")

# 4. 正确调用
result = division_tool.invoke({"dividend": 10, "divisor": 2})
print(f"除法结果：{result}")

import os
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel, Field
from langchain_core.tools import StructuredTool
from typing import Type

# 1. 定义包含业务逻辑的工具
class OrderQueryInput(BaseModel):
    """订单查询参数"""
    order_id: str = Field(description="订单编号，格式：ORD-2024-XXXX")
    include_details: bool = Field(default=False, description="是否包含商品明细")

class OrderQueryTool(StructuredTool):
    """订单查询工具"""
    name: str = "query_order"
    description: str = "查询电商平台订单状态和物流信息"
    args_schema: Type[BaseModel] = OrderQueryInput
    return_direct: bool = False

    def _run(self, order_id: str, include_details: bool = False) -> dict:
        # 模拟数据库查询
        order_db = {
            "ORD-2024-1234": {"status": "已发货", "express": "顺丰", "amount": 299},
            "ORD-2024-5678": {"status": "待付款", "express": "", "amount": 149},
        }

        # 处理查询逻辑
        if order_id not in order_db:
            return {"error": f"订单 {order_id} 不存在"}
        result = order_db[order_id]

        # 处理包含明细的情况
        if include_details:
            result["items"] = ["商品A × 2", "商品B × 1"]

        return result

# 2. 初始化模型
model = init_chat_model(
    model="openai:gpt-4o-mini",
    temperature=0,
    api_key=os.getenv("OPENAI_API_KEY")
)

# 3. 创建 ReAct Agent（自动处理工具调用）
agent = create_agent(
    model=model,
    tools=[OrderQueryTool()],  # 直接传入 StructuredTool 实例
    system_prompt="你是一个电商客服助手，使用工具查询订单信息，回答要友好且准确",
    checkpointer=InMemorySaver()
)

# 4. 执行并观察 ReAct 过程
async def run_agent():
    config = {"configurable": {"thread_id": "customer_001"},"recursion_limit": 15}# 最大 15 次迭代

    query = "请帮我查订单 ORD-2024-1234 的详细状态，包括商品明细"

    async for step in agent.astream(
        {"messages": [{"role": "user", "content": query}]},
        config=config,
        stream_mode="values"  # 流式输出模式，返回每一步的完整状态
    ):
        message = step["messages"][-1]
        message.pretty_print()
        print("-" * 50)

# 运行
await run_agent()

from langchain.agents import create_agent
from langchain_core.tools import tool

# 定义天气查询工具
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。"""
    weather_data = {
        "北京": "晴朗，气温25°C",
        "上海": "多云，气温28°C",
        "广州": "小雨，气温30°C"
    }
    return f"{city}的天气是：{weather_data.get(city, '未知')}"

# 定义数学计算工具
@tool
def calculate(expression: str) -> str:
    """计算一个数学表达式的结果。"""
    try:
        result = eval(expression)
        return f"计算结果是：{result}"
    except Exception as e:
        return f"计算出错：{str(e)}"

# 1. 初始化LLM
llm = load_chat_model(model="gpt-4o-mini",provider="openai")

# 2. 创建Agent
agent = create_agent(
    model=llm,
    tools=[get_weather, calculate],
    system_prompt="你是一个多功能的助手，可以查询天气和进行数学计算。"
)

# 3. 测试多工具调用
user_queries = [
    "北京和上海的天气怎么样？",
    "如果北京气温是25度，上海是28度，那么北京的温度比上海低多少度？"
]

# 4. 执行测试
for query in user_queries:
    print(f"用户: {query}")
    response = agent.invoke({
        "messages": [{"role": "user", "content": query}]
    })
    print(f"Agent: {response['messages'][-1].content}")
    print("-" * 50)

import getpass
import operator
from typing import Annotated, List, Union
import os

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langchain.agents import create_agent

# 引入 UI 库
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.markdown import Markdown

# 初始化控制台
console = Console()

# --- 第一步：定义工具 (和以前一样，这是 Core 标准) ---
# 定义天气查询工具
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。"""
    weather_data = {
        "北京": "晴朗，气温25°C",
        "上海": "多云，气温28°C",
        "广州": "小雨，气温30°C"
    }
    return f"{city}的天气是：{weather_data.get(city, '未知')}"

# 定义数学计算工具
@tool
def add(a: float, b: float) -> float:
    """计算两个数的和"""
    return a + b

tools = [get_weather, add]

# --- 第二步：初始化模型 (必须绑定工具) ---
model = ChatOpenAI(model="gpt-4o", temperature=0)

# --- 第三步：构建图 (使用 prebuilt 的 ReAct Agent) ---
# 在 LangChain 1.0+ 中，这是 AgentExecutor 的官方替代品
# 它自动构建了：State -> Model Node -> Tool Node -> Loop 逻辑
graph = create_agent(model, tools=tools)

# --- 第四步：编写“教学专用”的可视化流式运行器 ---
def run_demo_with_visualization(user_input: str):
    print("\n" + "="*50)
    console.print(f"[bold yellow]开始任务：[/bold yellow] {user_input}")

    messages = [HumanMessage(content=user_input)]

    # graph.stream 是 LangGraph 的核心
    # 它可以让我们看到状态流转的每一步 (step-by-step)
    step_count = 1

    # values 模式会返回当前的 message 列表状态
    for event in graph.stream({"messages": messages}, stream_mode="values"):
        # 获取最新的一条消息
        current_message = event["messages"][-1]

        # 1. 如果是人类的消息 (初始状态)
        if isinstance(current_message, HumanMessage):
            continue # 跳过，这是输入

        # 2. 如果是 AI 的消息 (思考与决策)
        if isinstance(current_message, AIMessage):
            # 检查是否有工具调用
            if current_message.tool_calls:
                # 提取工具调用的细节
                for tool_call in current_message.tool_calls:
                    console.print(Panel(
                        Text(f"🤔 AI 思考决定：需要调用外部工具\n"
                             f"🔧 工具名称: {tool_call['name']}\n"
                             f"📥 输入参数: {tool_call['args']}", style="bold cyan"),
                        title=f"Step {step_count}: 决策 (Decision)",
                        border_style="cyan"
                    ))
            else:
                # 如果没有工具调用，说明是最终回复
                console.print(Panel(
                    Markdown(current_message.content),
                    title=f"Step {step_count}: 最终回复 (Final Answer)",
                    border_style="green"
                ))
            step_count += 1

        # 3. 如果是工具的消息 (观察与结果)
        if isinstance(current_message, ToolMessage):
            console.print(Panel(
                Text(f"👀 工具返回结果 (Observation):\n{current_message.content}", style="italic white"),
                title=f"Step {step_count}: 执行与观察",
                border_style="magenta"
            ))
            step_count += 1

# --- 第五步：运行演示 ---
if __name__ == "__main__":
    # 这是一个多步任务：先算乘法，再查属性
    run_demo_with_visualization("查询一下北京和上海气温，并且计算一下北京的温度比上海低多少度？")

# 安装 MCP 适配器（关键依赖）\MCP 服务器开发库（如需自定义工具）
#!pip install langchain-mcp-adapters mcp

from langchain_mcp_adapters.client import MultiServerMCPClient   # 导入 MCP 客户端
import os
from langchain_core.tools import tool
from langchain.agents import create_agent

# 1. 初始化 MCP 客户端，只连接本地 MCP 服务器
    # 获取当前文件所在目录的绝对路径
mcp_server_path = os.path.join("mcp_server.py")
print(mcp_server_path)

# 2. 初始化 MCP 客户端，只连接本地 MCP 服务器
mcp_client = MultiServerMCPClient(
        {
            # 本地 Python MCP 服务器（stdio 传输）
            "math": {
                "transport": "stdio",
                "command": "python",
                "args": [mcp_server_path],  # 使用绝对路径
            },
            # 如果需要其他服务器，可以在这里添加
            # 注意：只添加确实在运行的服务器！否则会导致连接失败，需要先运行mcp_server.py文件！！！
        }
    )

# 3. 加载 MCP 工具
try:
    mcp_tools = await mcp_client.get_tools()
    print(f"✅ 成功加载 {len(mcp_tools)} 个 MCP 工具: {[t.name for t in mcp_tools]}")
except Exception as e:
    print(f"❌ 加载 MCP 工具失败: {e}")
    print("将只使用本地工具")
    mcp_tools = []

# 4. 定义天气查询工具
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。"""
    weather_data = {
        "北京": "晴朗，气温25°C",
        "上海": "多云，气温28°C",
        "广州": "小雨，气温30°C"
    }
    return f"{city}的天气是：{weather_data.get(city, '未知')}"

# 5. 合并所有工具
all_tools = [get_weather] + mcp_tools

# 6. 加载 ChatOpenAI 模型
llm = load_chat_model(model="gpt-4o-mini",provider="openai")

# 7. 创建Agent
agent = create_agent(
    model=llm,
    tools=all_tools,
    system_prompt="你是一个多功能的助手，可以查询天气和进行数学计算。"
)

# 8. 测试Agent，多个工具就可以使用ainvoke异步调用
response = await agent.ainvoke({
        "messages": [{"role": "user", "content": "查询一下北京和上海气温，并且计算一下北京的温度比上海低多少度？"}]
    })
print(f"Agent: {response['messages'][-1].content}")

from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.agents import create_agent

# 1. 正确的 MCP 配置格式（适用于 langchain_mcp_adapters）
# MultiServerMCPClient 需要的是扁平的字典结构，每个服务器是一个键值对
mcp_config = {
    # 本地 Python MCP 服务器
    "math": {
        "transport": "stdio",
        "command": "python",
        "args": ["mcp_server.py"]
    },
    # 高德地图 MCP 服务器
    "amap-maps": {
        "transport": "stdio",
        "command": "npx",
        "args": ["-y", "@amap/amap-maps-mcp-server"],
        "env": {
            "AMAP_MAPS_API_KEY": os.getenv("AMAP_MAPS_API_KEY"),
        }
    }
}

# 2. 创建 MCP 客户端
client = MultiServerMCPClient(mcp_config)
print("正在连接 MCP 服务器...")

# 3. client.get_tools() 会自动：
#   1. 调用所有服务器的 list_tools 接口
#   2. 将 MCP Tool Schema 转换为 LangChain StructuredTool
tools = await client.get_tools()
print(f"成功加载 {len(tools)} 个工具: {[t.name for t in tools]}")

# 4. 创建 Agent
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 直接将转换好的 tools 传给 create_agent
agent = create_agent(llm, tools,system_prompt="你是会调用工具进行天气查询、地图查询、网页部署的智能助手")

# 5. 运行 Agent
print("\n--- 开始测试 Agent ---")

# 6. 这里我们模拟一个请求（具体 prompt 取决于你的工具功能）
query = "请帮我搜索查询一下北京市今天的天气，并计算一下最大温差是多少度？"

inputs = {"messages": [HumanMessage(content=query)]}

async for chunk in agent.astream(inputs, stream_mode="values"):
    last_msg = chunk["messages"][-1]
    print(f"\n[{type(last_msg).__name__}]:")
    print(last_msg.content)

    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        print(f">>> 调用工具详情: {last_msg.tool_calls}")

MCP_CONFIG = {
    "weather": {
        "transport": "streamable_http",
        "url": "http://localhost:8000/mcp"  # HTTP 服务器地址
    }
}

# 定义意图分类系统提示
INTENT_SYSTEM_PROMPT = """
你是一个专业的意图分类器，请只返回以下类别之一：
- search
- pdf
- database
- math
- none

并严格只返回类别名，不要输出其它内容。
"""

# 创建一个意图识别模型
intent_llm = load_chat_model(
    model="gpt-4o-mini",    # 指定OpenAI的gpt-4o-mini模型
    provider="openai",      # 指定模型提供商为openai
)

@tool
def query_database(sql: str) -> str:
    """
        执行 SQL 查询，仅限内部业务数据库。
        参数：sql Sql语句。
        示例：如 select * from users limit 5
    """
    return f"模拟 SQL 执行：{sql}"

from langchain.tools import tool

@tool
def search_web(query: str) -> str:
    """Web 搜索工具，用于查询网络公开信息，不适用于内部数据.参数：query 用户查询，如 OpenAI发布会"""
    return f"模拟搜索结果：你搜索了 {query}"

@tool
def extract_pdf_text(path: str) -> str:
    """解析 PDF 文本文件。参数为文件的本地路径.参数：path 文件路径，如 /files/contract.pdf"""
    return f"模拟 PDF 内容：从 {path} 中解析出的内容"

@tool
def query_database(sql: str) -> str:
    """执行 SQL 查询，仅限内部业务数据库.参数：sql Sql语句，如 select * from users limit 5"""
    return f"模拟 SQL 执行：{sql}"

@tool
def calculate(expr: str) -> str:
    """计算数学表达式。适用于算式运算.参数：expr 数学表达式，如 (12+3)*(8-2)"""
    return str(eval(expr))

# 1.Tool 工具分组
TOOL_GROUPS = {
    "search": [search_web],
    "pdf": [extract_pdf_text],
    "database": [query_database],
    "math": [calculate],
}

# 2.创建一个意图识别模型
intent_llm = load_chat_model(
    model="gpt-4o-mini",    # 指定OpenAI的gpt-4o-mini模型
    provider="openai",      # 指定模型提供商为openai
)

# 3. 定义意图分类系统提示
INTENT_SYSTEM_PROMPT = """
你是一个专业的意图分类器，请只返回以下类别之一：
- search
- pdf
- database
- math
- none

并严格只返回类别名，不要输出其它内容。
"""

# 4. 定义意图分类函数
def classify_intent(user_query: str) -> str:
    result = intent_llm.invoke(
        [
            ("system", INTENT_SYSTEM_PROMPT),
            ("user", user_query)
        ]
    )
    return result.content.strip()

from langchain.agents import create_agent
import os
from dotenv import load_dotenv
load_dotenv()

# 5. 创建智能体函数
def create_agent_for_group(group: str):
    tools = TOOL_GROUPS.get(group, [])

    if not tools:
        return None

    model = load_chat_model(
        model="deepseek-chat",
        provider="deepseek",
    )

    agent = create_agent(
        model=model,tools=tools,system_prompt="你是一个 helpful assistant，可以使用工具回答问题。你必须严格根据工具描述选择工具！如果没有合适的工具，请回答“无合适工具”"
    )

    return agent

# 6. 路由智能体函数
def router_agent(user_query: str):
    # 1. 识别意图
    intent = classify_intent(user_query)
    print(f"[Router] 检测到意图: {intent}")

    # 2. 创建对应子 Agent
    sub_agent = create_agent_for_group(intent)

    if sub_agent is None:
        return "无法为该问题找到合适的工具或 Agent。"

    # 3. 调用子 Agent 执行任务
    result = sub_agent.invoke({
        "messages": [{"role": "user", "content": user_query}]
    })

    return result

res = router_agent("请帮我搜索一下今年Google最新的大模型版本的发布会")

# 7. 测试智能体
queries = [
        "请帮我搜索一下今年Google最新的大模型版本的发布会",
        "帮我解析一下这个PDF：/root/files/contract.pdf",
        "执行一个SQL：select * from products limit 5",
        "计算 (17+3)*(8-1)",
    ]

for q in queries:
    print("\n====== 用户问题 ======")
    print(q)
    print("====== Agent 回复 ======")
    print(router_agent(q)["messages"][2])

#system_prompt 在 ReAct 循环中的位置：

# System Prompt (固定前缀)
#    ↓
# 用户输入 → 模型推理 (Thought) → 工具调用 (Action) → 观察结果 (Observation)
#    ↓
# 循环直到满足终止条件 → 最终回答

from langchain.agents import create_agent
from langchain.tools import tool

# 1. 定义一个简单的天气查询工具
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。"""
    weather_data = {
        "北京": "晴朗，气温25°C",
        "上海": "多云，气温28°C",
        "广州": "小雨，气温30°C"
    }
    return f"{city}的天气是：{weather_data.get(city, '未知')}"

# 2. 静态 system_prompt（固定不变）
agent_static = create_agent(
    model="openai:gpt-4o-mini",
    tools=[get_weather],
    system_prompt=(
        "你是一个天气助手，回答不超过20字。\n"
        "调用工具时，严格按照以下格式：\n"
        "1. 使用 `get_weather(city: str)` 获取天气；\n"
        "2. 仅返回天气结果，不解释过程。"
    )
)

print("=== 静态 System Prompt ===")
response1 = agent_static.invoke({
    "messages": [{"role": "user", "content": "北京天气"}]
})
print(f"AI: {response1['messages'][-1].content}")

# 3. 动态提示词（通过中间件实现）
from langchain.agents.middleware import dynamic_prompt
from typing import TypedDict

# 4. 定义上下文结构
class Context(TypedDict):
    user_role: str  # 用户角色

# 5. 动态提示函数
@dynamic_prompt
def role_based_prompt(request):
    """根据用户角色生成不同提示词"""
    user_role = request.runtime.context.get("user_role", "user")

    if user_role == "expert":
        return "你是一个专业气象分析师，提供详细数据"
    elif user_role == "beginner":
        return "你是一个友善的导游，用简单语言解释"
    else:
        return "你是一个简洁的天气助手"

# 6. 创建动态 Agent
agent_dynamic = create_agent(
    model="openai:gpt-4o-mini",
    tools=[get_weather],
    middleware=[role_based_prompt],  # 注入动态提示
    context_schema=Context
)

print("\n=== 动态 System Prompt（专家角色）===")
response2 = agent_dynamic.invoke(
    {"messages": [{"role": "user", "content": "北京天气"}]},
    context={"user_role": "expert"}
)
print(f"AI: {response2['messages'][-1].content}")

print("\n=== 动态 System Prompt（新手角色）===")
response3 = agent_dynamic.invoke(
    {"messages": [{"role": "user", "content": "北京天气"}]},
    context={"user_role": "beginner"}
)
print(f"AI: {response3['messages'][-1].content}")

from langchain.agents import create_agent
from langchain_core.tools import tool

# 1. 定义天气查询工具
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息。"""
    weather_data = {
        "北京": "晴朗，气温25°C",
        "上海": "多云，气温28°C",
        "广州": "小雨，气温30°C"
    }
    return f"{city}的天气是：{weather_data.get(city, '未知')}"

# 2. 定义数学计算工具
@tool
def calculate(expression: str) -> str:
    """计算一个数学表达式的结果。"""
    try:
        result = eval(expression)
        return f"计算结果是：{result}"
    except Exception as e:
        return f"计算出错：{str(e)}"

# 3. 初始化LLM
llm = load_chat_model(model="gpt-4o-mini",provider="openai")

# 4. 创建Agent
agent = create_agent(
    model=llm,
    tools=[get_weather, calculate],
    system_prompt=("""
        你是一个多功能的 AI 助手，能够调用以下工具：
        1. `get_weather(city)`：查询指定城市的天气信息。参数 city 为城市名称（如“北京”）。
        2. `calculate(expression)`：计算数学表达式。参数 expression 为合法的 Python 表达式（如“25 - 28”）。
        请始终遵循以下最佳实践：
        • 当用户询问天气时，先提取城市名，再调用 `get_weather`，并返回自然语言总结。
        • 当用户需要计算时，先提取表达式，再调用 `calculate`，并给出易读的结果说明。
        • 若问题同时涉及天气与计算，按顺序依次调用对应工具，最后整合答案。
        • 禁止编造数据，必须调用工具获取结果后再回答。
        • 所有数字、单位、符号务必与工具返回保持一致，避免主观臆断。
        """
    )
)

# 5. 测试多工具调用
user_queries = [
    "北京和上海的天气怎么样？",
    "如果北京气温是25度，上海是28度，那么北京的温度比上海低多少度？"
]

# 6. 配置会话 ID
config = {"configurable": {"thread_id": "user_123"}}  # 会话 ID

# 7. 流式输出，实时观察推理过程
for step in agent.stream(
    {"messages": [{"role": "user", "content": "北京和上海的天气怎么样？"}]},
    config=config,
    stream_mode="values"    # 返回每个step步骤的完整消息列表，便于调试和观察
):
    # 获取最新消息并格式化打印
    message = step["messages"][-1]
    message.pretty_print()
    print("-" * 50)

# 监控Agent执行时间
# for chunk in agent.stream(..., stream_mode="values"):
#       steps += 1
#       elapsed = time.time() - start
#       print(f"步骤 {steps} 耗时：{elapsed:.2f}s")
#
# print(f"总耗时：{time.time() - start:.2f}s，总步骤：{steps}")


# 捕获中间步骤的错误
# for chunk in agent.stream(..., stream_mode="values"):
#     messages = chunk["messages"]
#     if messages[-1].type == "error":
#         print(f"步骤出错：{messages[-1].content}")
#         # 回滚到上一步的状态
#         last_valid_state = messages[-2]
#         # 重新执行...

import os
from langchain.agents import create_agent
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage, SystemMessage, trim_messages

# ============ 初始化 LLM ============
model = load_chat_model(model="gpt-4o-mini",provider="openai")

# ============ 定义工具函数 ============
@tool
def get_user_info(name: str) -> str:
    """查询用户信息，返回姓名、年龄、爱好"""
    user_db = {
        "陈明": {"age": 28, "hobby": "旅游、滑雪、喝茶"},
        "张三": {"age": 32, "hobby": "编程、阅读、电影"}
    }
    info = user_db.get(name, {"age": "未知", "hobby": "未知"})
    return f"姓名: {name}, 年龄: {info['age']}岁, 爱好: {info['hobby']}"

# ============ 1. 基础短期记忆：InMemorySaver ============
"""
开发环境使用内存存储，重启后记忆丢失。
关键参数：
- checkpointer: 记忆存储对象
- thread_id: 会话唯一标识（用户级隔离）
"""
def demo_inmemory_memory():
    print("=" * 60)
    print("场景 1: 内存记忆（开发环境）")
    print("=" * 60)

    # 创建内存检查点
    memory = InMemorySaver()

    # 创建 Agent（自动继承对话记忆能力）
    agent = create_agent(
        model=model,
        tools=[get_user_info],
        checkpointer=memory  # 启用短期记忆
    )

    # 配置：thread_id 作为会话 ID
    config = {"configurable": {"thread_id": "user_123"}}

    # 第一轮对话：自我介绍
    response1 = agent.invoke(
        {"messages": [{"role": "user", "content": "你好，我叫陈明，好久不见！"}]},
        config=config
    )
    print(f"用户：你好，我叫陈明，好久不见！")
    print(f"AI: {response1['messages'][-1].content}")
    print("-" * 40)

    # 第二轮对话：测试记忆
    response2 = agent.invoke(
        {"messages": [{"role": "user", "content": "请问你还记得我叫什么名字吗？"}]},
        config=config  # 使用相同 thread_id，自动携带上下文
    )
    print(f"用户：请问你还记得我叫什么名字吗？")
    print(f"AI: {response2['messages'][-1].content}")
    print("-" * 40)

    # 验证记忆状态
    state = agent.get_state(config)
    print(f"当前记忆轮次: {len(state.values['messages'])} 条消息")

    # 新开一个会话（不同 thread_id）
    config2 = {"configurable": {"thread_id": "user_456"}}
    response3 = agent.invoke(
        {"messages": [{"role": "user", "content": "我们之前聊过吗？"}]},
        config=config2
    )
    print(f"新会话 AI: {response3['messages'][-1].content}")  # 应无记忆

demo_inmemory_memory()

#系统安装postgresql
#brew install postgresql

#!pip install langgraph-checkpoint-postgres  # 生产环境使用

# 测试数据库是否连接正常
!psql -U myuser -d mydatabase -c "SELECT version();"

from langgraph.checkpoint.postgres import PostgresSaver

"""
生产环境使用数据库存储，支持：
- 持久化（重启不丢失）
- 多实例共享（分布式部署）
- 大规模并发
"""

print("\n" + "=" * 60)
print("场景 2: Postgres 持久化记忆（生产环境）")
print("=" * 60)

# 定义工具函数：查询用户信息
@tool
def get_user_info(name: str) -> str:
    """查询用户信息，返回姓名、年龄、爱好"""
    user_db = {
        "陈明": {"age": 28, "hobby": "旅游、滑雪、喝茶"},
        "张三": {"age": 32, "hobby": "编程、阅读、电影"}
    }
    info = user_db.get(name, {"age": "未知", "hobby": "未知"})
    return f"姓名: {name}, 年龄: {info['age']}岁, 爱好: {info['hobby']}"

# 创建模型
model = load_chat_model(model="gpt-4o-mini",provider="openai")

# 数据库连接字符串
DB_URI = "postgresql://myuser:123456@localhost:5432/mydatabase"

# 使用上下文管理器确保连接正确关闭
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    # 自动创建表结构（仅首次运行）
    checkpointer.setup()

    # 创建智能体
    agent = create_agent(
        model=model,
        tools=[get_user_info],
        checkpointer=checkpointer
    )

    # 配置线程 ID（用于区分不同用户）
    config = {"configurable": {"thread_id": "production_user_001"}}

    # 模拟用户注册流程
    agent.invoke(
        {"messages": [{"role": "user", "content": "我是新用户张三，请记录我的信息"}]},
        config=config
    )

    response = agent.invoke(
        {"messages": [{"role": "user", "content": "我是谁？"}]},
        config=config
    )
    print(f"AI: {response['messages'][-1].content}")

!psql -U myuser -d mydatabase -c "SELECT * FROM checkpoints WHERE thread_id = 'production_user_001' LIMIT 3;"

# ============ 核心模块导入 ============
import os
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver  # 短期记忆存储
from langchain_core.messages import trim_messages  # 消息裁剪工具
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# ============ 初始化 LLM 与工具 ============
model = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    temperature=0.7
)

# 1. 定义天气查询工具
@tool
def search_weather(city: str) -> str:
    """查询城市天气"""
    return f"{city}天气：晴，25°C"

# ============ 配置裁剪参数 ============
MAX_TOKENS = 100  # 根据模型上下文长度调整，上下文128K的话，一般设置为4000左右
TRIM_STRATEGY = "last"  # 保留最新消息
INCLUDE_SYSTEM = True  # 系统消息不参与裁剪

"""
tiktoken 是 OpenAI 官方 token 编码库，支持精确计数。
不同模型需使用不同编码：
- gpt-4o, gpt-4o-mini: o200k_base
- gpt-4-turbo: cl100k_base
- text-davinci-003: p50k_base
"""
import tiktoken

# 2. 定义 Token 编码器获取函数
def get_token_encoder(model_name: str = "gpt-4o-mini"):
    """获取 tiktoken 编码器实例"""
    try:
        # 自动映射模型到编码器
        encoding = tiktoken.encoding_for_model(model_name)
        print(f"✅ 已加载模型 '{model_name}' 的 tiktoken 编码器")
        return encoding
    except KeyError:
        # 如果模型不在官方列表，使用默认编码
        print(f"⚠️  模型 '{model_name}' 未在映射表中，使用默认编码 'o200k_base'")
        return tiktoken.get_encoding("o200k_base")

# 3. 初始化编码器（全局复用提升性能）
TOKEN_ENCODER = get_token_encoder("gpt-4o-mini")

# ============ 2. 精确 Token 计数函数 ============
def count_tokens_tiktoken(messages):
    """
    使用 tiktoken 精确计算消息列表的 token 总数

    参数:
        messages: List[BaseMessage] - 消息对象列表

    返回:
        int: 总 token 数

    说明:
        - 每个消息包含角色、内容和元数据，需完整编码
        - 不同消息格式（Human/AI/System）的 token 开销不同
        - 此函数模拟真实 API 调用的 token 计数逻辑
    """
    total_tokens = 0

    # 遍历每条消息，累加 token 数
    for message in messages:
        # 消息格式：role + content + 特殊标记
        # 典型格式：<|im_start|>role<|im_end|>content<|im_end|>

        # 角色 token（如 "user", "assistant", "system"）
        role_tokens = len(TOKEN_ENCODER.encode(message.type))

        # 内容 token（消息正文）
        content_tokens = len(TOKEN_ENCODER.encode(message.content))

        # 格式开销（OpenAI 消息边界标记）
        # 每条约 4 个特殊 token（开始、角色、结束、内容结束）
        format_overhead = 4

        total_tokens += role_tokens + content_tokens + format_overhead

    return total_tokens

# ============ 创建Agent ============
def create_trimmed_agent():
    """创建 Agent 并配置 InMemorySaver"""
    memory = InMemorySaver()

    agent = create_agent(
        model=model,
        tools=[search_weather],
        system_prompt="你是一个简洁的助手，记住用户提到的城市名称",
        checkpointer=memory  # 启用短期记忆
    )

    return agent

# ============ 手动裁剪并调用 Agent ============
def invoke_with_trim(agent, user_input: str, config: dict):
    """
    在调用 Agent 前手动裁剪上下文

    流程：
    1. 获取当前状态（所有历史消息）
    2. 使用 trim_messages 裁剪
    3. 构建新输入（裁剪后的消息 + 新消息）
    4. 调用 Agent
    """
    # 1. 获取当前记忆状态
    state = agent.get_state(config)
    existing_messages = state.values.get("messages", []) if state else []

    # 精确计算当前 token 数
    current_tokens = count_tokens_tiktoken(existing_messages)

    # 2. 如果有历史消息，先裁剪
    if existing_messages:
        print(f"裁剪前消息数: {len(existing_messages)}")

        # 核心：调用 trim_messages 进行裁剪
        trimmed_messages = trim_messages(
            existing_messages,                    # 待裁剪的消息列表
            max_tokens=MAX_TOKENS,                # 允许的最大 token 数，超过则触发裁剪
            token_counter=count_tokens_tiktoken,  # token 计数函数
            strategy=TRIM_STRATEGY,               # 裁剪策略，"last" 保留最新消息，"first" 保留最早消息
            include_system=INCLUDE_SYSTEM,        # 是否保留系统消息（通常必须保留）
            allow_partial=False,                  # False不允许部分消息，会尝试保留消息的完整性。如果无法在保持消息完整性的前提下将总token数裁剪到参数 max_tokens 以内，就会返回 空列表
            start_on="human"                      # 从 human 消息开始裁剪
        )

        # 计算裁剪后的 token 数
        new_tokens = count_tokens_tiktoken(trimmed_messages)

        print(f"裁剪后 token: {new_tokens}，裁剪后消息数: {len(trimmed_messages)}")
    else:
        trimmed_messages = []

    # 3. 添加新消息
    new_messages = trimmed_messages + [HumanMessage(content=user_input)]

    # 4. 调用 Agent（checkpointer 会自动保存新状态）
    response = agent.invoke(
        {"messages": new_messages},
        config=config
    )

    return response

# ============ 演示多轮对话与裁剪 ============
def demo_manual_trim():
    """演示手动裁剪上下文的多轮对话"""
    print("=" * 60)
    print("场景：手动 trim_messages + InMemorySaver")
    print("=" * 60)

    agent = create_trimmed_agent()
    config = {"configurable": {"thread_id": "trim_user_001"}}

    # 模拟多轮对话
    conversations = [
        "你好，我叫陈明",
        "查询北京天气",
        "上海呢？",
        "明天北京天气如何？",  # 此时会触发裁剪
        "我是谁？",  # 测试记忆是否保留
    ]

    for i, query in enumerate(conversations, 1):
        print(f"\n--- 第 {i} 轮 ---")
        print(f"用户: {query}")

        # 每次调用前自动裁剪
        response = invoke_with_trim(agent, query, config)

        print(f"AI: {response['messages'][-1].content}")
        print("-" * 40)

# ============ 运行演示 ============
if __name__ == "__main__":
    demo_manual_trim()

class ExtendedState(TypedDict):
    messages: list[BaseMessage]
    user_id: str           # 扩展：用户身份，用于权限控制和个性化
    session_id: str        # 扩展：会话标识，用于对话历史管理
    retry_count: int       # 扩展：重试次数，用于错误处理策略
    original_query: str    # 扩展：原始查询，用于日志和审计

# ============ 自定义 State 扩展 ============
"""
通过 TypedDict 扩展 AgentState，添加业务字段（用户ID、偏好等）。
LangChain 1.0 推荐使用 TypedDict 而非 Pydantic。
"""
from typing import TypedDict, Optional
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import InMemorySaver

# 定义自定义 State 结构
class CustomAgentState(AgentState):
    """扩展的 Agent 状态，包含业务上下文"""
    user_id: str  # 用户唯一标识
    preferences: dict  # 用户偏好（主题、语言等）
    visit_count: int  # 访问次数

# ============ 定义带状态访问的工具 ============
from langchain.tools import ToolRuntime
from langgraph.types import Command
from langchain.messages import ToolMessage

# 定义工具函数：更新用户偏好
@tool
def update_user_preference(runtime: ToolRuntime, theme: str) -> Command:
    """
    更新用户主题偏好，写入短期记忆

    ToolRuntime 提供对 state 和 context 的访问能力：
    - runtime.state: 当前状态（含自定义字段）
    - runtime.context: 调用上下文
    - runtime.tool_call_id: 工具调用ID
    """
    # 从当前状态获取偏好（如果不存在则初始化）
    current_prefs = runtime.state.get("preferences", {})
    current_prefs["theme"] = theme

    # 返回 Command 对象，指示状态更新
    return Command(update={
        "preferences": current_prefs,
        "messages": [
            ToolMessage(
                content=f"成功更新主题为: {theme}",
                tool_call_id=runtime.tool_call_id
            )
        ]
    })

# 定义工具函数：根据用户偏好生成问候
@tool
def greet_user(runtime: ToolRuntime) -> str:
    """根据用户偏好生成个性化问候"""
    user_name = runtime.state.get("user_id", "访客")
    prefs = runtime.state.get("preferences", {})
    theme = prefs.get("theme", "默认")

    return f"欢迎回来，{user_name}！当前主题: {theme}"

# ============ 创建带自定义状态的 Agent ============
def demo_custom_state():
    print("\n" + "=" * 60)
    print("场景 3: 自定义 State 扩展记忆维度")
    print("=" * 60)

    # 使用内存存储
    checkpointer = InMemorySaver()

    # 创建 Agent，指定自定义 state_schema
    agent = create_agent(
        model=model,
        tools=[update_user_preference, greet_user],
        state_schema=CustomAgentState,  # 关键：传入自定义状态类型
        checkpointer=checkpointer
    )

    # 配置线程 ID（用于区分不同用户）
    config = {"configurable": {"thread_id": "custom_state_user"}}

    # 第一轮：初始化用户信息
    result1 = agent.invoke(
        {
            "messages": [{"role": "user", "content": "设置主题为暗黑模式"}],
            "user_id": "user_789",  # 自定义字段
            "preferences": {"language": "zh-CN"},  # 初始偏好
            "visit_count": 1
        },
        config=config
    )
    print(f"第一轮: {result1['messages'][-1].content}")
    print("-" * 40)

    # 第二轮：读取记忆
    result2 = agent.invoke(
        {"messages": [{"role": "user", "content": "打个招呼"}]},
        config=config
    )
    print(f"第二轮: {result2['messages'][-1].content}")
    print("-" * 40)

    # 查看完整状态
    state = agent.get_state(config)
    print("当前记忆状态:")
    print(f"  用户ID: {state.values.get('user_id')}")
    print(f"  偏好: {state.values.get('preferences')}")
    print(f"  消息数: {len(state.values['messages'])}")

# ============ 运行自定义状态示例 ============
if __name__ == "__main__":
    demo_custom_state()