From 1e5eedaded7e66fc985b25a0bc3fb8958037e53a Mon Sep 17 00:00:00 2001 From: MFA-X-AI Date: Mon, 13 Jul 2026 15:25:44 +0800 Subject: [PATCH] fix(core): support async and sync generator methods in MethodProxy MethodProxy.__call__ unconditionally awaited the wrapped method, so any async generator method (e.g. LLMProtocol.generate_stream) called through the observability proxy raised TypeError. Dispatch on the method kind at call time and return a matching wrapper that emits CALL/EXCEPTION/RESULT around iteration, with RESULT carrying a chunk count. --- src/xaibo/core/exchange.py | 82 ++++++++++- tests/core/test_proxy_generator_methods.py | 160 +++++++++++++++++++++ 2 files changed, 238 insertions(+), 4 deletions(-) create mode 100644 tests/core/test_proxy_generator_methods.py diff --git a/src/xaibo/core/exchange.py b/src/xaibo/core/exchange.py index 8e915ec..b144ce8 100644 --- a/src/xaibo/core/exchange.py +++ b/src/xaibo/core/exchange.py @@ -1,3 +1,4 @@ +import inspect from itertools import chain from collections import defaultdict from typing import Type, Union, Any @@ -258,18 +259,29 @@ def _emit_event(self, event_type: EventType, result=None, arguments=None, except print("Exception during event handling") traceback.print_exc() - async def __call__(self, *args, **kwargs): + def __call__(self, *args, **kwargs): """Forward calls to the wrapped method. - + + Dispatches on the kind of the wrapped method so the proxy stays + transparent: coroutine methods return a coroutine to await, generator + methods return a wrapping generator to iterate. + Args: *args: Positional arguments to pass to wrapped method **kwargs: Keyword arguments to pass to wrapped method - + Returns: The result of calling the wrapped method """ + if inspect.isasyncgenfunction(self._method): + return self._call_async_generator(*args, **kwargs) + if inspect.isgeneratorfunction(self._method): + return self._call_generator(*args, **kwargs) + return self._call_async(*args, **kwargs) + + async def _call_async(self, *args, **kwargs): self._call_id += 1 - + # Emit call event self._emit_event( EventType.CALL, @@ -296,6 +308,68 @@ async def __call__(self, *args, **kwargs): return result + async def _call_async_generator(self, *args, **kwargs): + self._call_id += 1 + + # Emit call event + self._emit_event( + EventType.CALL, + arguments={"args": args, "kwargs": kwargs} + ) + + chunks = 0 + try: + async for chunk in self._method(*args, **kwargs): + chunks += 1 + yield chunk + except GeneratorExit: + # Consumer stopped iterating early; not an error + raise + except: + self._emit_event( + EventType.EXCEPTION, + exception=traceback.format_exc() + ) + traceback.print_exc() + raise + + # Emit result event + self._emit_event( + EventType.RESULT, + result={"stream": True, "chunks": chunks} + ) + + def _call_generator(self, *args, **kwargs): + self._call_id += 1 + + # Emit call event + self._emit_event( + EventType.CALL, + arguments={"args": args, "kwargs": kwargs} + ) + + chunks = 0 + try: + for chunk in self._method(*args, **kwargs): + chunks += 1 + yield chunk + except GeneratorExit: + # Consumer stopped iterating early; not an error + raise + except: + self._emit_event( + EventType.EXCEPTION, + exception=traceback.format_exc() + ) + traceback.print_exc() + raise + + # Emit result event + self._emit_event( + EventType.RESULT, + result={"stream": True, "chunks": chunks} + ) + def __repr__(self): return f"MethodProxy({self._method.__name__})" diff --git a/tests/core/test_proxy_generator_methods.py b/tests/core/test_proxy_generator_methods.py new file mode 100644 index 0000000..5030640 --- /dev/null +++ b/tests/core/test_proxy_generator_methods.py @@ -0,0 +1,160 @@ +import pytest + +from xaibo.core.exchange import Proxy +from xaibo.core.models.events import Event, EventType +from xaibo.primitives.modules.llm.mock import MockLLM +from xaibo.core.models.llm import LLMMessage + + +class DummyStreamer: + async def stream_method(self, count, prefix="chunk"): + for i in range(count): + yield f"{prefix}-{i}" + + async def failing_stream(self): + yield "first" + raise ValueError("stream broke") + + def sync_stream(self, count): + for i in range(count): + yield i + + async def regular_method(self): + return "hello" + + +@pytest.mark.asyncio +async def test_proxy_async_generator_streams_chunks(): + events = [] + + def event_handler(event: Event): + events.append(event) + + obj = DummyStreamer() + proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + chunks = [chunk async for chunk in proxy.stream_method(3, prefix="foo")] + assert chunks == ["foo-0", "foo-1", "foo-2"] + + # Should have generated 2 events (call and result) + assert len(events) == 2 + + call_event = events[0] + assert call_event.event_type == EventType.CALL + assert call_event.module_class == "DummyStreamer" + assert call_event.method_name == "stream_method" + assert call_event.arguments == {"args": (3,), "kwargs": {"prefix": "foo"}} + assert call_event.agent_id == "test-agent" + + result_event = events[1] + assert result_event.event_type == EventType.RESULT + assert result_event.method_name == "stream_method" + assert result_event.result == {"stream": True, "chunks": 3} + assert result_event.call_id == call_event.call_id + + +@pytest.mark.asyncio +async def test_proxy_async_generator_without_listeners(): + obj = DummyStreamer() + proxy = Proxy(obj) + + chunks = [chunk async for chunk in proxy.stream_method(2)] + assert chunks == ["chunk-0", "chunk-1"] + + +@pytest.mark.asyncio +async def test_proxy_async_generator_exception_mid_stream(): + events = [] + + def event_handler(event: Event): + events.append(event) + + obj = DummyStreamer() + proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + received = [] + with pytest.raises(ValueError, match="stream broke"): + async for chunk in proxy.failing_stream(): + received.append(chunk) + + assert received == ["first"] + + event_types = [e.event_type for e in events] + assert event_types == [EventType.CALL, EventType.EXCEPTION] + assert "stream broke" in events[1].exception + + +@pytest.mark.asyncio +async def test_proxy_async_generator_early_close(): + events = [] + + def event_handler(event: Event): + events.append(event) + + obj = DummyStreamer() + proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + gen = proxy.stream_method(10) + assert await anext(gen) == "chunk-0" + await gen.aclose() + + # Early close is not an error and the stream never completed + event_types = [e.event_type for e in events] + assert event_types == [EventType.CALL] + + +@pytest.mark.asyncio +async def test_proxy_sync_generator(): + events = [] + + def event_handler(event: Event): + events.append(event) + + obj = DummyStreamer() + proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + chunks = list(proxy.sync_stream(3)) + assert chunks == [0, 1, 2] + + event_types = [e.event_type for e in events] + assert event_types == [EventType.CALL, EventType.RESULT] + assert events[1].result == {"stream": True, "chunks": 3} + + +@pytest.mark.asyncio +async def test_proxy_regular_method_unchanged(): + events = [] + + def event_handler(event: Event): + events.append(event) + + obj = DummyStreamer() + proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + result = await proxy.regular_method() + assert result == "hello" + + event_types = [e.event_type for e in events] + assert event_types == [EventType.CALL, EventType.RESULT] + assert events[1].result == "hello" + + +@pytest.mark.asyncio +async def test_proxy_llm_generate_stream(): + """The flagship case: LLMProtocol.generate_stream through the observability proxy.""" + events = [] + + def event_handler(event: Event): + events.append(event) + + llm = MockLLM({"responses": [{"content": "Hello World"}], "streaming_chunk_size": 5}) + proxy = Proxy(llm, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module") + + messages = [LLMMessage.user("Hi")] + chunks = [chunk async for chunk in proxy.generate_stream(messages)] + + assert "".join(chunks) == "Hello World" + event_types = [e.event_type for e in events] + assert event_types == [EventType.CALL, EventType.RESULT] + assert events[1].result["stream"] is True + assert events[1].result["chunks"] == len(chunks)