135 lines
7.0 KiB
Python
135 lines
7.0 KiB
Python
"""Offline Dify attachment regression checks, including independent outlet failure."""
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from unittest import TestCase
|
|
|
|
import httpx
|
|
import model_gateway as gw
|
|
|
|
|
|
def outlet(kind='dify', **changes):
|
|
config = {'id': kind, 'name': kind, 'kind': kind, 'base_url': f'https://{kind}.invalid/v1',
|
|
'api_key': 'synthetic-key', 'model': 'audit', 'max_tokens': 300,
|
|
'temperature': .2, 'timeout_ms': 2000, 'max_inflight': 2}
|
|
config.update(changes)
|
|
return gw.Outlet(config=config, gate=asyncio.Semaphore(2))
|
|
|
|
|
|
def image(raw=b'audit-png', mime='image/png'):
|
|
return {'type': 'image_url', 'image_url': {'url': 'data:' + mime + ';base64,' + base64.b64encode(raw).decode()}}
|
|
|
|
|
|
def messages(*images):
|
|
return [{'role': 'user', 'content': [{'type': 'text', 'text': '请描述这些图片。'}, *images]}]
|
|
|
|
|
|
class GatewayFailureAudit(TestCase):
|
|
def run_request(self, handler, content, **kwargs):
|
|
async def run():
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
return await gw.call_outlet(client, kwargs.pop('outlet', outlet()), content, **kwargs)
|
|
return asyncio.run(run())
|
|
|
|
def test_plain_text_does_not_upload(self):
|
|
paths = []
|
|
def handler(request):
|
|
paths.append(request.url.path)
|
|
self.assertNotIn('files', json.loads(request.content))
|
|
return httpx.Response(200, json={'answer': '纯文字回复。'})
|
|
result = self.run_request(handler, [{'role': 'user', 'content': '在吗？'}])
|
|
self.assertEqual(result['text'], '纯文字回复。')
|
|
self.assertEqual(paths, ['/v1/chat-messages'])
|
|
|
|
def test_exact_endpoint_is_used_for_upload_root_and_chat(self):
|
|
paths = []
|
|
def handler(request):
|
|
paths.append(request.url.path)
|
|
return httpx.Response(200, json={'id': 'audit-upload'} if request.url.path.endswith('/files/upload') else {'answer': '成功。'})
|
|
result = self.run_request(handler, messages(image()), outlet=outlet(base_url='https://dify.invalid/custom/invoke', endpoint_mode='exact'))
|
|
self.assertEqual(result['text'], '成功。')
|
|
self.assertEqual(paths, ['/custom/invoke/files/upload', '/custom/invoke'])
|
|
|
|
def test_multiple_images_preserve_order_mime_bytes_and_same_user(self):
|
|
uploads, bodies = [], []
|
|
def handler(request):
|
|
if request.url.path.endswith('/files/upload'):
|
|
uploads.append(request.content)
|
|
self.assertIn(b'name="user"\r\n\r\nwechat-rpa\r\n', request.content)
|
|
return httpx.Response(200, json={'id': f'upload-{len(uploads)}'})
|
|
bodies.append(json.loads(request.content))
|
|
return httpx.Response(200, json={'answer': '两张图片已读取。'})
|
|
result = self.run_request(handler, messages(image(b'first-jpeg', 'image/jpeg'), image(b'second-webp', 'image/webp')))
|
|
self.assertEqual(result['error'], '')
|
|
self.assertEqual(len(uploads), 2)
|
|
self.assertIn(b'Content-Type: image/jpeg', uploads[0])
|
|
self.assertIn(b'first-jpeg', uploads[0])
|
|
self.assertIn(b'Content-Type: image/webp', uploads[1])
|
|
self.assertIn(b'second-webp', uploads[1])
|
|
self.assertEqual(bodies[0]['user'], 'wechat-rpa')
|
|
self.assertEqual([item['upload_file_id'] for item in bodies[0]['files']], ['upload-1', 'upload-2'])
|
|
self.assertNotIn('base64', bodies[0]['query'])
|
|
|
|
def test_old_turn_attachment_is_not_uploaded_for_a_text_followup(self):
|
|
seen = []
|
|
def handler(request):
|
|
seen.append(request.url.path)
|
|
return httpx.Response(200, json={'answer': '已说明。'})
|
|
content = [*messages(image()), {'role': 'assistant', 'content': '一张蓝色图片。'}, {'role': 'user', 'content': '谢谢'}]
|
|
result = self.run_request(handler, content)
|
|
self.assertEqual(result['error'], '')
|
|
self.assertEqual(seen, ['/v1/chat-messages'])
|
|
|
|
def test_same_inline_and_legacy_image_is_only_uploaded_once(self):
|
|
seen = []
|
|
def handler(request):
|
|
seen.append(request.url.path)
|
|
return httpx.Response(200, json={'id': 'one'} if request.url.path.endswith('/files/upload') else {'answer': '一张。'})
|
|
self.run_request(handler, messages(image()), image_b64=base64.b64encode(b'audit-png').decode())
|
|
self.assertEqual(seen.count('/v1/files/upload'), 1)
|
|
|
|
def _assert_upload_failure_is_local(self, upload_response):
|
|
async def run():
|
|
def handler(request):
|
|
if request.url.host == 'openai.invalid':
|
|
return httpx.Response(200, json={'choices': [{'message': {'content': '另一出口仍然成功。'}}]})
|
|
self.assertTrue(request.url.path.endswith('/files/upload'))
|
|
if isinstance(upload_response, type) and issubclass(upload_response, httpx.HTTPError):
|
|
raise upload_response('synthetic upload failure', request=request)
|
|
return upload_response
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
return await asyncio.gather(gw.call_outlet(client, outlet('dify'), messages(image())), gw.call_outlet(client, outlet('openai'), messages(image())))
|
|
results = asyncio.run(run())
|
|
self.assertEqual(results[0]['text'], '')
|
|
self.assertTrue(results[0]['error'])
|
|
self.assertEqual(results[1]['text'], '另一出口仍然成功。')
|
|
self.assertEqual(results[1]['error'], '')
|
|
|
|
def test_upload_invalid_json_is_local_to_dify_outlet(self):
|
|
self._assert_upload_failure_is_local(httpx.Response(200, text='invalid-json'))
|
|
|
|
def test_upload_non_object_json_is_local_to_dify_outlet(self):
|
|
for payload in ([{'id': 'wrong'}], 'wrong-shape', 7, None):
|
|
with self.subTest(payload=payload):
|
|
self._assert_upload_failure_is_local(httpx.Response(200, json=payload))
|
|
|
|
def test_upload_transport_failure_is_local_to_dify_outlet(self):
|
|
for error in (httpx.ReadTimeout, httpx.ConnectError):
|
|
with self.subTest(error=error.__name__):
|
|
self._assert_upload_failure_is_local(error)
|
|
|
|
def test_upload_http_failure_is_local_to_dify_outlet(self):
|
|
for status in (401, 413, 503):
|
|
with self.subTest(status=status):
|
|
self._assert_upload_failure_is_local(httpx.Response(status, text='synthetic failure'))
|
|
|
|
def test_remote_image_uses_file_reference_without_local_fetch(self):
|
|
requests = []
|
|
def handler(request):
|
|
requests.append(request)
|
|
body = json.loads(request.content)
|
|
self.assertEqual(body['files'], [{'type': 'image', 'transfer_method': 'remote_url', 'url': 'https://images.invalid/audit.png'}])
|
|
return httpx.Response(200, json={'answer': '已引用。'})
|
|
self.run_request(handler, messages({'type': 'image_url', 'image_url': {'url': 'https://images.invalid/audit.png'}}))
|
|
self.assertEqual(len(requests), 1)
|