vent_agent_2.0/extensions/main_fan/batch.py

"""固定监测预处理、PDF 提取和单次 Qwen 研判，默认只在本地保存。"""
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
import time
from uuid import uuid4
from zoneinfo import ZoneInfo

import config
from agents.workspace import session_root
from db import chat_store, main_fan_store, business_store, db_conn
from extensions.main_fan import data_export, discovery, exploration, mine_catalog, preparation, workflow
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.crawler_reports import open_report_source, report_source
from obs import bind_log_context, unbind_log_context, get_logger

logger = get_logger(__name__)


@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>'''
    source_note = '测定PDF使用测试原件模拟接口。' if state['pdf_mode'] == 'mock' else '测定PDF来自资料采集服务及只读挂载目录。'
    delivery_note = '结果先保存在本地，符合提交条件的结果进入推送队列；推送状态见业务任务记录。' if state['publish'] else '结果仅保存在本地，推送关闭。'
    body += '<p class="note">' + source_note + delivery_note + '</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=None, *, cache_root=None, executor=None, progress=None, max_mines=None):
        self.output = Path(output).resolve()
        self.pdf_path = Path(pdf_path).resolve() if pdf_path is not None else None
        self.source = report_source(self.pdf_path)
        self.publish = workflow.publish_enabled()
        self.cache_root = Path(cache_root or config.ROOT / 'data/main_fan/report_archives').resolve()
        self.executor = executor or workflow.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):
            started = time.monotonic()
            token = bind_log_context(batch_output=str(self.output))
            try:
                logger.info('主通风机批次启动', extra={'kv': {'event': 'batch.started',
                    'report_mode': self.source['mode'], 'publish': self.publish, 'lookback_days': request.lookback_days}})
                state = await self._run(request, snapshot)
                logger.info('主通风机批次结束', extra={'kv': {'event': 'batch.finished',
                    'batch_id': state['batch_id'], 'status': state['status'],
                    'duration_ms': round((time.monotonic() - started) * 1000)}})
                return state
            except Exception:
                logger.exception('主通风机批次失败', extra={'kv': {'event': 'batch.failed'}})
                raise
            finally:
                unbind_log_context(token)

    async def _run(self, request, snapshot):
        state_path = self.output / 'batch.json'
        if state_path.is_file():
            state = json.loads(state_path.read_text())
            old_execution = state.get('execution_mode') != workflow.VERSION
            if old_execution and not request.reassess:
                raise ValueError('这是旧执行版本批次，请新建工作流批次，或显式 reassess 重新研判并保留旧结果')
            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.get('report_source', {'mode': 'mock', 'sha256': state.get('pdf_sha256')}) != self.source
                    or state.get('publish', False) != self.publish
                    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',
                                             '推送记录.json', '给Agent的数据.json', '给Qwen的数据.json',
                                             'PDF提取结果.json', 'PDF筛选记录.json', 'Qwen原始响应.json',
                                             'Qwen最终响应.json', 'JSON格式修复记录.json',
                                             '监测规则判定.json', '最终结论核对.json', '研判数据.zip', '数据清单.json'):
                                    if (folder / name).is_file():
                                        shutil.copyfile(folder / name, revision / name)
                            entry['status'] = 'pending'
                            entry.pop('analysis_saved', None)
                            entry['assessment_revision'] = uuid4().hex
                    if old_execution:
                        # 新版本修正了编号识别；旧 plan 的风机归属不能直接沿用。
                        item.pop('plan', None)
                        item['shafts'] = []
                        item['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('本批次矿井数量超过配置上限，尚未启动任何矿井研判')
            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': self.publish, 'execution_mode': workflow.VERSION,
                'delivery': 'local_files_and_push_queue' if self.publish else 'local_files', 'pdf_mode': self.source['mode'],
                'report_source': self.source, 'pdf_sha256': self.source.get('sha256'),
                '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'
        state['data_preparation'] = 'fixed_code'
        state['execution_mode'] = workflow.VERSION
        self.persist(state)
        async with open_report_source(self.source, self.pdf_path) as report_client:
            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, report_mode=self.source['mode'], publish=self.publish)
                token = current_scope.set(scope)
                log_token = bind_log_context(batch_id=scope.batch_id, dept_id=item['dept_id'],
                    mine_name=item['mine_name'], session_id=scope.session)
                mine_started = time.monotonic()
                try:
                    logger.info('矿井处理开始', extra={'kv': {'event': 'batch.mine.started'}})
                    item.update(status='running', error='')
                    self.persist(state)
                    if not item.get('plan'):
                        data = self.input_for(scope, 'discovery')
                        _, prepared = self.prepare_workflow(scope, data, 'discovery')
                        plan = await discovery.discover(prepared['workflow_id'], scope.user, scope.session)
                        main_fan_store.update_workflow(prepared['workflow_id'], status='completed')
                        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', {'mode': 'fixed_code', 'plan': item['plan']})
                        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
                        scope.assessment_revision = entry.get('assessment_revision', '')
                        shaft_log = bind_log_context(shaft_id=shaft.shaft_id)
                        try:
                            entry.update(status='running', error='')
                            self.persist(state)
                            folder = mine_dir / shaft.shaft_id
                            # 推送入队失败后恢复时直接使用已落地的结果，避免重复调用模型。
                            if entry.get('analysis_saved') and (folder / '结果.json').is_file():
                                result = json.loads((folder / '结果.json').read_text(encoding='utf-8'))
                            else:
                                result = await self.call(scope, self.input_for(scope, 'assessment'), shaft.shaft_id)
                                self.export_shaft(scope, result)
                                entry['analysis_saved'] = True
                                self.persist(state)
                            entry['data_archive'] = await asyncio.to_thread(data_export.export, scope, result)
                            logger.info('研判所用数据已与结果一同归档', extra={'kv': {
                                'event': 'batch.data.exported', **entry['data_archive']}})
                            self.queue_delivery(scope, result)
                            (folder / 'failure.json').unlink(missing_ok=True)
                            entry.update(status='completed', result_code=result['platform_payload']['judgment']['resultCode'],
                                analysis_id=result['analysis_id'], workflow_id=result['workflow_id'])
                            logger.info('通风井研判已保存', extra={'kv': {'event': 'batch.shaft.completed',
                                'analysis_id': result['analysis_id'], 'publish': self.publish}})
                        except asyncio.CancelledError:
                            entry['status'] = 'interrupted'
                            raise
                        except Exception as exc:
                            logger.exception('通风井处理失败', extra={'kv': {'event': 'batch.shaft.failed'}})
                            entry.update(status='failed', error=type(exc).__name__ + ': ' + str(exc)[:500])
                            exploration._write(mine_dir / shaft.shaft_id / 'failure.json', {'error': entry['error']})
                        finally:
                            unbind_log_context(shaft_log)
                        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:
                    logger.exception('矿井处理失败', extra={'kv': {'event': 'batch.mine.failed'}})
                    item.update(status='failed', error=type(exc).__name__ + ': ' + str(exc)[:500])
                    exploration._write(mine_dir / 'failure.json', {'error': item['error']})
                finally:
                    logger.info('矿井处理结束', extra={'kv': {'event': 'batch.mine.finished',
                        'status': item['status'], 'duration_ms': round((time.monotonic() - mine_started) * 1000)}})
                    current_scope.reset(token)
                    # 包含真实请求、响应、PDF、脚本和证据，便于用户本地检查；没有服务端凭证。
                    try:
                        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)
                    finally:
                        unbind_log_context(log_token)
        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 [],
            'document_cache_policy': config.get_str('MAIN_FAN_PDF_CACHE_POLICY', 'auto'),
            'task_description': '本地批次主通风机判识',
            'measurements': {'batch': {'phase': phase, 'batch_id': scope.batch_id, 'data_preparation': 'fixed_code',
                'execution_mode': workflow.VERSION,
                'cutoff': scope.cutoff, 'lookback_days': scope.lookback_days,
                'assessment_revision': scope.assessment_revision,
                'shaft_plan': scope.shaft.model_dump(mode='json') if scope.shaft else None}}}

    @staticmethod
    def prepare_workflow(scope, data, phase):
        mode = data['measurements']['batch']['phase']
        seed = scope.batch_id + scope.mine['dept_id'] + mode + ':' + phase + ':' + workflow.VERSION
        if scope.assessment_revision:
            seed += ':' + scope.assessment_revision
        task_id = hashlib.sha256(seed.encode()).hexdigest()[:32]
        if mode == 'assessment':
            # 模型结果已入库但本地导出失败时，继续交付同一原件，不重跑 Qwen。
            for row in db_conn.query(
                    "SELECT id FROM main_fan_workflows WHERE user_name=? AND session_id=? AND status='completed' ORDER BY created_at DESC LIMIT 40",
                    (scope.user, scope.session)):
                previous = main_fan_store.workflow(row['id'], scope.user, scope.session)
                if (previous['context'].get('job_id') == task_id
                        and (previous.get('result') or {}).get('assessment', {}).get('mode') == workflow.VERSION):
                    exploration.validate_inputs(previous, scope.user, scope.session)
                    scope.active_workflow_id = previous['id']
                    return task_id, {**exploration.manifest(previous, scope.user, scope.session), 'completed_result': True}
        prepared = exploration.prepare(data, scope.user, scope.session, trusted=data, job_id=task_id)
        scope.active_workflow_id = prepared['workflow_id']
        return task_id, prepared

    async def call(self, scope, data, phase):
        started = time.monotonic()
        logger.info('固定代码准备分析输入', extra={'kv': {'event': 'batch.preparation.started'}})
        if self.progress:
            self.progress({'stage': 'fixed_data_preparation', 'dept_id': scope.mine['dept_id'], 'session_id': scope.session})
        _, prepared = self.prepare_workflow(scope, data, phase)
        if not prepared.get('completed_result'):
            await preparation.assessment(prepared['workflow_id'], scope.user, scope.session)
        logger.info('监测输入已准备，开始执行固定研判工作流', extra={'kv': {'event': 'batch.workflow.started',
            'workflow_id': prepared['workflow_id'], 'duration_ms': round((time.monotonic() - started) * 1000)}})
        result = await self.executor(workflow_id=prepared['workflow_id'], user_name=scope.user,
            session_id=scope.session, progress=self.progress)
        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, 'diagnostics': getattr(result, 'diagnostics', {})})
        if result.status != 'success' or not result.structured_output:
            reason = result.error or ('缺少工作流校验后的正式结果'
                                      if result.status == 'success' else result.status)
            detail = getattr(result, 'diagnostics', {}).get('error', '')
            raise ValueError('工作流研判未完成：' + reason + ('；' + detail if detail else ''))
        logger.info('工作流研判完成', extra={'kv': {'event': 'batch.workflow.completed',
            'workflow_id': prepared['workflow_id'], 'run_id': result.run_id, 'usage': result.usage,
            'duration_ms': round((time.monotonic() - started) * 1000)}})
        return result.structured_output

    def queue_delivery(self, scope, result):
        folder = scope.output / scope.shaft.shaft_id
        receipt = {'enabled': self.publish, 'status': 'local_only', 'task_id': result['platform_payload']['taskId']}
        if self.publish:
            if result.get('status') != 'completed':
                receipt['status'] = 'skipped_incomplete'
            else:
                from extensions.result_platform.settings import ResultPlatformSettings
                settings = ResultPlatformSettings.from_config()
                job = business_store.record_completed_result('main_fan', scope.user, scope.session,
                    self.input_for(scope, 'assessment'), result, settings.url)
                receipt.update(status=job['push_status'], job_id=job['id'])
                logger.info('本地结果已进入持久化推送队列', extra={'kv': {'event': 'batch.delivery.queued',
                    'task_id': job['id'], 'push_status': job['push_status']}})
        exploration._write(folder / '推送记录.json', receipt)

    @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': scope.publish, 'pdf_mode': scope.report_mode,
            'history_coverage': result.get('assessment', {}).get('history_coverage'),
            'report_archives': result.get('assessment', {}).get('report_archives', [])}
        exploration._write(folder / '来源与覆盖.json', metadata)