vent_agent_2.0/extensions/main_fan/exploration_tools.py

"""业务只保留网关取数；准备、沙箱执行和提交是宿主基础设施。"""
import asyncio
from functools import wraps
import inspect
import json
import math
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 equipment, exploration, realtime, timeseries
from extensions.shaanxi.client import GatewayError
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, GatewayError) 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))
            result = {'ok': False, 'error_type': type(exc).__name__, 'message': detail[:3000],
                      'next_action': '读取反馈、修正代码或补取原始数据后重试；不能把本次失败写成已完成。'}
            if isinstance(exc, GatewayError):
                result.update(retryable=exc.retryable, retry_after=exc.retry_after)
            return result
    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='按测点和时间查询历史时序')
@_action
async def query_fan_timeseries(workflow_id: str, node_id: str, dept_id: str,
                               start_time: str, end_time: str, runtime: ToolRuntime) -> dict:
    """查询主通风机相关测点的历史曲线，保存完整原始响应和规范化 JSON，不计算切换或研判结论。

    先 prepare_main_fan_task，使用它返回的 workflow_id。node_id、dept_id 必须来自用户或真实
    测点目录，不猜编号或设备归属。单次查一个点，可多次调用查询两台风机的不同点或时间窗口。
    时间支持 2026-10-19 11:00:00 或带时区的 ISO 格式；无时区按 Asia/Shanghai，精确到秒，
    end_time 必须晚于 start_time，默认最多24小时。网关凭证和 X-Agent-Code 由宿主配置提供。
    返回 sample_count、实际时间范围、质量计数、少量 preview、data_path、raw_path 和
    python_paths；空数据会明确返回0条。Python用python_paths中的相对路径读取完整samples，
    不以preview代替全量分析。文件已加入本轮可信来源，可在研判JSON的evidence中引用。
    此工具返回测点原始读数，一级/二级开停不自动解释为整机状态，单位未知时不补猜。
    """
    _, user, session = _identity(runtime)
    return await timeseries.query(workflow_id, node_id, dept_id, start_time, end_time, user, session)


@define_tool(group='主通风机', label='查询安全监控实时传感器列表')
@_action
async def query_fan_sensor_realtime(workflow_id: str, dept_id: str, runtime: ToolRuntime,
                                    sensor_type_list: StringList | None = None,
                                    page_num: int = 1, page_size: int = 100) -> dict:
    """通过 /aqjk/node/getNodeRealtimeDetails 查询本部门传感器目录及当前值，一次查询一页。

    先准备本轮workflow_id；dept_id、类型编码来自用户或真实目录。sensor_type_list可传
    ["0005","0013","1001","1010","0001,0011,0043","0002"]，宿主按索引构造键值对数组。
    不传类型数组时不添加类型筛选，不把无返回类型解释为未安装。page_num从1开始，默认每页100条。
    返回row_count、total、has_more、next_page、少量preview及完整data_path/raw_path。
    按分页信息继续查询，不能把单页当全部。Python用python_paths.data读取JSON的payload原始字段。
    未识别列表结构时row_count为null，需读取完整响应，不当成空数据。文件可作为报告证据引用。
    安全监控测点与重大设备测点编码体系分别核对，设备归属及单位不凭名称补猜。
    """
    _, user, session = _identity(runtime)
    return await realtime.query_sensors(workflow_id, dept_id, page_num, page_size, sensor_type_list, user, session)


@define_tool(group='主通风机', label='查询主通风机重大设备实时列表')
@_action
async def query_fan_equipment_realtime(workflow_id: str, dept_id: str, runtime: ToolRuntime,
                                       page_num: int = 1, page_size: int = 100) -> dict:
    """通过 /zdsb/sbnode/selectSbNodeList 查询主通风机系统重大设备的实时数据，一次查询一页。

    使用准备工具返回的workflow_id与真实dept_id。宿主固定devOfSystem=30，自动补齐鉴权和
    当前业务X-Agent-Code；page_num从1开始，默认page_size=100。返回完整数据文件及分页信息。
    可以从真实返回中查找电机、轴承温度、风量等测点，但不保证一定存在，字段与单位以原文为准。
    Python用python_paths.data读取payload，保留设备及其测点关联；preview不能替代全量或全部分页。
    列表结构未知时会明确提示，不能按0条处理。取得重大设备测点编号后，可调用
    query_fan_equipment_timeseries查询对应历史，不直接套用安全监控曲线接口的测点编号。
    返回文件是本轮可信输入，可在报告evidence中引用。
    """
    _, user, session = _identity(runtime)
    return await realtime.query_equipment(workflow_id, dept_id, page_num, page_size, user, session)


@define_tool(group='主通风机', label='查询重大设备测点历史时序')
@_action
async def query_fan_equipment_timeseries(workflow_id: str, node_id: str, dept_id: str,
                                         start_time: str, end_time: str, runtime: ToolRuntime) -> dict:
    """通过 /zdsb/sbnode/selectSbNodeCurve 查询电机、轴承温度等重大设备测点的历史曲线。

    先准备本轮workflow_id。node_id、dept_id来自用户或重大设备实时列表，不能混用其他部门或
    安全监控测点编号。宿主使用单数nodeId参数，自动补齐凭证和业务身份。
    起止时间支持YYYY-MM-DD HH:MM:SS或带时区ISO时间，无时区按Asia/Shanghai，默认最多24小时。
    原样保存完整上游响应；data_path文件的payload就是原始业务响应。Python使用python_paths.data，
    自行核对实际曲线字段、时间、测点与单位再分析。工具不猜字段、不假设数组项数等于采样点数。
    支持多次查询与报告证据引用；失败反馈ok=false，不伪造空曲线。
    """
    _, user, session = _identity(runtime)
    return await equipment.query_timeseries(workflow_id, node_id, dept_id, start_time, end_time, 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('提交报告应单独调用，不能与其他工具并行执行')
    from extensions.main_fan.batch_context import current_scope
    scope = current_scope.get()
    if scope is not None:
        workflow = exploration._load(workflow_id, user, session, allow_completed=True)
        context = workflow['context']
        if not context.get('report_queries') or 'history_manifest' not in context:
            raise ValueError('本地井级批次必须先查询测定报告并回查本井40天开停和风量历史')
        from extensions.main_fan.report_archive import VERSION
        if any(item.get('archive_version') != VERSION for item in context.get('report_archives', [])):
            raise ValueError('旧报告档案缺少实际全页读取核验，请重新fetch_main_fan_reports，并按当前契约完整读取后归档；不得把旧的已读声明当作事实。')
        candidate = exploration.AgentAssessment.model_validate(exploration._read_json(
            python_sandbox.workspace_file(user, session, calculations_path)))
        switching = candidate.details.get('switch_analysis', {})
        flow = candidate.details.get('flow_analysis', {})
        actual_flow = [row for row in context['history_manifest'].get('equipment_sample_counts', {}).values() if row['role'] == 'flow']
        flow_count = flow.get('raw_sample_count', flow.get('sample_count'))
        if actual_flow and (flow.get('status') not in {'evaluated', 'insufficient_data', 'no_data', 'unrecognized_structure'}
                or (not isinstance(flow_count, int) and not (flow.get('status') == 'unrecognized_structure' and flow_count is None))):
            raise ValueError('请在details.flow_analysis中记录status、raw_sample_count及有效样本数，按实际data[].sbNodeCurve[].dataTime/realData核对风量样本；未知结构的数量用null，不当成空数据。')
        if any(row['recognized_samples'] > 0 for row in actual_flow) and (flow.get('status') == 'no_data' or flow_count in (None, 0)):
            raise ValueError('实际风量sbNodeCurve内已有历史记录。请读取完整嵌套数组并区分全0读数、无效数据和真正空数据，不能声明没有历史数据。')
        if flow.get('status') == 'no_data' and any(row['unrecognized_windows'] for row in actual_flow):
            raise ValueError('仍有未识别的风量响应结构，不能将其宣称为空数据；请自编代码核对或明确记录解析受限。')
        if (switching.get('status') not in {'evaluated', 'insufficient_data'}
                or not isinstance(switching.get('events'), list) or not switching.get('method')
                or 'limitations' not in switching):
            raise ValueError('请在details.switch_analysis中交付status(evaluated/insufficient_data)、method、events数组和limitations。不能只用开停占比代替切换识别；候选事件也须列出实际状态变化时刻和证据。')
        if switching.get('switch_detected') and not switching['events']:
            raise ValueError('宣称发现切换时必须给出具体事件时刻及测点证据；仅有开停占比不能认定已切换。')
        from extensions.main_fan.gateway_queries import parse_time
        for event in switching['events']:
            if not isinstance(event, dict) or event.get('status') not in {'confirmed', 'candidate'}:
                raise ValueError('每个切换事件须为对象，status为confirmed或candidate')
            if not all(key in event for key in ('from_stopped_at', 'to_started_at', 'interval_minutes', 'evidence')):
                raise ValueError('切换事件须含from_stopped_at、to_started_at、interval_minutes和evidence；无法确定的时刻或间隔填null并说明限制')
            for key in ('from_stopped_at', 'to_started_at'):
                if event[key] is not None:
                    parse_time(event[key])
            if not event.get('from_stopped_at') and not event.get('to_started_at') and not event.get('state_changes'):
                raise ValueError('候选事件也须保留至少一条实际状态变化时刻')
            interval = event['interval_minutes']
            if interval is not None and (type(interval) not in (int, float) or not math.isfinite(interval) or interval < 0):
                raise ValueError('切换间隔须为非负有限分钟数，无法计算填null；先启动后停机单列重叠并按Skill处理')
        if candidate.result_code == 'NORMAL' and (not context['history_manifest']['coverage_complete']
                or scope.shaft.identity_status != 'confirmed'):
            raise ValueError('历史覆盖不完整或通风井归属未确认时不能提交NORMAL')
    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')]})