vent_agent_2.0/agents/middleware.py

"""自定义中间件。

2.0 只保留一个必要中间件：动态时间。其余能力尽量用框架内置
（压缩用框架的摘要中间件、HITL 用 interrupt_on），不自造轮子。
"""

from __future__ import annotations

import re
from typing import Any, Callable

from langchain.agents.middleware.types import AgentMiddleware, ModelRequest, ModelResponse

from agents.prompts import build_time_block
from obs import get_logger

logger = get_logger(__name__)

# 系统提示词里的时间段落。编译期写入的是占位内容，每次调用前整体替换。
_TIME_BLOCK_RE = re.compile(r"\n## 系统时间\n.*?(?=\n## |\Z)", re.DOTALL)


def _refresh_time(request: ModelRequest) -> ModelRequest:
    """把系统提示词里的「## 系统时间」段落刷新为当前时间。

    没有系统消息、或 content 是结构化块（多模态）时原样返回：不去动它，
    避免破坏块结构。
    """
    message = request.system_message
    if message is None:
        return request

    content = message.content
    if not isinstance(content, str):
        return request

    refreshed = _TIME_BLOCK_RE.sub(build_time_block(), content, count=1)
    if refreshed == content and "## 系统时间" not in content:
        refreshed = content + build_time_block()

    if refreshed == content:
        return request
    return request.override(system_message=message.model_copy(update={"content": refreshed}))


class DynamicTimeMiddleware(AgentMiddleware):
    """每次模型调用前刷新系统提示词里的「## 系统时间」段落。

    为什么需要：agent 是编译期单例、长驻进程。如果不刷新，系统提示词里的
    时间会永久停在 agent 构建的那一刻——用户第二天问「今天几号」，模型会
    拿到昨天的日期，且因为它「看得见」一个具体日期，比不知道更危险。

    **同步与异步两套钩子都要实现**：只实现 ``wrap_model_call`` 时，用
    ``astream``/``ainvoke`` 调用会在运行期抛 NotImplementedError——而且报错
    发生在模型节点内部，离「你少写了个异步方法」这个原因很远。本服务全程
    走异步路径，这里两个都实现，同步的实现顺带留给脚本与测试用。
    """

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> Any:
        return handler(_refresh_time(request))

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Any],
    ) -> Any:
        return await handler(_refresh_time(request))


__all__ = ["DynamicTimeMiddleware"]