mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-23 08:55:44 +08:00
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
import pytest
|
|
|
|
from langbot_plugin.api.entities.builtin.provider.message import Message, MessageChunk
|
|
from langbot_plugin.api.entities.builtin.platform.message import File, MessageChain
|
|
from langbot.pkg.pipeline.process.stream_results import coalesce_stream_results
|
|
|
|
|
|
async def source(*items):
|
|
for item in items:
|
|
yield item
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminal_snapshot_is_delivered_once_with_attachments():
|
|
text = '<think>reasoning</think>answer'
|
|
completed = Message(role='assistant', content=text)
|
|
completed.attachments = MessageChain([File(url='https://example.com/report.txt', name='report.txt')])
|
|
results = [
|
|
item
|
|
async for item in coalesce_stream_results(
|
|
source(
|
|
MessageChunk(role='assistant', content='answer', all_content=text, is_final=True),
|
|
completed,
|
|
)
|
|
)
|
|
]
|
|
assert results == [completed]
|
|
assert results[0].attachments is completed.attachments
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('with_snapshot', [False, True])
|
|
async def test_progressive_chunks_are_not_buffered(with_snapshot):
|
|
first = MessageChunk(role='assistant', content='hel', all_content='hel')
|
|
final = MessageChunk(role='assistant', content='lo', all_content='hello', is_final=True)
|
|
completed = Message(role='assistant', content='hello')
|
|
stream = coalesce_stream_results(source(first, final, *([completed] if with_snapshot else [])))
|
|
assert await anext(stream) == first
|
|
assert [item async for item in stream] == [completed if with_snapshot else final]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_only_and_distinct_messages_are_preserved():
|
|
final = MessageChunk(role='assistant', content='first', is_final=True)
|
|
other = Message(role='assistant', content='second')
|
|
results = [item async for item in coalesce_stream_results(source(final, other))]
|
|
assert [item.content for item in results] == ['first', 'second']
|
|
assert [item async for item in coalesce_stream_results(source(other))] == [other]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_blank_terminal_delta_uses_completed_snapshot():
|
|
completed = Message(role='assistant', content='answer')
|
|
results = [
|
|
item
|
|
async for item in coalesce_stream_results(
|
|
source(
|
|
MessageChunk(role='assistant', is_final=True),
|
|
completed,
|
|
)
|
|
)
|
|
]
|
|
assert results == [completed]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failures_propagate_without_successful_final():
|
|
async def failing():
|
|
yield MessageChunk(role='assistant', content='partial')
|
|
raise RuntimeError('provider failed')
|
|
|
|
stream = coalesce_stream_results(failing())
|
|
assert (await anext(stream)).content == 'partial'
|
|
with pytest.raises(RuntimeError, match='provider failed'):
|
|
await anext(stream)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cumulative_stream_ends_with_last_snapshot_even_when_final_is_blank():
|
|
results = [
|
|
item
|
|
async for item in coalesce_stream_results(
|
|
source(
|
|
MessageChunk(role='assistant', content='hel'),
|
|
MessageChunk(role='assistant', content='hello'),
|
|
MessageChunk(role='assistant', is_final=True),
|
|
)
|
|
)
|
|
]
|
|
assert [item.all_content for item in results] == ['hel', 'hello', 'hello']
|
|
assert results[-1].is_final
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_round_completion_keeps_one_transcript_and_attachments():
|
|
first = '<think>plan</think>I will create a file.'
|
|
last = '<think>file created</think>Here is your file.'
|
|
completed = Message(role='assistant', content=last)
|
|
completed.attachments = MessageChain([File(name='report.txt', base64='aGVsbG8=')])
|
|
results = [
|
|
item
|
|
async for item in coalesce_stream_results(
|
|
source(
|
|
MessageChunk(role='assistant', content=first, is_final=False),
|
|
MessageChunk(role='assistant', content=first + last, is_final=True),
|
|
completed,
|
|
)
|
|
)
|
|
]
|
|
assert len(results) == 2
|
|
assert results[-1].content == first + last
|
|
assert results[-1].attachments is completed.attachments
|