vent_agent_2.0/extensions/main_fan/middleware.py

"""文档传输兼容层；原件解析与渲染仍由 Agent 的沙箱代码完成。"""
from pathlib import PurePosixPath
from threading import RLock

from langchain.agents.middleware.types import AgentMiddleware
from langchain_core.messages import ToolMessage

_page_lock = RLock()


def record_page_read(request, response):
    """记录实际文件工具读取及当时指纹；不接受模型自报已读页数。"""
    from extensions.main_fan.batch_context import current_scope
    scope = current_scope.get()
    call = request.tool_call
    if scope is None or not scope.active_workflow_id or call['name'] != 'read_file':
        return
    virtual = call.get('args', {}).get('file_path', '')
    if not isinstance(virtual, str) or not virtual.startswith('/workspace/'):
        return
    extension = PurePosixPath(virtual).suffix.lower()
    if extension not in {'.png', '.jpg', '.jpeg', '.webp', '.txt'}:
        return
    update = getattr(response, 'update', None)
    messages = [response] if isinstance(response, ToolMessage) else update.get('messages', []) if isinstance(update, dict) else []
    if not messages or any(getattr(message, 'status', 'error') == 'error' for message in messages):
        return
    from runtime.python_sandbox import digest, workspace_file
    from db import main_fan_store
    path = workspace_file(scope.user, scope.session, virtual)
    if not path.is_file():
        return
    offset = call.get('args', {}).get('offset') or 0
    limit = call.get('args', {}).get('limit') or 2000
    complete = extension != '.txt' or (offset == 0 and len(path.read_text().splitlines()) <= limit)
    with _page_lock:
        workflow = main_fan_store.workflow(scope.active_workflow_id, scope.user, scope.session)
        workflow['context'].setdefault('page_reads', {})[virtual] = {'sha256': digest(path), 'complete': complete}
        main_fan_store.update_workflow(scope.active_workflow_id, context=workflow['context'])


def _pdf_read_feedback(request):
    call = request.tool_call
    from extensions.main_fan.batch_context import current_scope
    scope = current_scope.get()
    queries = {'query_fan_sensor_realtime', 'query_fan_equipment_realtime', 'query_fan_timeseries',
               'query_fan_equipment_timeseries', 'fetch_main_fan_data', 'submit_main_fan_mine_plan',
               'query_main_fan_history', 'fetch_main_fan_reports'}
    if scope is not None and scope.active_workflow_id and call['name'] in queries:
        from db import main_fan_store
        from json import dumps
        prepared = main_fan_store.workflow(scope.active_workflow_id, scope.user, scope.session)['context'].get('prepared_data')
        if prepared:
            # 本地批次的取数只由宿主执行；兼容旧提示中的批量工具调用，返回现成收据。
            receipt = (prepared['history'] if call['name'] == 'query_main_fan_history' else
                       prepared['reports'] if call['name'] == 'fetch_main_fan_reports' else None)
            return ToolMessage(content=dumps(receipt or {'ok': False,
                'message': '本批次已由固定代码完成测点发现、历史取数和数据整理。直接读取agent_data_path分析，无需重新规划取数。',
                'agent_data_path': prepared['agent_data_path']}, ensure_ascii=False),
                tool_call_id=call['id'], name=call['name'], status='success' if receipt else 'error')
    if scope is not None and call['name'] in {'ask_user', 'request_plan_approval'}:
        return ToolMessage(content=(
            '当前是用户已授权的本地批次。使用本轮消息明确给出的workflow_id；前一阶段的ID已经归档。'
            '使用宿主已准备的监测数据和报告原件继续分析，数据归属不足时如实提交UNDETERMINED及限制，'
            '无需询问是否继续，不要调用人工确认工具。'
        ), tool_call_id=call['id'], name=call['name'], status='error')
    if call['name'] != 'read_file':
        return None
    path = call.get('args', {}).get('file_path', '')
    if not isinstance(path, str) or PurePosixPath(path).suffix.lower() != '.pdf':
        return None
    # 当前 OpenAI 兼容网关不接受框架的 PDF file 内容块。给模型可恢复反馈，
    # 不在此处隐式 OCR 或调用旧解析器，也不读取这份原件。
    return ToolMessage(content=(
        'PDF 是二进制原件，当前模型接口不能直接接收 PDF 附件。请按 Skill 自己编写 Python，'
        '通过 run_python 使用 pypdfium2 取得页数和文本；扫描页渲染为 PNG，再用 read_file 查看 PNG。'
        '代码内将 /workspace/ 前缀移除，使用工作区相对路径。'
    ), tool_call_id=call['id'], name='read_file', status='error')


class MainFanSandboxMiddleware(AgentMiddleware):
    def wrap_tool_call(self, request, handler):
        feedback = _pdf_read_feedback(request)
        if feedback is not None:
            return feedback
        response = handler(request)
        record_page_read(request, response)
        return response

    async def awrap_tool_call(self, request, handler):
        feedback = _pdf_read_feedback(request)
        if feedback is not None:
            return feedback
        response = await handler(request)
        record_page_read(request, response)
        return response