from langgraph.agent import create_agent
from langgraph.store.postgres import PostgresStore
from langgraph.prebuilt.tool_runtime import ToolRuntime
from typing import NotRequired
from langgraph.graph.state import AgentState

# 1. 初始化模型
from langchain.chat_models import init_chat_model
model = init_chat_model(model="gpt-4o-mini")

# 2. 数据库连接地址
DB_URL = "postgresql://postgres:123456@localhost:5432/langchadb"

# 3. 使用上下文管理器初始化 PostgresStore
with PostgresStore.from_conn_string(DB_URL) as store:
    store.setup()  # 首次运行创建相关表

    # 4. 自定义 AgentState
    class CustomAgentState(AgentState):
        user_id: NotRequired[str]

    # 5. 定义工具（与 InMemoryStore 版完全相同）
    def save_user_info(name: str, runtime: ToolRuntime) -> str:
        """将客户信息保存在长期记忆中。"""
        user_id = runtime.state["user_id"]
        runtime.store.put(
            namespace=("users",),
            key=user_id,
            value={"name": name}
        )
        return "saved"

    def get_user_info(runtime: ToolRuntime) -> str:
        """从长期记忆中读取客户信息。"""
        user_id = runtime.state["user_id"]
        item = runtime.store.get(namespace=("users",), key=user_id)
        if item:
            return str(item.value)
        return "INFO NOT FOUND"

    # 6. 系统提示词
    SYSTEM_PROMPT = """
    用户提及个人信息时，使用工具保存用户信息。
    如果用户询问个人信息时，尝试使用工具检索/读取用户信息。
    """

    # 7. 创建 Agent
    agent = create_agent(
        model=model,
        tools=[save_user_info, get_user_info],
        store=store,
        system_prompt=SYSTEM_PROMPT,
        state_schema=CustomAgentState
    )

    # 8. 第一次会话：存储信息
    response = agent.invoke({
        "messages": [{"role": "user", "content": "你好，我是小花"}],
        "user_id": "user-1"
    })
    print(response["messages"][-1].content)
    # "很高兴认识你，小花！"

    # 9. 第二次会话：读取信息
    response = agent.invoke({
        "messages": [{"role": "user", "content": "我是谁？"}],
        "user_id": "user-1"
    })
    print(response["messages"][-1].content)
    # "你是小花。"