vent_agent_2.0/extensions/main_fan/data_export.py

"""将单次研判的实际输入和执行证据归档到同一个结果目录。"""
import hashlib
import json
from pathlib import PurePosixPath
import re
import shutil
from uuid import uuid4
from zipfile import ZipFile, ZipInfo, ZIP_DEFLATED, ZIP_STORED

from agents.workspace import session_root
from db import main_fan_store
from extensions.main_fan.exploration import _write
from runtime.python_sandbox import digest, workspace_file


def export(scope, result):
    workflow = main_fan_store.workflow(result['workflow_id'], scope.user, scope.session)
    context = workflow['context']
    folder = scope.output / scope.shaft.shaft_id
    folder.mkdir(parents=True, exist_ok=True)
    files, skipped = {}, []

    def add(virtual, sha):
        if not isinstance(virtual, str) or not virtual.startswith('/workspace/') or '..' in PurePosixPath(virtual).parts:
            raise ValueError('研判数据归档包含非法工作区路径')
        if not isinstance(sha, str) or not re.fullmatch(r'[0-9a-f]{64}', sha):
            raise ValueError('研判数据归档缺少有效文件指纹')
        if virtual in files and files[virtual][1] != sha:
            raise ValueError('同一研判数据路径出现不同指纹')
        files[virtual] = (workspace_file(scope.user, scope.session, virtual), sha,
                          str(PurePosixPath(virtual).relative_to('/')))

    def derived(virtual, sha):
        source = workspace_file(scope.user, scope.session, virtual)
        # 多轮脚本可以覆盖或删除中间结果；原始输入始终严格核验，
        # 已不存在的中间版本不能用当前同名文件冒充，记录在清单中。
        if not source.is_file() or digest(source) != sha:
            skipped.append({'source': virtual, 'sha256': sha, 'reason': '中间输出已被覆盖或删除，执行时指纹见运行收据'})
        else:
            add(virtual, sha)

    for key in ('exports', 'query_exports'):
        for virtual, sha in context.get(key, {}).items():
            add(virtual, sha)
    for row in context['task_input'].get('materials', []):
        add(row['path'], row['sha256'])
    for virtual, row in context.get('page_reads', {}).items():
        derived(virtual, row['sha256'])
    root = session_root(scope.user, scope.session).resolve()
    outputs = {}
    for receipt in context.get('python_runs', []):
        run_id = receipt['run_id']
        if not re.fullmatch(r'[0-9a-f]{32}', run_id):
            raise ValueError('研判执行记录 ID 无效')
        for name in ('script.py', 'receipt.json'):
            path = root / 'python_runs' / run_id / name
            if path.is_symlink() or not path.resolve().is_relative_to(root / 'python_runs') or not path.is_file():
                raise ValueError('研判执行记录缺失或越界')
            key = 'python_runs/' + run_id + '/' + name
            files[key] = (path, digest(path), key)
        for output in receipt.get('outputs', []):
            outputs[output['path']] = output['sha256']
    for virtual, sha in outputs.items():
        derived(virtual, sha)

    manifest = {'workflow_id': result['workflow_id'], 'analysis_id': result['analysis_id'],
        'mine': {key: scope.mine[key] for key in ('dept_id', 'mine_id', 'mine_name')},
        'shaft_id': scope.shaft.shaft_id, 'cutoff': scope.cutoff, 'lookback_days': scope.lookback_days,
        'note': 'workspace/ 对应本会话 /workspace/；历史原始请求、响应和采样在 history-evidence.zip 内，按 manifest 查阅。',
        'files': [], 'unavailable_intermediates': skipped}
    temporary = folder / ('.研判数据-' + uuid4().hex + '.writing')
    try:
        with ZipFile(temporary, 'w', allowZip64=True) as archive:
            for virtual, (source, expected, member) in sorted(files.items()):
                if not source.is_file():
                    raise ValueError('研判数据缺失：' + virtual)
                info = ZipInfo(member)
                info.compress_type = ZIP_STORED if source.suffix.lower() == '.zip' else ZIP_DEFLATED
                sha, size = hashlib.sha256(), 0
                with source.open('rb') as incoming, archive.open(info, 'w', force_zip64=True) as outgoing:
                    while chunk := incoming.read(1024 * 1024):
                        sha.update(chunk)
                        size += len(chunk)
                        outgoing.write(chunk)
                if sha.hexdigest() != expected:
                    raise ValueError('研判数据内容与分析时的指纹不一致：' + virtual)
                manifest['files'].append({'source': virtual, 'path': member, 'sha256': expected, 'size': size})
            archive.writestr('数据清单.json', json.dumps(manifest, ensure_ascii=False, indent=2))
        sha = digest(temporary)
        target = folder / '研判数据.zip'
        temporary.replace(target)
    finally:
        temporary.unlink(missing_ok=True)
    _write(folder / '数据清单.json', {**manifest, 'archive': target.name, 'sha256': sha})
    agent_path = context.get('prepared_data', {}).get('agent_data_path')
    if agent_path:
        with ZipFile(target) as archive, archive.open(files[agent_path][2]) as source:
            with (folder / '给Agent的数据.json').open('wb') as output:
                shutil.copyfileobj(source, output)
    # 首次研判和最终采纳响应分别保存，格式修复不覆盖原始响应。
    workflow_files = context.get('assessment_workflow', {}).get('files', {})
    for original, display in (('qwen-input.json', '给Qwen的数据.json'),
                              ('pdf-extraction.json', 'PDF提取结果.json'),
                              ('pdf-page-selection.json', 'PDF筛选记录.json'),
                              ('qwen-response.json', 'Qwen原始响应.json'),
                              ('qwen-final-response.json', 'Qwen最终响应.json'),
                              ('json-repair-history.json', 'JSON格式修复记录.json'),
                              ('monitoring-assessment.json', '监测规则判定.json'),
                              ('decision-resolution.json', '最终结论核对.json')):
        virtual = workflow_files.get(original)
        if virtual:
            source = workspace_file(scope.user, scope.session, virtual)
            if virtual not in files or digest(source) != files[virtual][1]:
                raise ValueError('工作流导出文件与本轮指纹不一致：' + original)
            shutil.copyfile(source, folder / display)
    metadata = folder / '来源与覆盖.json'
    if metadata.is_file():
        value = json.loads(metadata.read_text(encoding='utf-8'))
        value['data_archive'] = {'path': target.name, 'sha256': sha, 'manifest': '数据清单.json'}
        _write(metadata, value)
    return {'path': target.name, 'sha256': sha, 'file_count': len(files), 'size': target.stat().st_size}