292 lines
9.6 KiB
Python
292 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""webtest-agent - AI 网页测试智能体 (Flask 服务)"""
|
|
import os
|
|
import threading
|
|
|
|
from flask import Flask, request, jsonify, send_from_directory, abort
|
|
from flask_cors import CORS
|
|
|
|
import config
|
|
from db import (init_db, create_task, get_task, list_tasks, load_step_logs,
|
|
update_task, list_llm_configs, get_llm_config,
|
|
get_default_llm_config, create_llm_config, update_llm_config,
|
|
delete_llm_config, set_default_llm_config, llm_config_to_dict)
|
|
from agent import TaskRunner
|
|
|
|
app = Flask(__name__, static_folder='static', static_url_path='/static')
|
|
CORS(app)
|
|
|
|
init_db()
|
|
|
|
_runner_lock = threading.Lock()
|
|
_runners = {} # task_id -> TaskRunner
|
|
_semaphore = threading.Semaphore(config.MAX_CONCURRENT_TASKS)
|
|
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return send_from_directory('static', 'index.html')
|
|
|
|
|
|
@app.route('/health')
|
|
def health():
|
|
default_cfg = get_default_llm_config()
|
|
return jsonify({'status': 'ok', 'version': '1.0.0',
|
|
'llm_model': default_cfg.get('model') if default_cfg else None,
|
|
'llm_name': default_cfg.get('name') if default_cfg else None,
|
|
'llm_vision': bool(default_cfg.get('vision')) if default_cfg else None,
|
|
'concurrent': len(_runners)})
|
|
|
|
|
|
# ========== 大模型配置管理 ==========
|
|
|
|
|
|
@app.route('/api/llm-configs', methods=['GET'])
|
|
def api_list_llm_configs():
|
|
return jsonify({'configs': list_llm_configs()})
|
|
|
|
|
|
def _parse_llm_fields(data, partial=False):
|
|
"""校验并提取模型配置字段，返回 (fields, error)"""
|
|
fields = {}
|
|
for key, label in (('name', '名称'), ('base_url', 'Base URL'),
|
|
('api_key', 'API Key'), ('model', '模型名')):
|
|
val = (data.get(key) or '').strip()
|
|
if partial and key not in data:
|
|
continue
|
|
if not val:
|
|
return None, f'缺少{label}'
|
|
fields[key] = val
|
|
if 'vision' in data:
|
|
fields['vision'] = 1 if data.get('vision') else 0
|
|
if 'temperature' in data:
|
|
try:
|
|
fields['temperature'] = float(data.get('temperature'))
|
|
except (TypeError, ValueError):
|
|
return None, '温度必须是数字'
|
|
if 'timeout' in data:
|
|
try:
|
|
fields['timeout'] = max(10, min(int(data.get('timeout')), 600))
|
|
except (TypeError, ValueError):
|
|
return None, '超时必须是数字'
|
|
if 'is_default' in data:
|
|
fields['is_default'] = 1 if data.get('is_default') else 0
|
|
return fields, None
|
|
|
|
|
|
@app.route('/api/llm-configs', methods=['POST'])
|
|
def api_create_llm_config():
|
|
data = request.get_json(silent=True) or {}
|
|
fields, err = _parse_llm_fields(data)
|
|
if err:
|
|
return jsonify({'error': err}), 400
|
|
cid = create_llm_config(
|
|
name=fields['name'], base_url=fields['base_url'],
|
|
api_key=fields['api_key'], model=fields['model'],
|
|
vision=fields.get('vision', 0),
|
|
temperature=fields.get('temperature', 0.2),
|
|
timeout=fields.get('timeout', 120),
|
|
is_default=fields.get('is_default', False))
|
|
return jsonify({'ok': True, 'id': cid}), 201
|
|
|
|
|
|
@app.route('/api/llm-configs/<int:cid>', methods=['PUT'])
|
|
def api_update_llm_config(cid):
|
|
if not get_llm_config(cid):
|
|
abort(404)
|
|
data = request.get_json(silent=True) or {}
|
|
fields, err = _parse_llm_fields(data, partial=True)
|
|
if err:
|
|
return jsonify({'error': err}), 400
|
|
update_llm_config(cid, **fields)
|
|
return jsonify({'ok': True})
|
|
|
|
|
|
@app.route('/api/llm-configs/<int:cid>', methods=['DELETE'])
|
|
def api_delete_llm_config(cid):
|
|
row = get_llm_config(cid)
|
|
if not row:
|
|
abort(404)
|
|
if row.get('is_default'):
|
|
return jsonify({'error': '默认配置不能删除，请先设置其他配置为默认'}), 400
|
|
delete_llm_config(cid)
|
|
return jsonify({'ok': True})
|
|
|
|
|
|
@app.route('/api/llm-configs/<int:cid>/default', methods=['POST'])
|
|
def api_set_default_llm_config(cid):
|
|
if not get_llm_config(cid):
|
|
abort(404)
|
|
set_default_llm_config(cid)
|
|
return jsonify({'ok': True})
|
|
|
|
|
|
@app.route('/api/llm-configs/<int:cid>/test', methods=['POST'])
|
|
def api_test_llm_config(cid):
|
|
row = get_llm_config(cid)
|
|
if not row:
|
|
abort(404)
|
|
return _test_llm(llm_config_to_dict(row))
|
|
|
|
|
|
@app.route('/api/llm-configs/test-form', methods=['POST'])
|
|
def api_test_llm_config_form():
|
|
"""测试未保存的表单配置（前端弹窗里点「测试连接」用）"""
|
|
data = request.get_json(silent=True) or {}
|
|
base_url = (data.get('base_url') or '').strip()
|
|
api_key = (data.get('api_key') or '').strip()
|
|
model = (data.get('model') or '').strip()
|
|
if not base_url or not api_key or not model:
|
|
return jsonify({'error': '缺少 base_url / api_key / model'}), 400
|
|
return _test_llm({
|
|
'base_url': base_url,
|
|
'api_key': api_key,
|
|
'model': model,
|
|
'temperature': data.get('temperature', 0.2),
|
|
'timeout': data.get('timeout', 30),
|
|
})
|
|
|
|
|
|
def _test_llm(cfg):
|
|
from llm import chat, LLMError
|
|
try:
|
|
reply = chat(
|
|
[{'role': 'user', 'content': '请只回复两个字：正常'}],
|
|
cfg=cfg, max_tokens=20, timeout=30)
|
|
return jsonify({'ok': True, 'reply': (reply or '')[:200]})
|
|
except LLMError as e:
|
|
return jsonify({'ok': False, 'error': str(e)}), 400
|
|
|
|
|
|
@app.route('/api/tasks', methods=['POST'])
|
|
def api_create_task():
|
|
data = request.get_json(silent=True) or {}
|
|
url = (data.get('url') or '').strip()
|
|
goal = (data.get('goal') or '').strip()
|
|
if not url:
|
|
return jsonify({'error': '缺少 url 参数'}), 400
|
|
if not goal:
|
|
return jsonify({'error': '缺少 goal（测试目标）参数'}), 400
|
|
if not url.startswith(('http://', 'https://')):
|
|
url = 'https://' + url
|
|
max_steps = int(data.get('max_steps') or config.DEFAULT_MAX_STEPS)
|
|
timeout = int(data.get('timeout') or config.DEFAULT_TIMEOUT)
|
|
max_steps = max(1, min(max_steps, 100))
|
|
timeout = max(30, min(timeout, 3600))
|
|
|
|
# 选择分析模型：指定 id → 默认配置
|
|
cfg_row = None
|
|
if data.get('llm_config_id'):
|
|
cfg_row = get_llm_config(int(data.get('llm_config_id')))
|
|
if not cfg_row:
|
|
return jsonify({'error': '指定的模型配置不存在'}), 400
|
|
else:
|
|
cfg_row = get_default_llm_config()
|
|
if not cfg_row:
|
|
return jsonify({'error': '尚未配置任何大模型，请先在「大模型配置」中添加'}), 400
|
|
llm_cfg = llm_config_to_dict(cfg_row)
|
|
|
|
tid = create_task(url, goal, max_steps, timeout,
|
|
llm_config_id=cfg_row['id'],
|
|
llm_name=cfg_row['name'],
|
|
vision=cfg_row['vision'])
|
|
|
|
def _launch():
|
|
with _semaphore:
|
|
if get_task(tid) and get_task(tid).get('status') == 'stopped':
|
|
return
|
|
runner = TaskRunner(tid, url, goal, max_steps, timeout, llm_cfg)
|
|
with _runner_lock:
|
|
_runners[tid] = runner
|
|
runner.start()
|
|
runner.join()
|
|
with _runner_lock:
|
|
_runners.pop(tid, None)
|
|
|
|
threading.Thread(target=_launch, daemon=True).start()
|
|
return jsonify({'task_id': tid, 'status': 'queued',
|
|
'url': url, 'goal': goal,
|
|
'llm_name': cfg_row['name'],
|
|
'vision': bool(cfg_row['vision'])}), 202
|
|
|
|
|
|
@app.route('/api/tasks', methods=['GET'])
|
|
def api_list_tasks():
|
|
tasks = list_tasks(limit=50)
|
|
for t in tasks:
|
|
t['created'] = _fmt_time(t.get('created_at'))
|
|
t['finished'] = _fmt_time(t.get('finished_at'))
|
|
return jsonify({'tasks': tasks})
|
|
|
|
|
|
@app.route('/api/tasks/<tid>', methods=['GET'])
|
|
def api_get_task(tid):
|
|
t = get_task(tid)
|
|
if not t:
|
|
abort(404)
|
|
t['created'] = _fmt_time(t.get('created_at'))
|
|
t['finished'] = _fmt_time(t.get('finished_at'))
|
|
t['steps_log'] = load_step_logs(tid)
|
|
runner = _runners.get(tid)
|
|
if runner:
|
|
t['current'] = runner.current
|
|
return jsonify(t)
|
|
|
|
|
|
@app.route('/api/tasks/<tid>/stop', methods=['POST'])
|
|
def api_stop_task(tid):
|
|
runner = _runners.get(tid)
|
|
if runner:
|
|
runner.stop()
|
|
update_task(tid, status='stopped', result='stopped')
|
|
return jsonify({'ok': True, 'message': '停止请求已发送'})
|
|
t = get_task(tid)
|
|
if not t:
|
|
abort(404)
|
|
if t['status'] == 'queued':
|
|
update_task(tid, status='stopped', result='stopped')
|
|
return jsonify({'ok': True, 'message': '已取消排队任务'})
|
|
return jsonify({'ok': False, 'message': '任务不在运行中'})
|
|
|
|
|
|
@app.route('/api/tasks/<tid>/report')
|
|
def api_report(tid):
|
|
path = os.path.join(config.TASKS_DIR, tid, 'report.html')
|
|
if os.path.exists(path):
|
|
return send_from_directory(config.TASKS_DIR, f'{tid}/report.html')
|
|
abort(404)
|
|
|
|
|
|
@app.route('/api/tasks/<tid>/report.json')
|
|
def api_report_json(tid):
|
|
t = get_task(tid)
|
|
if not t:
|
|
abort(404)
|
|
return jsonify({
|
|
'task': t,
|
|
'steps': load_step_logs(tid),
|
|
'report_url': f'/api/tasks/{tid}/report',
|
|
})
|
|
|
|
|
|
@app.route('/api/tasks/<tid>/screenshot/<name>')
|
|
def api_screenshot(tid, name):
|
|
if not name or '..' in name or '/' in name:
|
|
abort(400)
|
|
path = os.path.join(config.TASKS_DIR, tid, name)
|
|
if os.path.exists(path):
|
|
return send_from_directory(config.TASKS_DIR, f'{tid}/{name}')
|
|
abort(404)
|
|
|
|
|
|
def _fmt_time(ts):
|
|
if not ts:
|
|
return None
|
|
import time
|
|
return time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(ts))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print(f'webtest-agent 启动: http://0.0.0.0:{config.PORT}')
|
|
app.run(host=config.HOST, port=config.PORT, threaded=True)
|