|
@@ -233,6 +233,45 @@ class IncrementalTimingClient:
|
|
|
return None
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+class CloseFailingProviderIterator:
|
|
|
|
|
+ def __init__(
|
|
|
|
|
+ self,
|
|
|
|
|
+ *,
|
|
|
|
|
+ next_error: BaseException | None = None,
|
|
|
|
|
+ items: tuple[StreamItem, ...] = (),
|
|
|
|
|
+ ) -> None:
|
|
|
|
|
+ self.next_error = next_error
|
|
|
|
|
+ self.items = list(items)
|
|
|
|
|
+
|
|
|
|
|
+ def __aiter__(self) -> "CloseFailingProviderIterator":
|
|
|
|
|
+ return self
|
|
|
|
|
+
|
|
|
|
|
+ async def __anext__(self) -> StreamItem:
|
|
|
|
|
+ if self.next_error is not None:
|
|
|
|
|
+ raise self.next_error
|
|
|
|
|
+ if self.items:
|
|
|
|
|
+ return self.items.pop(0)
|
|
|
|
|
+ raise StopAsyncIteration
|
|
|
|
|
+
|
|
|
|
|
+ async def aclose(self) -> None:
|
|
|
|
|
+ raise RuntimeError("close failed")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class IteratorTimingClient:
|
|
|
|
|
+ def __init__(self, iterator: AsyncIterator[StreamItem]) -> None:
|
|
|
|
|
+ self.iterator = iterator
|
|
|
|
|
+
|
|
|
|
|
+ def stream_chat(
|
|
|
|
|
+ self,
|
|
|
|
|
+ messages: list[ChatMessage],
|
|
|
|
|
+ tools: list[dict[str, Any]],
|
|
|
|
|
+ params: AgentParams,
|
|
|
|
|
+ tool_choice: dict[str, Any] | None = None,
|
|
|
|
|
+ ) -> AsyncIterator[StreamItem]:
|
|
|
|
|
+ del messages, tools, params, tool_choice
|
|
|
|
|
+ return self.iterator
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
class SelectiveUsageBenchmarkClient:
|
|
class SelectiveUsageBenchmarkClient:
|
|
|
def __init__(
|
|
def __init__(
|
|
|
self,
|
|
self,
|
|
@@ -523,6 +562,78 @@ async def test_timing_client_finishes_timing_when_stream_is_cancelled():
|
|
|
assert inner.closed is True
|
|
assert inner.closed is True
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_timing_client_preserves_provider_cancellation_when_close_fails():
|
|
|
|
|
+ iterator = CloseFailingProviderIterator(next_error=asyncio.CancelledError())
|
|
|
|
|
+ client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
|
|
|
|
|
+
|
|
|
|
|
+ with pytest.raises(asyncio.CancelledError):
|
|
|
|
|
+ async for _ in client.stream_chat(
|
|
|
|
|
+ messages=[],
|
|
|
|
|
+ tools=[],
|
|
|
|
|
+ params=AgentParams(model="benchmark-model"),
|
|
|
|
|
+ ):
|
|
|
|
|
+ pass
|
|
|
|
|
+
|
|
|
|
|
+ assert len(client.timings) == 1
|
|
|
|
|
+ assert client.timings[0].first_item_kind is None
|
|
|
|
|
+ assert client.timings[0].usage is None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_timing_client_preserves_provider_error_when_close_fails():
|
|
|
|
|
+ iterator = CloseFailingProviderIterator(
|
|
|
|
|
+ next_error=RuntimeError("provider failed")
|
|
|
|
|
+ )
|
|
|
|
|
+ client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
|
|
|
|
|
+
|
|
|
|
|
+ with pytest.raises(RuntimeError, match="provider failed"):
|
|
|
|
|
+ async for _ in client.stream_chat(
|
|
|
|
|
+ messages=[],
|
|
|
|
|
+ tools=[],
|
|
|
|
|
+ params=AgentParams(model="benchmark-model"),
|
|
|
|
|
+ ):
|
|
|
|
|
+ pass
|
|
|
|
|
+
|
|
|
|
|
+ assert len(client.timings) == 1
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_timing_client_propagates_close_error_after_normal_completion():
|
|
|
|
|
+ iterator = CloseFailingProviderIterator()
|
|
|
|
|
+ client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
|
|
|
|
|
+
|
|
|
|
|
+ with pytest.raises(RuntimeError, match="close failed"):
|
|
|
|
|
+ async for _ in client.stream_chat(
|
|
|
|
|
+ messages=[],
|
|
|
|
|
+ tools=[],
|
|
|
|
|
+ params=AgentParams(model="benchmark-model"),
|
|
|
|
|
+ ):
|
|
|
|
|
+ pass
|
|
|
|
|
+
|
|
|
|
|
+ assert len(client.timings) == 1
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_timing_client_propagates_close_error_when_consumer_closes_stream():
|
|
|
|
|
+ iterator = CloseFailingProviderIterator(
|
|
|
|
|
+ items=(StreamItem.raw_response_chunk({"provider": "first"}),)
|
|
|
|
|
+ )
|
|
|
|
|
+ client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
|
|
|
|
|
+ stream = client.stream_chat(
|
|
|
|
|
+ messages=[],
|
|
|
|
|
+ tools=[],
|
|
|
|
|
+ params=AgentParams(model="benchmark-model"),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ assert (await anext(stream)).kind == "raw_chunk"
|
|
|
|
|
+ with pytest.raises(RuntimeError, match="close failed"):
|
|
|
|
|
+ await stream.aclose()
|
|
|
|
|
+
|
|
|
|
|
+ assert len(client.timings) == 1
|
|
|
|
|
+ assert client.timings[0].first_item_kind == "raw_chunk"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.asyncio
|
|
|
async def test_timing_client_indexes_concurrent_calls_by_start_order():
|
|
async def test_timing_client_indexes_concurrent_calls_by_start_order():
|
|
|
inner = ConcurrentTimingClient()
|
|
inner = ConcurrentTimingClient()
|