vent_agent_2.0/agents/middleware.py

"""动态时间与 Skill 任务的工具提交收口；压缩和 HITL 由框架处理。"""

from __future__ import annotations

import re
from typing import Any, Callable

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

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))


class SkillCompletionMiddleware(AgentMiddleware):
    """提交工具写入正式结果后结束；校验失败仍交还模型修正。"""

    def before_agent(self, state, runtime):
        return {'structured_response': None}

    async def abefore_agent(self, state, runtime):
        return self.before_agent(state, runtime)

    def wrap_model_call(self, request, handler):
        # 完成通过提交工具表达，避免模型只写一段计划就结束队列任务。
        return handler(request.override(tool_choice='any'))

    async def awrap_model_call(self, request, handler):
        return await handler(request.override(tool_choice='any'))

    @hook_config(can_jump_to=['end'])
    def before_model(self, state, runtime):
        if state.get('structured_response') is not None:
            return {'jump_to': 'end'}
        return None

    @hook_config(can_jump_to=['end'])
    async def abefore_model(self, state, runtime):
        return self.before_model(state, runtime)


__all__ = ["DynamicTimeMiddleware", "SkillCompletionMiddleware"]