diff --git a/agent/relay_llm.py b/agent/relay_llm.py index c6c6affa78c..85286c47dae 100644 --- a/agent/relay_llm.py +++ b/agent/relay_llm.py @@ -706,23 +706,7 @@ class AnthropicStreamAccumulator: def response(self, base: Any = None) -> Any: """Return the attribute-shaped response consumed by Hermes.""" assembled = self.finalize() - if base is not None and base.__class__.__module__ == "unittest.mock": - base_payload = {} - for key in ( - "id", - "type", - "role", - "model", - "content", - "stop_reason", - "stop_sequence", - "usage", - ): - value = getattr(base, key, None) - if value is not None and value.__class__.__module__ != "unittest.mock": - base_payload[key] = _jsonable(value) - else: - base_payload = _jsonable(base) + base_payload = _jsonable(base) if not isinstance(base_payload, dict): base_payload = {} content = assembled.pop("content", []) @@ -1072,25 +1056,25 @@ def _codec(relay: Any, metadata: dict[str, Any] | None) -> Any: def _jsonable(value: Any) -> Any: if value is None or isinstance(value, (str, int, float, bool)): return value - # Test doubles synthesize arbitrary callable attributes such as - # ``model_dump``. Treat them as opaque instead of recursively invoking an - # endless chain of child mocks. - if value.__class__.__module__ == "unittest.mock": - return str(value) if isinstance(value, dict): return {str(key): _jsonable(item) for key, item in value.items()} if isinstance(value, (list, tuple, set)): return [_jsonable(item) for item in value] - model_dump = getattr(value, "model_dump", None) + model_dump = getattr(type(value), "model_dump", None) if callable(model_dump): try: - return _jsonable(model_dump(mode="json")) + return _jsonable(value.model_dump(mode="json")) except Exception: pass try: - return _jsonable(vars(value)) + attributes = { + str(key): item + for key, item in vars(value).items() + if not str(key).startswith("_") + } except (TypeError, AttributeError): return str(value) + return _jsonable(attributes) if attributes else str(value) def _namespace(value: Any) -> Any: diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py index 32e4788a974..0266d21bd2f 100644 --- a/tests/agent/test_relay_llm.py +++ b/tests/agent/test_relay_llm.py @@ -281,6 +281,52 @@ def test_anthropic_stream_accumulator_merges_terminal_usage(): } +def test_anthropic_stream_accumulator_merges_plain_provider_object(): + accumulator = relay_llm.AnthropicStreamAccumulator() + accumulator.observe({ + "type": "message_start", + "message": { + "id": "message-1", + "type": "message", + "role": "assistant", + "model": "claude-test", + "usage": {"input_tokens": 10}, + }, + }) + accumulator.observe({ + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": "hello"}, + }) + + response = accumulator.response( + SimpleNamespace( + id="message-1", + type="message", + role="assistant", + model="claude-test", + content=[], + stop_reason=None, + usage={"input_tokens": 10}, + ) + ) + + assert response.id == "message-1" + assert response.content[0].text == "hello" + assert response.usage.input_tokens == 10 + + +def test_jsonable_does_not_probe_dynamic_attributes(): + class DynamicProviderObject: + def __getattr__(self, name): + raise AssertionError(f"unexpected dynamic attribute lookup: {name}") + + def __str__(self): + return "opaque-provider-object" + + assert relay_llm._jsonable(DynamicProviderObject()) == "opaque-provider-object" + + def test_non_stream_preserves_raw_provider_response_identity(relay_turn): _relay, _turn = relay_turn raw_response = SimpleNamespace(model="test-model", content="raw")