diff --git a/src/langbot/pkg/platform/sources/aiocqhttp.py b/src/langbot/pkg/platform/sources/aiocqhttp.py index 3cb55d89d..fc57a5661 100644 --- a/src/langbot/pkg/platform/sources/aiocqhttp.py +++ b/src/langbot/pkg/platform/sources/aiocqhttp.py @@ -16,6 +16,14 @@ from ...utils import image import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger +def _normalize_base64_payload(value: str) -> str: + if value.startswith('base64://'): + return value.removeprefix('base64://') + if value.startswith('data:') and ';base64,' in value: + return value.split(';base64,', 1)[1] + return value + + class AiocqhttpMessageConverter(abstract_platform_adapter.AbstractMessageConverter): @staticmethod async def yiri2target( @@ -35,7 +43,7 @@ class AiocqhttpMessageConverter(abstract_platform_adapter.AbstractMessageConvert elif type(msg) is platform_message.Image: arg = '' if msg.base64: - arg = msg.base64 + arg = _normalize_base64_payload(msg.base64) msg_list.append(aiocqhttp.MessageSegment.image(f'base64://{arg}')) elif msg.url: arg = msg.url @@ -50,7 +58,7 @@ class AiocqhttpMessageConverter(abstract_platform_adapter.AbstractMessageConvert elif type(msg) is platform_message.Voice: arg = '' if msg.base64: - arg = msg.base64 + arg = _normalize_base64_payload(msg.base64) msg_list.append(aiocqhttp.MessageSegment.record(f'base64://{arg}')) elif msg.url: arg = msg.url @@ -62,7 +70,10 @@ class AiocqhttpMessageConverter(abstract_platform_adapter.AbstractMessageConvert for node in msg.node_list: msg_list.extend((await AiocqhttpMessageConverter.yiri2target(node.message_chain))[0]) elif isinstance(msg, platform_message.File): - msg_list.append({'type': 'file', 'data': {'file': msg.url, 'name': msg.name}}) + file = msg.url or msg.path + if not file and msg.base64: + file = f'base64://{_normalize_base64_payload(msg.base64)}' + msg_list.append({'type': 'file', 'data': {'file': file, 'name': msg.name}}) elif isinstance(msg, platform_message.Face): if msg.face_type == 'face': msg_list.append(aiocqhttp.MessageSegment.face(msg.face_id)) @@ -433,9 +444,7 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter) elif isinstance(component, platform_message.Image): img_data = {} if component.base64: - b64 = component.base64 - if b64.startswith('data:'): - b64 = b64.split(',', 1)[-1] if ',' in b64 else b64 + b64 = _normalize_base64_payload(component.base64) img_data['file'] = f'base64://{b64}' elif component.url: img_data['file'] = component.url diff --git a/tests/unit_tests/platform/test_aiocqhttp_message_converter.py b/tests/unit_tests/platform/test_aiocqhttp_message_converter.py new file mode 100644 index 000000000..55f835c97 --- /dev/null +++ b/tests/unit_tests/platform/test_aiocqhttp_message_converter.py @@ -0,0 +1,105 @@ +import pytest + +import langbot_plugin.api.entities.builtin.platform.message as platform_message +from langbot.pkg.platform.sources.aiocqhttp import AiocqhttpAdapter, AiocqhttpMessageConverter + + +async def _convert_single(component: platform_message.MessageComponent): + chain = platform_message.MessageChain([component]) + message, _, _ = await AiocqhttpMessageConverter.yiri2target(chain) + return message[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('payload', 'expected'), + [ + ('data:image/jpeg;base64,raw-image', 'base64://raw-image'), + ('raw-image', 'base64://raw-image'), + ('base64://raw-image', 'base64://raw-image'), + ], +) +async def test_image_base64_payload_is_normalized(payload, expected): + segment = await _convert_single(platform_message.Image(base64=payload)) + + assert segment.type == 'image' + assert segment.data['file'] == expected + + +@pytest.mark.asyncio +async def test_voice_data_uri_base64_payload_is_normalized(): + segment = await _convert_single(platform_message.Voice(base64='data:audio/wav;base64,raw-voice')) + + assert segment.type == 'record' + assert segment.data['file'] == 'base64://raw-voice' + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('component', 'expected'), + [ + ( + platform_message.File(name='report.txt', base64='data:text/plain;base64,raw-file'), + {'file': 'base64://raw-file', 'name': 'report.txt'}, + ), + ( + platform_message.File(name='report.txt', base64='raw-file'), + {'file': 'base64://raw-file', 'name': 'report.txt'}, + ), + ( + platform_message.File(name='a.txt', url='http://example.com/a.txt'), + {'file': 'http://example.com/a.txt', 'name': 'a.txt'}, + ), + ( + platform_message.File(name='a.txt', path='/tmp/a.txt'), + {'file': '/tmp/a.txt', 'name': 'a.txt'}, + ), + ], +) +async def test_file_message_uses_available_file_source(component, expected): + segment = await _convert_single(component) + + assert segment.type == 'file' + assert segment.data == expected + + +@pytest.mark.asyncio +async def test_forward_image_base64_payload_is_normalized(): + forward = platform_message.Forward( + node_list=[ + platform_message.ForwardMessageNode( + sender_id='10001', + sender_name='Tester', + message_chain=platform_message.MessageChain( + [platform_message.Image(base64='data:image/png;base64,raw-forward-image')] + ), + ) + ] + ) + messages = [] + + class Logger: + async def info(self, _message): + return None + + async def error(self, _message): + return None + + class Bot: + async def call_action(self, action, **kwargs): + assert action == 'send_forward_msg' + messages.append(kwargs) + + platform = AiocqhttpAdapter.model_construct( + bot_account_id='10000', + config={}, + logger=Logger(), + bot=Bot(), + ) + + await platform._send_forward_message(1000, forward) + + assert messages[0]['messages'][0]['data']['content'][0] == { + 'type': 'image', + 'data': {'file': 'base64://raw-forward-image'}, + }