Memory Basics - main.py

"""
LangChain 1.0 - Memory Basics (内存管理基础)
==========================================

本模块重点讲解 LangGraph 提供的内存管理机制，包括：
1. InMemorySaver - LangGraph 提供的内存管理
2. checkpointer 参数 - 为 Agent 添加内存
3. thread_id - 会话管理
4. 多轮对话状态保持
"""

from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from init_model import get_chat_model

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

# 导入自定义工具
from tools import weather
from tools import web_search
from tools import calculator

# 初始化数据
weather_tool = weather.get_weather
web_search_tool = web_search.web_search
calculator_tool = calculator.calculator
chat_model = get_chat_model()


# 创建一个简单的工具
@tool
def get_user_info(user_id: str) -> str:
    """获取用户信息"""
    users = {"123": "张三，25岁，工程师", "456": "李四，30岁，设计师"}
    return users.get(user_id, "用户不存在")


# ============================================================================
# 示例 简单客服机器人
# ============================================================================
def demo_simple_customer_service():
    """
    演示简单客服机器人
    """
    agent = create_agent(
        model=chat_model,
        tools=[get_user_info],
        checkpointer=InMemorySaver(),  # 添加内存
    )

    config = {"configurable": {"thread_id": "session_1"}}

    print("\n第一轮：查询用户信息")
    response1 = agent.invoke(
        {"messages": [{"role": "user", "content": "查询用户 123 的信息"}]},
        config=config,
    )
    print(f"Agent: {response1['messages'][-1].content}")

    print("\n第二轮：询问之前的信息")
    response2 = agent.invoke(
        {"messages": [{"role": "user", "content": "你之前查询的用户 123 信息是？"}]},
        config=config,
    )
    print(f"Agent: {response2['messages'][-1].content}")


if __name__ == "__main__":
    demo_simple_customer_service()