010-1-自定义时间中间件

举例我想实现一个时间中间件，模型调用时知道当前时间，日期，节假日，距离上次会话的时间，以增加agent 回复内容的精准，知道用户，比如很短时间想要恢复，或者间隔了很久，知道是早上晚上，工作日，节假日，这样的中间件如何写

🎯 时间上下文中间件

核心代码

import holidays
from datetime import datetime, timezone
from zoneinfo import ZoneInfo
from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse
from langchain.messages import SystemMessage
from typing import Callable
import pytz


class TimeContextMiddleware:
    """时间上下文中间件 - 让 Agent 感知时间、日期、节假日、会话间隔"""

    def __init__(self, timezone: str = "Asia/Shanghai"):
        self.timezone = ZoneInfo(timezone)
        self.country_holidays = holidays.CN()  # 中国节假日

    def _get_session_interval(self, state) -> str:
        """计算距离上次会话的时间间隔"""
        messages = state.get("messages", [])
        if len(messages) < 2:
            return "这是会话的开始"

        # 获取最后一次用户消息的时间（需要预先存储）
        last_message_time = state.get("last_message_timestamp")
        if not last_message_time:
            return "无法获取上次会话时间"

        now = datetime.now(self.timezone)
        delta = now - last_message_time

        minutes = delta.total_seconds() / 60
        hours = minutes / 60
        days = delta.days

        if minutes < 1:
            return "刚刚（几秒前）"
        elif minutes < 60:
            return f"{int(minutes)}分钟前"
        elif hours < 24:
            return f"{int(hours)}小时前"
        else:
            return f"{days}天前"

    def _get_time_context(self, state) -> str:
        """构建完整的时间上下文"""
        now = datetime.now(self.timezone)
        hour = now.hour
        weekday = now.weekday()  # 0=周一, 6=周日
        is_weekend = weekday >= 5

        # 时段判断
        if 5 <= hour < 12:
            period = "早晨/上午"
        elif 12 <= hour < 14:
            period = "中午"
        elif 14 <= hour < 18:
            period = "下午"
        elif 18 <= hour < 22:
            period = "晚上"
        else:
            period = "深夜/凌晨"

        # 节假日判断
        is_holiday = now.date() in self.country_holidays
        holiday_name = self.country_holidays.get(now.date()) if is_holiday else None

        # 会话间隔
        interval = self._get_session_interval(state)

        return f"""
## 当前时间上下文
- 当前时间：{now.strftime('%Y-%m-%d %H:%M:%S %A')}
- 时段：{period}
- 是否周末：{'是' if is_weekend else '否'}
- 是否节假日：{'是（' + holiday_name + '）' if is_holiday else '否'}
- 距离上次对话：{interval}

## 时间敏感建议
- {"深夜时间，回复应简洁，避免打扰感" if period == "深夜/凌晨" else ""}
- {"周末/节假日，回复可更轻松友好" if is_weekend or is_holiday else ""}
- {"用户刚提问过，回复可简洁" if "分钟前" in interval and int(interval.split('分钟')[0]) < 5 else ""}
- {"用户很久没来，可礼貌问候或简要回顾" if "天前" in interval else ""}
"""

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """注入时间上下文到系统提示词"""
        # 获取现有系统消息
        existing_content = request.system_message.content_blocks if request.system_message else []

        # 构建时间上下文
        time_context = self._get_time_context(request.state)

        # 追加到系统提示词
        new_content = list(existing_content) + [
            {"type": "text", "text": time_context}
        ]

        new_system_message = SystemMessage(content=new_content)
        updated_request = request.override(system_message=new_system_message)

        # 可选：更新 state 中的时间戳（供下次使用）
        response = handler(updated_request)

        return response

    # 异步版本
    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        existing_content = request.system_message.content_blocks if request.system_message else []
        time_context = self._get_time_context(request.state)
        new_content = list(existing_content) + [{"type": "text", "text": time_context}]
        new_system_message = SystemMessage(content=new_content)
        updated_request = request.override(system_message=new_system_message)
        return await handler(updated_request)

---

使用方式

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

# 创建 agent 时需要存储消息时间戳
class TimestampState(AgentState):
    last_message_timestamp: NotRequired[datetime]

# 使用中间件
agent = create_agent(
    model="gpt-4o",
    tools=[...],
    state_schema=TimestampState,  # 需要 state 存储时间戳
    checkpointer=InMemorySaver(),
    middleware=[TimeContextMiddleware(timezone="Asia/Shanghai")],
)

# 调用时，中间件会自动注入时间上下文
config = {"configurable": {"thread_id": "user_123"}}
result = agent.invoke(
    {"messages": [HumanMessage("今天有什么新闻？")]},
    config
)
# Agent 知道是今天、什么时段、是否周末/节假日

---

增强版：自动更新时间戳

class TimeContextWithTrackingMiddleware(TimeContextMiddleware):
    """增强版：自动记录/更新时间戳"""

    def wrap_model_call(self, request, handler):
        # 注入时间上下文
        time_context = self._get_time_context(request.state)
        # ... 同上述逻辑 ...

    def after_model(self, state, runtime):
        """模型调用后，更新时间戳"""
        return {
            "last_message_timestamp": datetime.now(self.timezone)
        }

---

💡 一句话总结

wrap_model_call 钩子动态读取当前时间、节假日、会话间隔，通过 request.override() 追加到系统提示词，让 Agent 在回答时知晓完整的时间上下文。