vent_agent_2.0/extensions/main_fan/batch.py

"""按矿复用main_fan Agent的本地批次：先发现、再逐井多轮执行，结果不进入推送队列。"""
import asyncio
from contextlib import contextmanager
from datetime import datetime
import fcntl
import hashlib
from html import escape
import json
from pathlib import Path
import shutil
from uuid import uuid4
from zoneinfo import ZoneInfo

import config
from agents.workspace import session_root
from db import chat_store
from extensions.main_fan import exploration, mine_catalog
from extensions.main_fan.gateway_queries import parse_time
from extensions.main_fan.batch_context import BatchScope, current_scope
from extensions.main_fan.batch_models import BatchRequest, MinePlan
from extensions.main_fan.report_api import ReportClient, mock_report_client
from runtime import execution
from runtime.python_sandbox import digest, workspace_file


@contextmanager
def batch_lock(folder):
    folder.mkdir(parents=True, exist_ok=True)
    with (folder / '.run.lock').open('a') as stream:
        try:
            fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise ValueError('该批次正在执行，不能重复启动') from None
        try:
            yield
        finally:
            fcntl.flock(stream, fcntl.LOCK_UN)


def exportable_files(directory, names):
    # 沙箱中的代码可创建链接；宿主导出时不跟随链接读取工作区外的文件。
    return [name for name in names if (Path(directory) / name).is_symlink()
            or not ((Path(directory) / name).is_file() or (Path(directory) / name).is_dir())]


def export_index(folder, state):
    rows = []
    for mine in state['mines']:
        for shaft in mine.get('shafts', []):
            base = mine['dept_id'] + '/' + shaft['shaft_id'] + '/'
            report = '<a href="' + base + '报告.md">查看报告</a>' if shaft['status'] == 'completed' else escape(shaft.get('error', '待处理'))
            rows.append('<tr><td>' + escape(mine['mine_name']) + '</td><td>' + escape(shaft['shaft_name']) +
                '</td><td>' + escape(shaft.get('result_code', shaft['status'])) + '</td><td>' + report + '</td></tr>')
        if not mine.get('shafts'):
            rows.append('<tr><td>' + escape(mine['mine_name']) + '</td><td>矿级发现</td><td>' +
                escape(mine['status']) + '</td><td>' + escape(mine.get('error', '待处理')) + '</td></tr>')
    body = '''<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>主通风机本地批次</title>
<style>body{max-width:1100px;margin:48px auto;padding:0 24px;font:16px/1.7 system-ui;color:#17303b;background:#f5f7f8}
table{border-collapse:collapse;width:100%;background:white}td,th{padding:12px;border-bottom:1px solid #dbe3e5;text-align:left}
a{color:#076b78}.note{padding:16px;background:#fff2d7;border-radius:8px}code{font-size:14px}</style>
<h1>主通风机判识 · 本地预览</h1>'''
    body += '<p class="note">测定PDF使用既有测试原件的模拟接口，保留原件真实身份。结果仅写入本地文件，未推送。</p>'
    body += '<p>批次：<code>' + escape(state['batch_id']) + '</code>　状态：' + escape(state['status']) + '</p>'
    body += '<p>截止时刻：' + escape(state['cutoff']) + '，回查 ' + str(state['lookback_days']) + ' 天。</p>'
    verification = state.get('local_verification')
    if verification:
        body += '<p class="note"><strong>已交付本地数据核验。</strong> <a href="' + escape(verification['html_path'], quote=True) + '\">打开核验报告</a>。'
        body += 'Agent补充复核待模型权限恢复，旧AI初稿已标记未通过复核。此核验不作为正式业务结论。</p>'
    body += '<p><a href="batch.json">批次记录</a> · <a href="mine-catalog.json">冻结矿井清单</a></p>'
    body += '<table><thead><tr><th>煤矿</th><th>通风井/待确认对象</th><th>结果</th><th>文件</th></tr></thead><tbody>' + ''.join(rows) + '</tbody></table></html>'
    (folder / 'index.html').write_text(body, encoding='utf-8')


class BatchRunner:
    def __init__(self, output, pdf_path, *, cache_root=None, executor=None, progress=None, max_mines=None):
        self.output = Path(output).resolve()
        self.pdf_path = Path(pdf_path).resolve()
        self.cache_root = Path(cache_root or config.ROOT / 'data/main_fan/report_archives').resolve()
        self.executor = executor or execution.execute
        self.progress = progress
        self.max_mines = max_mines

    def persist(self, state):
        exploration._write(self.output / 'batch.json', state)
        export_index(self.output, state)

    async def run(self, request, *, snapshot=None):
        request = BatchRequest.model_validate(request)
        if len(set(request.dept_ids)) != len(request.dept_ids):
            raise ValueError('选择的煤矿部门不能重复')
        with batch_lock(self.output):
            return await self._run(request, snapshot)

    async def _run(self, request, snapshot):
        state_path = self.output / 'batch.json'
        if state_path.is_file():
            state = json.loads(state_path.read_text())
            if (state.get('requested_dept_ids', state['dept_ids']) != request.dept_ids
                    or state.get('all_producing_mines', False) != request.all_producing_mines
                    or state['lookback_days'] != request.lookback_days
                    or state['pdf_sha256'] != digest(self.pdf_path)
                    or state['user_name'] != request.user_name
                    or (request.cutoff and parse_time(request.cutoff).replace(microsecond=0).isoformat() != state['cutoff'])
                    or (request.snapshot_id and request.snapshot_id != state['snapshot_id'])):
                raise ValueError('恢复批次的矿井、时段、报告或用户与原记录不一致')
            snapshot = json.loads((self.output / 'mine-catalog.json').read_text())
            if request.reassess:
                for item in state['mines']:
                    if item.get('plan'):
                        item['status'] = 'pending'
                        for entry in item['shafts']:
                            folder = self.output / item['dept_id'] / entry['shaft_id']
                            if entry.get('analysis_id') and folder.is_dir():
                                revision = folder / 'revisions' / entry['analysis_id']
                                revision.mkdir(parents=True, exist_ok=True)
                                for name in ('报告.md', '结果.json', '来源与覆盖.json', '待审阅报文.json'):
                                    if (folder / name).is_file():
                                        shutil.copyfile(folder / name, revision / name)
                            entry['status'] = 'pending'
                state['status'] = 'running'
        else:
            snapshot = snapshot or (mine_catalog.load(request.snapshot_id) if request.snapshot_id
                                    else await mine_catalog.refresh())
            selected = {mine['dept_id']: mine for mine in snapshot['mines']}
            dept_ids = list(selected) if request.all_producing_mines else request.dept_ids
            if self.max_mines is not None and len(dept_ids) > self.max_mines:
                raise ValueError('本批次矿井数量超过配置上限，尚未启动任何矿井Agent')
            if any(dept not in selected for dept in dept_ids):
                raise ValueError('选择的部门不在本批次正常生产煤矿清单中')
            cutoff = (parse_time(request.cutoff)
                      if request.cutoff else datetime.now(ZoneInfo('Asia/Shanghai'))).replace(microsecond=0).isoformat()
            state = {'schema_version': 1, 'batch_id': uuid4().hex, 'snapshot_id': snapshot['snapshot_id'],
                'status': 'running', 'publish': False, 'delivery': 'local_files', 'pdf_mode': 'mock',
                'pdf_sha256': digest(self.pdf_path), 'cutoff': cutoff, 'lookback_days': request.lookback_days,
                'dept_ids': dept_ids, 'requested_dept_ids': request.dept_ids,
                'all_producing_mines': request.all_producing_mines, 'user_name': request.user_name,
                'mines': [{**{key: selected[dept][key] for key in ('dept_id', 'mine_id', 'mine_name')},
                           'status': 'pending', 'shafts': []} for dept in dept_ids]}
            exploration._write(self.output / 'mine-catalog.json', snapshot)
        if state['status'] != 'completed':
            state['status'] = 'running'
        self.persist(state)
        async with mock_report_client(self.pdf_path) as http:
            report_client = ReportClient(http)
            for item in state['mines']:
                if item['status'] == 'completed':
                    continue
                mine_dir = self.output / item['dept_id']
                mine_dir.mkdir(exist_ok=True)
                if not item.get('session_id'):
                    session = chat_store.create_session(request.user_name, 'full', item['mine_name'] + ' 主通风机月度判识')
                    item['session_id'] = session['id']
                elif not chat_store.get_session_for_user(item['session_id'], request.user_name):
                    raise ValueError('恢复批次的原矿级会话不存在，请使用原业务数据库')
                scope = BatchScope(state['batch_id'], item, request.user_name, item['session_id'],
                    mine_dir, self.cache_root, state['cutoff'], state['lookback_days'], report_client,
                    progress=self.progress)
                token = current_scope.set(scope)
                try:
                    item.update(status='running', error='')
                    self.persist(state)
                    if not item.get('plan'):
                        data = self.input_for(scope, 'discovery')
                        result = await self.call(scope, data, 'discovery')
                        plan = MinePlan.model_validate(result['assessment']['plan'])
                        item['plan'] = plan.model_dump(mode='json')
                        item['shafts'] = [{**shaft.model_dump(mode='json'), 'status': 'pending'} for shaft in plan.shafts]
                        exploration._write(mine_dir / 'discovery.json', result)
                        self.persist(state)
                    plan = MinePlan.model_validate(item['plan'])
                    for shaft, entry in zip(plan.shafts, item['shafts']):
                        if entry['status'] == 'completed':
                            continue
                        scope.shaft = shaft
                        entry.update(status='running', error='')
                        self.persist(state)
                        try:
                            result = await self.call(scope, self.input_for(scope, 'assessment'), shaft.shaft_id)
                            self.export_shaft(scope, result)
                            entry.update(status='completed', result_code=result['platform_payload']['judgment']['resultCode'],
                                analysis_id=result['analysis_id'], workflow_id=result['workflow_id'])
                        except asyncio.CancelledError:
                            entry['status'] = 'interrupted'
                            raise
                        except Exception as exc:
                            entry.update(status='failed', error=type(exc).__name__ + ': ' + str(exc)[:500])
                            exploration._write(mine_dir / shaft.shaft_id / 'failure.json', {'error': entry['error']})
                        self.persist(state)
                    item['status'] = 'completed' if all(row['status'] == 'completed' for row in item['shafts']) else 'partial'
                except asyncio.CancelledError:
                    item['status'] = 'interrupted'
                    state['status'] = 'interrupted'
                    self.persist(state)
                    raise
                except Exception as exc:
                    item.update(status='failed', error=type(exc).__name__ + ': ' + str(exc)[:500])
                    exploration._write(mine_dir / 'failure.json', {'error': item['error']})
                finally:
                    current_scope.reset(token)
                    # 包含真实请求、响应、PDF、脚本和证据，便于用户本地检查；没有服务端凭证。
                    source = session_root(scope.user, scope.session)
                    for directory in ('workspace', 'python_runs'):
                        if (source / directory).is_dir():
                            shutil.copytree(source / directory, mine_dir / directory,
                                            dirs_exist_ok=True, ignore=exportable_files)
                    self.persist(state)
        state['status'] = 'completed' if all(row['status'] == 'completed' for row in state['mines']) else 'partial'
        state['finished_at'] = datetime.now(ZoneInfo('Asia/Shanghai')).isoformat()
        self.persist(state)
        return state

    @staticmethod
    def input_for(scope, phase):
        return {'mine_code': scope.mine['mine_id'], 'mine_name': scope.mine['mine_name'],
            'dept_id': scope.mine['dept_id'], 'use_gateway': True,
            'fan_pair': scope.shaft.fan_pair if scope.shaft else [],
            'task_description': '本地批次主通风机判识',
            'measurements': {'batch': {'phase': phase, 'batch_id': scope.batch_id,
                'cutoff': scope.cutoff, 'lookback_days': scope.lookback_days,
                'shaft_plan': scope.shaft.model_dump(mode='json') if scope.shaft else None}}}

    async def call(self, scope, data, phase):
        if self.progress:
            self.progress({'stage': phase, 'dept_id': scope.mine['dept_id'], 'session_id': scope.session})
        mode = data['measurements']['batch']['phase']
        task_id = hashlib.sha256((scope.batch_id + scope.mine['dept_id'] + mode + ':' + phase).encode()).hexdigest()[:32]
        prepared = exploration.prepare(data, scope.user, scope.session, trusted=data, job_id=task_id)
        scope.active_workflow_id = prepared['workflow_id']
        message = ('这是主通风机本地批次，先读取业务Skill及references/batch.md。完整输入已由宿主注入，'
            '调用prepare_main_fan_task后读取input.json和其中measurements.batch。'
            '本轮阶段为' + ('矿级发现：查询本矿两个实时列表并取齐分页，自主确认通风井和测点，最后submit_main_fan_mine_plan。'
                if mode == 'discovery' else
                '井级分析：fetch_main_fan_reports核对PDF及历史档案；query_main_fan_history按天查满开停和风量。'
                '首次PDF完整读取后archive_main_fan_report保存提取，复用档案时跳过PDF识别。'
                '新的档案契约要求page_sources和每页实际read_file记录，渲染不能代替阅读。合订PDF中的每份报告都要单独提取，不能只整理第一份。'
                '自编Python完成时序、稳定窗口分析和报告，最后submit_main_fan_report。') +
            '历史取数后先读取agent_data_path的观测状态区间、变化边界、异常和风量统计。'
            '各日原件在单个ZIP中，按manifest的archive成员用Python按需复核；不要逐份读取或打印数十万条采样。'
            '切换分析必须基于具体时间而不只是开停占比。在details.switch_analysis中交付status、method、events、limitations；'
            '每个event记录候选或确认状态、原机停机时刻、另一台启动时刻、间隔和测点证据。'
            '无法确认整机时仍逐条保留相关测点的状态变化时刻作为候选，说明限制，不能仅凭占比宣称已发现切换。'
            '重大设备历史已有真实样例data[].sbNodeCurve[]，其中dataTime为时刻、realData为原始读数。'
            '优先采用整理结果中的实际样本数和质量信息；这些嵌套数组有数据时不能称为“无历史数据”，全为0也须报告实际样本和数据有效性限制。'
            '结果仅存本地，不调用推送。模拟PDF保留原件真实身份，禁止移植为本矿正式报告。'
            '本地任务已获授权，身份或数据不足时提交受限结果，不请求人工确认。'
            '\n本轮已准备好的唯一workflow_id：' + prepared['workflow_id'] +
            '；请使用此ID，前一阶段的ID已经结束。输入文件：' + prepared['input_path'] +
            '\n矿井身份：' + json.dumps({key: data[key] for key in ('mine_code', 'mine_name', 'dept_id')}, ensure_ascii=False))
        result = await self.executor(agent_name='main_fan', user_name=scope.user, session_id=scope.session,
            message=message, kind='business', graph_config={'recursion_limit': config.get_int('AGENT_EXPLORATION_RECURSION_LIMIT', 480),
                'configurable': {'business_task_input': data, 'business_task_id': task_id}})
        exploration._write(scope.output / ('execution-' + phase + '.json'), {
            'status': result.status, 'run_id': result.run_id, 'error': result.error, 'output': result.output,
            'usage': result.usage, 'interrupts': getattr(result, 'interrupts', []),
            'structured_output': result.structured_output})
        if result.status != 'success' or not result.structured_output:
            raise ValueError('Agent未完成实际提交：' + (result.error or result.status))
        return result.structured_output

    @staticmethod
    def export_shaft(scope, result):
        folder = scope.output / scope.shaft.shaft_id
        folder.mkdir(parents=True, exist_ok=True)
        exploration._write(folder / '结果.json', result)
        (folder / '报告.md').write_text(result['result'], encoding='utf-8')
        payload = result.get('platform_payload')
        if payload:
            exploration._write(folder / '待审阅报文.json', payload)
        metadata = {'mine': {key: scope.mine[key] for key in ('dept_id', 'mine_id', 'mine_name')},
            'shaft': scope.shaft.model_dump(mode='json'), 'cutoff': scope.cutoff,
            'lookback_days': scope.lookback_days, 'publish': False, 'pdf_mode': 'mock',
            'history_coverage': result.get('assessment', {}).get('history_coverage'),
            'report_archives': result.get('assessment', {}).get('report_archives', [])}
        exploration._write(folder / '来源与覆盖.json', metadata)