mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-30 04:15:45 +08:00
48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langbot_plugin.api.entities.builtin.provider.message import FunctionCall, MessageChunk, ToolCall
|
|
|
|
from pkg.model_calling import StreamingModelCaller
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fast_stream_coalesces_without_losing_text_reasoning_or_arguments(monkeypatch):
|
|
monkeypatch.setattr("pkg.model_calling.time.monotonic", lambda: 1.0)
|
|
|
|
async def provider(**kwargs):
|
|
for i in range(1000):
|
|
yield MessageChunk(
|
|
role="assistant",
|
|
content="字",
|
|
provider_specific_fields={"reasoning_content": "想"},
|
|
tool_calls=[ToolCall(id="call", type="function", function=FunctionCall(name="tool", arguments="a"))],
|
|
is_final=i == 999,
|
|
)
|
|
|
|
caller = StreamingModelCaller(SimpleNamespace(invoke_llm_stream=provider), ["model"], [])
|
|
values = [chunk async for chunk, _ in caller.stream() if chunk.content]
|
|
assert len(values) == 2
|
|
assert values[-1].content == "字" * 1000
|
|
assert values[-1].provider_specific_fields["reasoning_content"] == "想" * 1000
|
|
assert values[-1].tool_calls[0].function.arguments == "a" * 1000
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_after_first_snapshot_closes_underlying_provider():
|
|
closed = asyncio.Event()
|
|
|
|
async def provider(**kwargs):
|
|
try:
|
|
yield MessageChunk(role="assistant", content="first")
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
closed.set()
|
|
|
|
caller = StreamingModelCaller(SimpleNamespace(invoke_llm_stream=provider), ["model"], [])
|
|
stream = caller.stream()
|
|
await anext(stream)
|
|
await stream.aclose()
|
|
assert closed.is_set()
|