from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph
from langgraph.prebuilt.chat_agent_executor import AgentState
from langchain_core.messages import RemoveMessage
from typing import Optional

# 1. 创建模型和 Agent
model = ...  # ChatOpenAI 或其他模型
agent = create_react_agent(model, tools=[])

# 2. 定义裁剪中间件 — 使用 @before_model
#    在每次调用模型之前触发
def trim_messages(state: AgentState, runtime) -> Optional[dict]:
    """消息裁剪中间件：保留首条 + 最近 N 条消息"""
    messages = state["messages"]
    
    # 消息数 <= 3 时无需裁剪
    if len(messages) <= 3:
        return None
    
    # 保留第一条消息（通常是系统消息 / 角色设定）
    first_message = messages[0]
    
    # 根据消息总数奇偶性决定保留最近几条
    length = len(messages)
    if length % 2 == 0:
        # 偶数条：保留最近 3 条
        recent_messages = messages[-3:]
    else:
        # 奇数条：保留最近 4 条
        recent_messages = messages[-4:]
    
    # 组装新消息列表
    new_messages = [first_message] + list(recent_messages)
    
    # 返回裁剪指令：删除所有旧消息，替换为新列表
    return {
        "messages": [
            RemoveMessage(id=m.id),  # 删除原有消息
            *new_messages             # 添加裁剪后消息
        ]
    }

# 3. 注册中间件到 Agent
#    （将 trim_messages 作为 agent 的 middleware 或通过 graph 绑定）

# 4. 多次对话测试
config = {"configurable": {"thread_id": "1"}}

# 第1次 — 消息数 <= 3，不裁剪
agent.invoke(
    {"messages": [HumanMessage(content="你好，我是老王")]},
    config=config
)

# 第2次 — 消息数 <= 3，不裁剪
agent.invoke(
    {"messages": [HumanMessage(content="从现在起你可以叫我小王")]},
    config=config
)

# 第3次 — 消息数 = 5，触发裁剪（奇数→保留最近4条+首条=5条，实际未裁）
agent.invoke(
    {"messages": [HumanMessage(content="今天天气不错")]},
    config=config
)

# 第4次 — 消息数 = 7，触发裁剪（奇数→保留4条+首条=5条，裁掉2条旧消息）
agent.invoke(
    {"messages": [HumanMessage(content="告诉我你记得我是谁？")]},
    config=config
)