vent_agent_2.0/runtime/test_lab/service.py

"""工作台测试服务：选择输入、归档请求、复用业务队列、预览真实输出。"""
from copy import deepcopy
import hashlib
import json
from pathlib import Path
import shutil
import threading
import uuid

from agents.workspace import session_root
from db import business_store, chat_store, db_conn, test_lab_store
from extensions.registry import load_extensions
from runtime.business_tasks import get_controller, submit
from runtime.test_lab import files
from runtime.test_lab.presets import PRESETS, main_fan_input, preset
from runtime.test_lab.spreadsheets import result_tables, write_workbook
from runtime.test_lab.main_fan_simulation import MainFanSimulation, build_input

_submit_lock=threading.Lock()


def catalog(user):
    agents=[]
    for spec in load_extensions().business_agents.values():
        agents.append({'name':spec.name,'display_name':spec.display_name,'pipeline_available':spec.test_runner is not None,
            'presets':list(PRESETS) if spec.name=='main_fan' else []})
    return {'agents':agents,'assets':files.list_assets(user),
        'main_fan_simulation':MainFanSimulation().model_dump(mode='json')}


def get_preset(agent, identifier):
    load_extensions().business(agent)
    if agent!='main_fan':
        raise ValueError('该智能体还没有内置样例，可通过请求JSON输入测试数据')
    return {**preset(identifier),'input':main_fan_input(identifier)}


def _merge_fragment(data, fragment):
    for kind, rows in fragment.items():
        if not rows:
            continue
        if kind=='reports':
            data[kind]=rows
        elif kind=='fan_state_records':
            ids={r['fan_id'] for r in rows}
            data[kind]=[r for r in data.get(kind,[]) if r['fan_id'] not in ids]+rows
        elif kind=='operating_series':
            keys={(r['fan_id'],r['metric']) for r in rows}
            data[kind]=[r for r in data.get(kind,[]) if (r['fan_id'],r['metric']) not in keys]+rows
        elif kind=='devices':
            ids={r['device_code'] for r in rows}
            data[kind]=[r for r in data.get(kind,[]) if r['device_code'] not in ids]+rows
        else:
            data[kind]=[*data.get(kind,[]),*rows]


def start_run(request, user):
    with _submit_lock:
        return _start_run(request,user)


def _start_run(request,user):
    get_controller()
    spec=load_extensions().business(request.agent_name)
    simulation=request.main_fan_simulation
    if simulation is not None:
        if spec.name!='main_fan' or request.execution_mode!='agent':
            raise ValueError('三项输入测试必须使用主通风机完整智能体')
        if request.input or request.input_asset_id:
            raise ValueError('三项输入测试仅使用所选PDF及本次运行时序，请勿混入其他请求数据')
    if request.execution_mode=='pipeline' and spec.test_runner is None:
        raise ValueError('该智能体暂不支持业务流程模式，请选择完整智能体模式')
    source='test_pipeline' if request.execution_mode=='pipeline' else 'test_agent'
    if request.request_id:
        existing=db_conn.query_one('SELECT * FROM business_jobs WHERE id=?',
            (hashlib.sha256(f'{spec.name}\0{user}\0{source}\0{request.request_id}'.encode()).hexdigest()[:32],))
        if existing:
            previous=test_lab_store.run(existing['id'],user)
            snapshot=json.loads(previous['request_json'])
            if snapshot['submitted']!=request.model_dump(mode='json'):
                raise ValueError('同一请求标识已用于不同的测试参数')
            return public_run(existing['id'],user)
    data=build_input(simulation) if simulation is not None else deepcopy(request.input)
    if request.input_asset_id:
        input_path,_=files.owned_file(request.input_asset_id,user)
        if input_path.suffix.lower()!='.json':
            raise ValueError('请求数据来源必须是JSON文件')
    if data.get('materials'):
        raise ValueError('测试附件请从资料库选择，不要在请求JSON中手写文件路径')
    selected=[]
    slot_ids={key:getattr(simulation,key) for key in ('switch_asset_id','flow_asset_id')} if simulation else {}
    if simulation and len(set([*request.asset_ids,*filter(None,slot_ids.values())])) != len(request.asset_ids)+sum(bool(v) for v in slot_ids.values()):
        raise ValueError('每份测试资料只能选择一次，请分别选择报告、开停和风量文件')
    for identifier in [*request.asset_ids,*filter(None,slot_ids.values())]:
        path,description=files.owned_file(identifier,user)
        if path.suffix.lower() not in files.DOCUMENTS|{'.json'}:
            raise ValueError('当前主通风机仅接收PDF、图片和JSON作为研判输入；Excel文件可在线预览')
        description['sha256']=hashlib.sha256(path.read_bytes()).hexdigest()
        if simulation:
            slot=next((key for key,value in slot_ids.items() if value==identifier),'report')
            if slot=='report' and path.suffix.lower()!='.pdf':
                raise ValueError('测定报告输入请选择PDF文件')
            if slot!='report' and path.suffix.lower()!='.json':
                raise ValueError('开停曲线和风量数据上传文件必须为JSON')
            description['input_slot']=slot
        selected.append((path,description))
    if len({description['sha256'] for _,description in selected})!=len(selected):
        raise ValueError('所选资料中存在内容相同的文件，请只保留一份')
    if data.get('use_gateway') and selected:
        raise ValueError('网关取数样例会自行取得文件，请取消手动选择的资料，或改用PDF资料研判')
    combined={}
    for path,description in selected:
        if path.suffix.lower()=='.json':
            raw=files.json_value(path)
            try:
                fragment=files.series_fragment(raw,'test-lab://assets/'+description['id'])
            except (KeyError,TypeError) as exc:
                raise ValueError('运行时序JSON缺少明确的设备、时间或单位字段') from exc
            if fragment is None:
                allowed={'reports','devices','points','operating_series','fan_state_records','switch_records','confirmed_bindings'}
                if not raw or set(raw)-allowed or not all(isinstance(v,list) for v in raw.values()):
                    raise ValueError('完整请求JSON请先点击“用作请求”；也可直接选择开停、风量或业务事实片段JSON')
                fragment=raw
            if simulation:
                required='fan_state_records' if description['input_slot']=='switch_asset_id' else 'operating_series'
                allowed={required,'devices'} if required=='fan_state_records' else {required}
                if not fragment.get(required) or any(rows and key not in allowed for key,rows in fragment.items()):
                    raise ValueError('开停与风量文件须分别提供对应的时序，不能混入其他资料')
                expected_ids=set(data['fan_pair'])
                rows=fragment[required]
                if {row.get('fan_id') for row in rows}!=expected_ids or any(
                    row.get('mine_code')!=data['mine_code'] or row.get('mine_name')!=data['mine_name'] for row in rows):
                    raise ValueError('上传时序必须明确对应本次矿井及1号、2号通风机')
                if required=='operating_series' and (len(rows)!=2 or any(row.get('metric')!='flow' for row in rows)):
                    raise ValueError('风量文件须分别提供两台风机各一条风量时序')
                # 上传数据完整替换该输入槽，不能与参数生成的另一半序列拼接。
                data[required]=[]
            for kind,rows in fragment.items():
                combined.setdefault(kind,[]).extend(rows)
    _merge_fragment(data,combined)
    document_count=sum(path.suffix.lower() in files.DOCUMENTS for path,_ in selected)
    if simulation and not document_count:
        raise ValueError('请先选择或上传主备风机测定报告PDF')
    if request.scenario_id=='documents' and not document_count:
        raise ValueError('请至少选择一份PDF或图片资料')
    if document_count and any(r.get('evidence',{}).get('source','').startswith('test-lab://') for r in data.get('reports',[])):
        raise ValueError('使用PDF原件时请选择“PDF资料研判”，避免合成报告参数与原件混用')
    # 文件复制和提交在新会话下完成，任务只看到本次明确选择的资料。
    spec.input_schema.model_validate(data)
    session=chat_store.create_session(user,'full','智能体测试 · '+request.label[:100])
    business_store.bind_session(session['id'],spec.name)
    root=session_root(user,session['id'])
    attached=[]
    for path,description in selected:
        if path.suffix.lower() not in files.DOCUMENTS:
            continue
        destination=root/'workspace/uploads'/f"{description['id']}_{path.name}"
        destination.parent.mkdir(parents=True,exist_ok=True)
        shutil.copyfile(path,destination)
        digest=hashlib.sha256(destination.read_bytes()).hexdigest()
        attached.append({'path':'/workspace/uploads/'+destination.name,'filename':path.name,'sha256':digest})
    if attached:
        data['materials']=attached
    with db_conn.transaction():
        job=submit(spec.name,data,user_name=user,request_id=request.request_id,source=source,publish=False)
        business_store.update(job['id'],session_id=session['id'])
        snapshot={'submitted':request.model_dump(mode='json'),'input':data,
            'assets':[description for _,description in selected],
            'input_asset_id':request.input_asset_id,'publish':False}
        if simulation:
            # 每次是新的会话/文件根；仅归档本次原始输入，agent不会挂载历史workspace。
            input_dir=root/'workspace/test_inputs'
            input_dir.mkdir(parents=True,exist_ok=True)
            for name,key in (('主备风机开停曲线.json','fan_state_records'),('主备风机风量数据.json','operating_series')):
                (input_dir/name).write_text(json.dumps({key:data[key]},ensure_ascii=False,indent=2),encoding='utf-8')
            snapshot['isolation']={'fresh_session':True,'fresh_workspace':True,
                'document_cache_policy':'off','history_reused':False}
        if request.input_asset_id:
            _,description=files.owned_file(request.input_asset_id,user)
            snapshot['input_asset']=description
        test_lab_store.save_run(job,request.label,request.scenario_id,request.execution_mode,snapshot)
    return public_run(job['id'],user)


def _result(job):
    return json.loads(job['result_json']) if job.get('result_json') else None


def artifact_files(job):
    result=_result(job)
    if not result:
        return []
    root=session_root(job['user_name'],job['session_id']).resolve()
    artifacts=[]
    for item in result.get('artifacts',[]):
        virtual=item.get('path','')
        if not virtual.startswith('/workspace/'):
            continue
        path=(root/virtual.lstrip('/')).resolve()
        if not path.is_relative_to(root/'workspace') or not path.is_file() or path.suffix.lower() not in {'.xlsx','.csv'}:
            continue
        key='returned-'+hashlib.sha256(virtual.encode()).hexdigest()[:24]
        artifacts.append({'id':key,'filename':item.get('filename') or path.name,'origin':'returned','kind':'excel','path':path})
    artifacts.append({'id':'result_excel','filename':'研判结果.xlsx','origin':'generated','kind':'excel'})
    return artifacts


def artifact(identifier, key, user):
    test_lab_store.run(identifier,user)
    job=business_store.get(identifier)
    match=next((a for a in artifact_files(job) if a['id']==key),None)
    if match is None:
        raise ValueError('结果文件不存在')
    if key=='result_excel':
        path=session_root(user,job['session_id'])/'workspace/test_exports'/identifier/'研判结果.xlsx'
        if not path.is_file():
            path.parent.mkdir(parents=True,exist_ok=True)
            temporary=path.with_name(uuid.uuid4().hex+'.xlsx')
            try:
                write_workbook(temporary,result_tables(_result(job)))
                temporary.replace(path)
            finally:
                temporary.unlink(missing_ok=True)
        match['path']=path
    return match


def public_run(identifier,user):
    record=test_lab_store.run(identifier,user)
    job=business_store.get(identifier)
    result=_result(job)
    workflow=None
    if job.get('session_id'):
        workflow=db_conn.query_one('SELECT id,stages_json FROM main_fan_workflows WHERE session_id=? AND user_name=? ORDER BY created_at DESC LIMIT 1', (job['session_id'],user))
    artifacts=[]
    for item in artifact_files(job):
        artifacts.append({**{k:v for k,v in item.items() if k!='path'},
            'url':f"/api/test-lab/runs/{identifier}/artifacts/{item['id']}/content",
            'preview_url':f"/api/test-lab/runs/{identifier}/artifacts/{item['id']}/preview"})
    return {'id':identifier,'label':record['label'],'agent_name':record['agent_name'],
        'scenario_id':record['scenario_id'],'execution_mode':record['execution_mode'],
        'status':job['status'],'error':job['error'],'session_id':job['session_id'],
        'created_at':record['created_at'],'finished_at':job['finished_at'],
        'request':json.loads(record['request_json']),'result':result,'artifacts':artifacts,
        'stages':[{'name':s['name'],'status':s['status']} for s in json.loads(workflow['stages_json'])] if workflow else []}