Files
wehub-resource-sync c3bf08ac8d
K8s Workspace Integration Tests / k8s-workspace-tests (push) Waiting to run
Pre-commit / run (ubuntu-latest) (push) Waiting to run
Python Unittest Coverage / test (macos-15, 3.11) (push) Waiting to run
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Waiting to run
Python Unittest Coverage / test (windows-latest, 3.11) (push) Waiting to run
Web UI / check (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 12:39:27 +08:00

470 lines
17 KiB
Python

# -*- coding: utf-8 -*-
"""Unit tests for :class:`agentscope.model.ChatModelBase.__call__` — the
retry / accumulation / interrupt wrapper around ``_call_api``."""
import asyncio
from unittest.async_case import IsolatedAsyncioTestCase
from utils import AnyString, MockModel
from agentscope.message import TextBlock, ThinkingBlock, ToolCallBlock, UserMsg
from agentscope.model import ChatResponse, FinishedReason
def _dump(chat_response: ChatResponse) -> dict:
"""Normalize a ``ChatResponse`` (a ``dict``-subclass with pydantic
blocks inside) into a plain dict suitable for
``assertDictEqual`` / ``assertListEqual`` comparison."""
d = dict(chat_response)
d["content"] = [b.model_dump(mode="json") for b in d["content"]]
if d["usage"] is not None:
d["usage"] = dict(d["usage"])
return d
def _expected(
content: list,
is_last: bool,
finished_reason: FinishedReason = FinishedReason.COMPLETED,
usage: dict | None = None,
) -> dict:
"""Build the expected serialized ``ChatResponse`` dict, with
``AnyString`` placeholders for auto-generated fields (id,
created_at)."""
return {
"content": content,
"is_last": is_last,
"id": AnyString(),
"created_at": AnyString(),
"type": "chat_response",
"usage": usage,
"finished_reason": finished_reason,
"metadata": {},
}
class ChatModelBaseCallTest(IsolatedAsyncioTestCase):
"""Test ``ChatModelBase.__call__`` end-to-end with a ``MockModel``.
Covers the four scenarios that ``__call__`` must handle:
1. Non-stream success — the underlying ``ChatResponse`` is returned
verbatim.
2. Non-stream ``CancelledError`` raised from inside ``_call_api`` —
converted to a ``ChatResponse`` with
``finished_reason=INTERRUPTED``.
3. Stream success — every delta is forwarded, followed by an
accumulated final ``ChatResponse`` with ``is_last=True`` and
``finished_reason=COMPLETED``.
4. Stream ``CancelledError`` raised while consuming the underlying
async generator — deltas produced so far are forwarded, followed
by an accumulated final ``ChatResponse`` with
``finished_reason=INTERRUPTED``.
"""
async def asyncSetUp(self) -> None:
"""The async setup method."""
self.model = MockModel(model="mock-model")
self.messages = [UserMsg(name="user", content="hi")]
# ------------------------------------------------------------------
# 1) non-stream success
# ------------------------------------------------------------------
async def test_non_stream_success(self) -> None:
"""Non-stream ``_call_api`` returns a ``ChatResponse``; the base
class must return it unchanged."""
response = ChatResponse(
content=[TextBlock(text="hello", id="t1")],
is_last=True,
)
self.model.set_responses([response])
result = await self.model(messages=self.messages)
self.assertIsInstance(result, ChatResponse)
self.assertDictEqual(
_dump(result),
_expected(
content=[
{"type": "text", "text": "hello", "id": "t1"},
],
is_last=True,
),
)
# ------------------------------------------------------------------
# 2) non-stream CancelledError raised from inside _call_api
# ------------------------------------------------------------------
async def test_non_stream_cancelled_error(self) -> None:
"""``CancelledError`` raised from inside a non-stream
``_call_api`` is translated into an empty ``ChatResponse`` with
``finished_reason=INTERRUPTED``."""
self.model.set_responses([asyncio.CancelledError()])
result = await self.model(messages=self.messages)
self.assertIsInstance(result, ChatResponse)
self.assertDictEqual(
_dump(result),
_expected(
content=[],
is_last=True,
finished_reason=FinishedReason.INTERRUPTED,
),
)
# ------------------------------------------------------------------
# 3) stream success
# ------------------------------------------------------------------
async def test_stream_success_with_final_accumulation(self) -> None:
"""A well-behaved stream of deltas is forwarded chunk-by-chunk;
the base class appends a final accumulated ``ChatResponse``
with ``is_last=True``."""
deltas = [
ChatResponse(
content=[TextBlock(text="hello", id="t1")],
is_last=False,
id="chunk-1",
),
ChatResponse(
content=[TextBlock(text=" world", id="t1")],
is_last=False,
id="chunk-2",
),
]
self.model.set_responses([deltas])
gen = await self.model(messages=self.messages)
collected = [_dump(c) async for c in gen]
self.assertListEqual(
collected,
[
# delta 1 — passed through verbatim
_expected(
content=[
{"type": "text", "text": "hello", "id": "t1"},
],
is_last=False,
),
# delta 2 — passed through verbatim
_expected(
content=[
{"type": "text", "text": " world", "id": "t1"},
],
is_last=False,
),
# accumulated final — synthesised by ChatModelBase
_expected(
content=[
{
"type": "text",
"text": "hello world",
"id": "t1",
},
],
is_last=True,
finished_reason=FinishedReason.COMPLETED,
),
],
)
# ------------------------------------------------------------------
# 4) stream CancelledError raised while consuming the generator
# ------------------------------------------------------------------
async def test_stream_cancelled_error_mid_consumption(self) -> None:
"""``CancelledError`` raised mid-stream is caught inside the
base class's ``_stream`` wrapper. Deltas produced before the
cancellation are still forwarded, and a final accumulated
``ChatResponse`` with ``finished_reason=INTERRUPTED`` is
appended."""
deltas = [
ChatResponse(
content=[TextBlock(text="hello", id="t1")],
is_last=False,
id="chunk-1",
),
asyncio.CancelledError(),
# anything after the exception must not be reached
ChatResponse(
content=[TextBlock(text=" world", id="t1")],
is_last=False,
id="chunk-2",
),
]
self.model.set_responses([deltas])
gen = await self.model(messages=self.messages)
collected = [_dump(c) async for c in gen]
self.assertListEqual(
collected,
[
# the one delta yielded before the cancellation
_expected(
content=[
{"type": "text", "text": "hello", "id": "t1"},
],
is_last=False,
),
# accumulated final with INTERRUPTED — content reflects
# only the deltas received before the cancellation
_expected(
content=[
{"type": "text", "text": "hello", "id": "t1"},
],
is_last=True,
finished_reason=FinishedReason.INTERRUPTED,
),
],
)
# ------------------------------------------------------------------
# 5) stream CancelledError — thinking + text content interrupted
# ------------------------------------------------------------------
async def test_stream_cancelled_error_with_thinking_and_text(
self,
) -> None:
"""Same as (4) but the underlying stream produces multiple
``ThinkingBlock`` and ``TextBlock`` deltas (across the same
block ids) before the cancellation. The accumulated final
response must merge the deltas by block id in their original
order."""
deltas = [
# thinking (part 1)
ChatResponse(
content=[
ThinkingBlock(thinking="let me ", id="think-1"),
],
is_last=False,
id="chunk-1",
),
# thinking (part 2, same block-id, plus signature)
ChatResponse(
content=[
ThinkingBlock(thinking="think...", id="think-1"),
],
is_last=False,
id="chunk-2",
),
# text (part 1, new block-id)
ChatResponse(
content=[TextBlock(text="hello", id="text-1")],
is_last=False,
id="chunk-3",
),
asyncio.CancelledError(),
# unreachable
ChatResponse(
content=[TextBlock(text=" world", id="text-1")],
is_last=False,
id="chunk-4",
),
]
self.model.set_responses([deltas])
gen = await self.model(messages=self.messages)
collected = [_dump(c) async for c in gen]
self.assertListEqual(
collected,
[
_expected(
content=[
{
"type": "thinking",
"thinking": "let me ",
"id": "think-1",
},
],
is_last=False,
),
_expected(
content=[
{
"type": "thinking",
"thinking": "think...",
"id": "think-1",
},
],
is_last=False,
),
_expected(
content=[
{"type": "text", "text": "hello", "id": "text-1"},
],
is_last=False,
),
# accumulated final — thinking merged, text merged, no
# trailing " world" (it came after the cancellation)
_expected(
content=[
{
"type": "thinking",
"thinking": "let me think...",
"id": "think-1",
},
{
"type": "text",
"text": "hello",
"id": "text-1",
},
],
is_last=True,
finished_reason=FinishedReason.INTERRUPTED,
),
],
)
# ------------------------------------------------------------------
# 6) stream CancelledError — thinking + text + tool_call interrupted
# ------------------------------------------------------------------
async def test_stream_cancelled_error_with_thinking_text_tool_call(
self,
) -> None:
"""Same as (5) but the stream also produces a
``ToolCallBlock`` whose ``input`` string is streamed across
multiple deltas. The cancellation happens after the partial
tool_call input; the accumulated final response must contain
the tool_call with the concatenated ``input`` received so far."""
deltas = [
# thinking
ChatResponse(
content=[
ThinkingBlock(thinking="planning...", id="think-1"),
],
is_last=False,
id="chunk-1",
),
# text
ChatResponse(
content=[
TextBlock(text="calling tool", id="text-1"),
],
is_last=False,
id="chunk-2",
),
# tool_call (part 1: name + partial input)
ChatResponse(
content=[
ToolCallBlock(
id="tool-1",
name="get_weather",
input='{"city":"',
),
],
is_last=False,
id="chunk-3",
),
# tool_call (part 2: input continuation, same block-id)
ChatResponse(
content=[
ToolCallBlock(
id="tool-1",
name="get_weather",
input='Beijing"',
),
],
is_last=False,
id="chunk-4",
),
asyncio.CancelledError(),
# unreachable — the closing "}" never arrives
ChatResponse(
content=[
ToolCallBlock(
id="tool-1",
name="get_weather",
input="}",
),
],
is_last=False,
id="chunk-5",
),
]
self.model.set_responses([deltas])
gen = await self.model(messages=self.messages)
collected = [_dump(c) async for c in gen]
self.assertListEqual(
collected,
[
_expected(
content=[
{
"type": "thinking",
"thinking": "planning...",
"id": "think-1",
},
],
is_last=False,
),
_expected(
content=[
{
"type": "text",
"text": "calling tool",
"id": "text-1",
},
],
is_last=False,
),
_expected(
content=[
{
"type": "tool_call",
"id": "tool-1",
"name": "get_weather",
"input": '{"city":"',
"state": "pending",
"suggested_rules": [],
},
],
is_last=False,
),
_expected(
content=[
{
"type": "tool_call",
"id": "tool-1",
"name": "get_weather",
"input": 'Beijing"',
"state": "pending",
"suggested_rules": [],
},
],
is_last=False,
),
# accumulated final — thinking / text preserved,
# tool_call.input concatenated to the partial JSON
# received before the cancellation (no closing "}")
_expected(
content=[
{
"type": "thinking",
"thinking": "planning...",
"id": "think-1",
},
{
"type": "text",
"text": "calling tool",
"id": "text-1",
},
{
"type": "tool_call",
"id": "tool-1",
"name": "get_weather",
"input": '{"city":"Beijing"',
"state": "pending",
"suggested_rules": [],
},
],
is_last=True,
finished_reason=FinishedReason.INTERRUPTED,
),
],
)
async def asyncTearDown(self) -> None:
"""The async teardown method."""