149 lines
7.5 KiB
Python
149 lines
7.5 KiB
Python
"""Guard-only routing: configured vision model, no customer-answer fanout/judge."""
|
|
import asyncio
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest import TestCase, mock
|
|
|
|
import httpx
|
|
import admin_backend
|
|
import model_gateway as gw
|
|
|
|
|
|
def provider(key, kind='openai'):
|
|
return {'id': key, 'name': key.upper(), 'kind': kind,
|
|
'base_url': f'https://{key}.invalid/v1', 'api_key': 'synthetic', 'model': key}
|
|
|
|
|
|
class GuardPlanAudit(TestCase):
|
|
def catalog(self, vision='v'):
|
|
cat = gw.Catalog(Path('unused-audit.db'))
|
|
cat.outlets = {key: gw.Outlet(provider(key), asyncio.Semaphore(2)) for key in ('a', 'b', 'v', 'j')}
|
|
cat.roles = {'answer_ids': 'a,b', 'judge_id': 'j', 'vision_id': vision, 'judge_mode': 'arbitrate', 'version': 7}
|
|
return cat
|
|
|
|
def test_guard_selects_only_configured_vision_and_no_judge(self):
|
|
answers, judge, _, version = self.catalog().plan(purpose='guard')
|
|
self.assertEqual([a.id for a in answers], ['v'])
|
|
self.assertIsNone(judge)
|
|
self.assertEqual(version, 7)
|
|
|
|
def test_guard_without_usable_vision_falls_back_to_primary_only(self):
|
|
for missing in ('', 'missing'):
|
|
with self.subTest(vision=missing):
|
|
answers, judge, _, _ = self.catalog(missing).plan(purpose='guard')
|
|
self.assertEqual([a.id for a in answers], ['a'])
|
|
self.assertIsNone(judge)
|
|
|
|
def test_comfyui_is_not_a_guard_vision_model(self):
|
|
cat = self.catalog()
|
|
cat.outlets['v'].config['kind'] = 'comfyui'
|
|
self.assertEqual([a.id for a in cat.plan(purpose='guard')[0]], ['a'])
|
|
|
|
def test_guard_works_with_vision_even_without_answer_roles(self):
|
|
cat = self.catalog()
|
|
cat.roles['answer_ids'] = ''
|
|
self.assertEqual([a.id for a in cat.plan(purpose='guard')[0]], ['v'])
|
|
self.assertEqual(cat.plan()[0], [])
|
|
|
|
def test_normal_chat_plan_is_unchanged(self):
|
|
answers, judge, mode, _ = self.catalog().plan()
|
|
self.assertEqual([a.id for a in answers], ['a', 'b'])
|
|
self.assertEqual(judge.id, 'j')
|
|
self.assertEqual(mode, 'arbitrate')
|
|
|
|
|
|
class GuardEndpointAudit(TestCase):
|
|
def setUp(self):
|
|
self.directory = tempfile.TemporaryDirectory(prefix='guard-audit-')
|
|
self.addCleanup(self.directory.cleanup)
|
|
self.db_path = Path(self.directory.name) / 'audit.db'
|
|
self.db = admin_backend.Database(self.db_path)
|
|
self.db.initialize('InitialAdmin123')
|
|
self.uid = self.db.authenticate('admin', 'InitialAdmin123')['id']
|
|
for key in ('a', 'b', 'v', 'j'):
|
|
self.db.save_model_provider(provider(key), self.uid, '127.0.0.1')
|
|
self.roles = {'answer_ids': 'a,b', 'judge_id': 'j', 'vision_id': 'v', 'judge_mode': 'arbitrate'}
|
|
self.db.save_model_roles(self.roles, self.uid, '127.0.0.1')
|
|
self.token, _ = self.db.create_desktop_session(
|
|
{'user_id': 'guard-audit', 'sn': 'guard-audit', 'nickname': '合成测试'},
|
|
device_id='audit', device_name='audit', app_version='1', ip='127.0.0.1')
|
|
|
|
def request(self, purpose='guard', *, broken=False, image=False):
|
|
seen = []
|
|
async def run():
|
|
app = gw.create_app(self.db_path)
|
|
await app.gateway_startup()
|
|
await app.state.client.aclose()
|
|
def upstream(request):
|
|
seen.append(request)
|
|
if broken:
|
|
return httpx.Response(401, text='synthetic-auth-failure')
|
|
if request.url.path.endswith('/files/upload'):
|
|
return httpx.Response(200, json={'id': 'audit-file'})
|
|
if request.url.host == 'j.invalid':
|
|
text = '{"winner":"A","score":0.8,"risk":"low"}'
|
|
else:
|
|
text = '{"state":"chat_ready","action":"none"}'
|
|
if request.url.path.endswith('/chat-messages'):
|
|
return httpx.Response(200, json={'answer': text})
|
|
return httpx.Response(200, json={'choices': [{'message': {'content': text}}]})
|
|
app.state.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream))
|
|
content = [{'type': 'text', 'text': '判断聊天界面'}, {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,YXVkaXQ='}}] if image else '判断聊天界面'
|
|
body = {'purpose': purpose, 'customer_text': '判断聊天界面', 'messages': [{'role': 'user', 'content': content}], 'task_id': 'audit-guard'}
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://gateway') as client:
|
|
result = await client.post('/v1/answer', headers={'Authorization': 'Bearer ' + self.token}, json=body)
|
|
await app.gateway_log_queue.join()
|
|
return result
|
|
finally:
|
|
await app.gateway_shutdown()
|
|
await app.state.client.aclose()
|
|
with mock.patch('model_gateway.run_judge', wraps=gw.run_judge) as judge:
|
|
response = asyncio.run(run())
|
|
return response, seen, judge.call_count
|
|
|
|
def test_guard_calls_only_vision_and_never_customer_judge(self):
|
|
response, seen, judged = self.request(image=True)
|
|
self.assertEqual(response.status_code, 200, response.text)
|
|
self.assertEqual([r.url.host for r in seen], ['v.invalid'])
|
|
self.assertEqual(response.json()['chosen'], 'V')
|
|
self.assertEqual(len(response.json()['candidates']), 1)
|
|
self.assertFalse(response.json()['judge']['participated'])
|
|
self.assertEqual(judged, 0)
|
|
|
|
def test_guard_without_vision_keeps_legacy_primary_compatibility(self):
|
|
self.db.save_model_roles({**self.roles, 'vision_id': ''}, self.uid, '127.0.0.1')
|
|
response, seen, judged = self.request()
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual([r.url.host for r in seen], ['a.invalid'])
|
|
self.assertEqual(judged, 0)
|
|
|
|
def test_disabled_vision_uses_primary_compatibility(self):
|
|
self.db.save_model_provider({**provider('v'), 'enabled': False}, self.uid, '127.0.0.1')
|
|
response, seen, _ = self.request()
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assertEqual([r.url.host for r in seen], ['a.invalid'])
|
|
|
|
def test_guard_failure_does_not_silently_switch_to_chat_answers(self):
|
|
response, seen, judged = self.request(broken=True)
|
|
self.assertEqual(response.status_code, 502)
|
|
self.assertEqual([r.url.host for r in seen], ['v.invalid'])
|
|
self.assertEqual(judged, 0)
|
|
|
|
def test_configured_dify_vision_receives_uploaded_screenshot(self):
|
|
self.db.save_model_provider(provider('v', 'dify'), self.uid, '127.0.0.1')
|
|
response, seen, judged = self.request(image=True)
|
|
self.assertEqual(response.status_code, 200, response.text)
|
|
self.assertEqual([r.url.host for r in seen], ['v.invalid', 'v.invalid'])
|
|
self.assertEqual([r.url.path for r in seen], ['/v1/files/upload', '/v1/chat-messages'])
|
|
self.assertEqual(json.loads(seen[-1].content)['files'][0]['upload_file_id'], 'audit-file')
|
|
self.assertEqual(judged, 0)
|
|
|
|
def test_normal_image_chat_retains_answer_fanout_and_judge(self):
|
|
response, seen, judged = self.request(purpose='chat', image=True)
|
|
self.assertEqual(response.status_code, 200, response.text)
|
|
self.assertEqual(sorted(r.url.host for r in seen), ['a.invalid', 'b.invalid', 'j.invalid'])
|
|
self.assertEqual(len(response.json()['candidates']), 2)
|
|
self.assertEqual(judged, 1)
|