vent_agent_2.0/extensions/main_fan/exploration_tools.py

"""业务只保留网关取数；准备、沙箱执行和提交是宿主基础设施。"""
import asyncio
from functools import wraps
import inspect
import json
import subprocess
from typing import Annotated, Any
from weakref import WeakValueDictionary

from langchain.tools import ToolRuntime
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from pydantic import BeforeValidator, ValidationError

from extensions import define_tool
from extensions.main_fan import exploration
from runtime import python_sandbox

_locks = WeakValueDictionary()


def _list_argument(value):
    # 部分兼容模型把数组参数再次编码成 JSON 字符串；保持工具 schema 为数组，容忍这层编码。
    if isinstance(value, str):
        value = json.loads(value)
    return value


def _workspace_argument(value):
    if isinstance(value, str) and value and not value.startswith('/'):
        return '/workspace/' + value.removeprefix('./')
    return value


WorkspacePath = Annotated[str, BeforeValidator(_workspace_argument)]
WorkspacePaths = Annotated[list[WorkspacePath], BeforeValidator(_list_argument)]
StringList = Annotated[list[str], BeforeValidator(_list_argument)]


def _identity(runtime):
    identity = runtime.config.get('configurable', {})
    if not identity.get('workspace_user') or not identity.get('workspace_session'):
        raise ValueError('探索工具缺少服务端会话身份')
    return identity, identity['workspace_user'], identity['workspace_session']


def _action(function):
    """同会话串行写入；可纠正的交付错误交还 Agent 继续探索。"""
    signature = inspect.signature(function)

    @wraps(function)
    async def wrapped(*args, **kwargs):
        try:
            runtime = signature.bind(*args, **kwargs).arguments['runtime']
            _, user, session = _identity(runtime)
            key = (user, session)
            lock = _locks.get(key)
            if lock is None:
                lock = _locks[key] = asyncio.Lock()
            async with lock:
                return await function(*args, **kwargs)
        except (ValueError, OSError, subprocess.SubprocessError) as exc:
            detail = (json.dumps(exc.errors(include_input=False, include_url=False), ensure_ascii=False, default=str)
                      if isinstance(exc, ValidationError) else str(exc))
            return {'ok': False, 'error_type': type(exc).__name__, 'message': detail[:3000],
                    'next_action': '读取反馈、修正代码或补取原始数据后重试；不能把本次失败写成已完成。'}
    return wrapped


@define_tool(group='自主探索', label='准备任务与可信输入')
@_action
async def prepare_main_fan_task(target: dict[str, Any] | str, runtime: ToolRuntime) -> dict:
    """绑定任务身份，导出本轮原始输入、原件路径、规则、身份和交付 schema；不解析文档或计算。

    target 传用户原目标或明确给出的输入。队列任务使用服务端已校验输入，不自行重造确认或时序。
    后续文档解析、数值分析和报告生成均由你按 Skill 写 Python，在沙箱中执行。
    """
    identity, user, session = _identity(runtime)
    trusted = identity.get('business_task_input') if identity.get('business_agent') == 'main_fan' else None
    return await asyncio.to_thread(exploration.prepare, target, user, session, trusted,
                                   identity.get('business_task_id'))


@define_tool(group='主通风机', label='取得本次网关资料')
@_action
async def fetch_main_fan_data(workflow_id: str, runtime: ToolRuntime) -> dict:
    """use_gateway=true 时取得已配置网关的设备、原件和时序，更新 input.json/data.json；不解析或分析。"""
    _, user, session = _identity(runtime)
    return await exploration.fetch(workflow_id, user, session)


@define_tool(group='自主探索', label='在沙箱执行 Agent 编写的 Python')
@_action
async def run_python(script_path: WorkspacePath, runtime: ToolRuntime, arguments: StringList | None = None,
                       output_paths: WorkspacePaths | None = None, timeout_seconds: int = 60,
                       workflow_id: str | None = None) -> dict:
    """执行你用 write_file/edit_file 编写的 /workspace/*.py，返回真实 stdout/stderr、退出码和 run_id。

    代码当前目录为本会话 workspace，使用相对路径。可用标准库、numpy、pypdfium2、Pillow、cv2；
    无网络、不传服务端凭证、不支持临时安装依赖。自己编写 PDF 文本提取/渲染、计算及报告生成代码。
    对扫描件可输出 PNG 并用 read_file 看图，之后自行编写提取与分析脚本，不调用旧业务算法。
    必须传 workflow_id 关联本轮交付，声明 output_paths 保存输出哈希；原件与可信输入不要修改。
    报告 Markdown、研判 JSON 和附件都须由成功脚本生成，并在提交时引用相应 run_id。
    """
    _, user, session = _identity(runtime)
    if workflow_id:
        workflow = await asyncio.to_thread(exploration._load, workflow_id, user, session)
        await asyncio.to_thread(exploration.validate_inputs, workflow, user, session)
    receipt = await python_sandbox.run(user, session, script_path, arguments or (), output_paths or (), timeout_seconds)
    if workflow_id:
        await asyncio.to_thread(exploration.record_python, workflow_id, user, session, receipt)
        receipt['python_paths'] = exploration.python_paths(workflow)
    if not receipt['ok']:
        receipt['next_action'] = '代码执行失败；先按 stderr 最后的异常类型和行号修正代码，再重新运行。'
        if 'FileNotFoundError' in receipt['stderr'] or 'PermissionError' in receipt['stderr']:
            receipt['next_action'] += ('代码内使用 python_paths 中的相对路径；不要用 /workspace/ 绝对路径，'
                                       '也不要把 exploration/任务ID/input.json 缩成 input.json。')
        elif 'AttributeError' in receipt['stderr']:
            receipt['next_action'] += '文档库 API 用法见 Skill 的 references/document-reading.md，先核对实际可用方法。'
    elif receipt['missing_outputs'] or receipt['unchanged_outputs']:
        receipt['next_action'] = ('脚本退出码为0，但部分声明文件未生成或本次未改写。核对 stdout 中实际保存位置，'
                                  '将正确的 /workspace/ 输出路径传入 output_paths 后重新运行生成代码，再提交。')
    return receipt


@define_tool(group='自主探索', label='核验沙箱交付并归档')
@_action
async def submit_main_fan_report(workflow_id: str, report_path: WorkspacePath, calculations_path: WorkspacePath,
                                  runtime: ToolRuntime, python_run_ids: StringList,
                                  artifacts: WorkspacePaths | None = None) -> Any:
    """提交沙箱脚本实际生成的 Markdown 和 AgentAssessment JSON，不重跑文档解析、业务计算或报告模板。

    格式见准备工具返回的 submission_schema_path。宿主只核验身份、原件指纹、来源引用、文件格式和
    本轮成功执行输出。python_run_ids 必填，报告、JSON 和附件必须都在相应成功执行的 output_paths 中。
    不生成 MainFanResult/analysis_id。失败可修正重提；成功提交作为最后操作，须单独调用。
    """
    _, user, session = _identity(runtime)
    calls = getattr(runtime.state.get('messages', [])[-1], 'tool_calls', []) if runtime.state.get('messages') else []
    if len(calls) > 1:
        raise ValueError('提交报告应单独调用，不能与其他工具并行执行')
    result = await asyncio.to_thread(exploration.submit, workflow_id, user, session, report_path,
        calculations_path, python_run_ids, artifacts or ())
    receipt = {'ok': True, 'status': result['status'], 'analysis_id': result['analysis_id'],
               'workflow_id': workflow_id, 'sources': result['sources'], 'message': '报告已归档，本轮研判完成。'}
    return Command(update={'structured_response': result, 'messages': [ToolMessage(
        content=json.dumps(receipt, ensure_ascii=False), tool_call_id=runtime.tool_call_id,
        name='submit_main_fan_report')]})