Python: updated declarative samples and handling of non-pydantic response formats (#5022)

* updated declarative samples and handling of non-pydantic response formats

* fixed from comments

* update docstring
This commit is contained in:
Eduard van Valkenburg
2026-04-01 21:16:00 +02:00
committed by GitHub
Unverified
parent 6acab3d1d6
commit 519bb0cb2b
21 changed files with 370 additions and 90 deletions
@@ -1026,20 +1026,13 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
session_context=context["session_context"],
suppress_response_id=context["suppress_response_id"],
)
response_format = context["chat_options"].get("response_format")
if not (
response_format is not None and isinstance(response_format, type) and issubclass(response_format, BaseModel)
):
response_format = None
return AgentResponse(
messages=response.messages,
response_id=None if context["suppress_response_id"] else response.response_id,
created_at=response.created_at,
usage_details=response.usage_details,
value=response.value,
response_format=response_format,
response_format=context["chat_options"].get("response_format"),
continuation_token=response.continuation_token,
raw_representation=response,
additional_properties=response.additional_properties,
@@ -1125,10 +1118,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
response_format: Any | None = None,
) -> AgentResponse[Any]:
"""Finalize response updates into a single AgentResponse."""
output_format_type = response_format if isinstance(response_format, type) else None
return AgentResponse.from_updates( # pyright: ignore[reportUnknownVariableType]
updates,
output_format_type=output_format_type,
output_format_type=response_format,
)
@staticmethod
@@ -345,10 +345,9 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
response_format: Any | None = None,
) -> ChatResponse[Any]:
"""Finalize response updates into a single ChatResponse."""
output_format_type = response_format if isinstance(response_format, type) else None
return ChatResponse.from_updates( # pyright: ignore[reportUnknownVariableType]
updates,
output_format_type=output_format_type,
output_format_type=response_format,
)
def _build_response_stream(
@@ -2327,7 +2327,6 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
return _get_response()
response_format = mutable_options.get("response_format") if mutable_options else None
output_format_type: type[BaseModel] | None = response_format if isinstance(response_format, type) else None
stream_result_hooks: list[Callable[[ChatResponse], Any]] = []
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
@@ -2485,6 +2484,6 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]:
# Note: stream_result_hooks are already run via inner stream's get_final_response()
# We don't need to run them again here
return ChatResponse.from_updates(updates, output_format_type=output_format_type)
return ChatResponse.from_updates(updates, output_format_type=response_format)
return ResponseStream(_stream(), finalizer=_finalize)
+80 -28
View File
@@ -299,6 +299,7 @@ ToolModeT = TypeVar("ToolModeT", bound="ToolMode")
AgentResponseT = TypeVar("AgentResponseT", bound="AgentResponse")
ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None, covariant=True)
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
StructuredResponseFormat = type[BaseModel] | Mapping[str, Any] | None
CreatedAtT = str # Use a datetimeoffset type? Or a more specific type like datetime.datetime?
@@ -1949,6 +1950,24 @@ class ContinuationToken(TypedDict):
# endregion
def _parse_structured_response_value(text: str, response_format: Any | None) -> Any | None:
if response_format is None:
return None
if isinstance(response_format, type) and issubclass(response_format, BaseModel):
return response_format.model_validate_json(text)
if isinstance(response_format, Mapping):
try:
return json.loads(text)
except json.JSONDecodeError as exc:
raise ValueError(f"Response text is not valid JSON: {exc}") from exc
logger.warning(
"Unable to parse structured response value, use either a Pydantic model or a dict defining the schema, "
"received response_format type: %s",
type(response_format), # type: ignore[reportUnknownArgumentType]
)
return None
class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
"""Represents the response to a chat request.
@@ -2014,7 +2033,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
finish_reason: FinishReasonLiteral | FinishReason | None = None,
usage_details: UsageDetails | None = None,
value: ResponseModelT | None = None,
response_format: type[BaseModel] | None = None,
response_format: StructuredResponseFormat = None,
continuation_token: ContinuationToken | None = None,
additional_properties: dict[str, Any] | None = None,
raw_representation: Any | None = None,
@@ -2058,7 +2077,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
self.finish_reason = finish_reason
self.usage_details = usage_details
self._value: ResponseModelT | None = value
self._response_format: type[BaseModel] | None = response_format
self._response_format: StructuredResponseFormat = response_format
self._value_parsed: bool = value is not None
self.additional_properties = (
_restore_compaction_annotation_in_additional_properties(additional_properties) or {}
@@ -2087,6 +2106,15 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
output_format_type: type[ResponseModelBoundT],
) -> ChatResponse[ResponseModelBoundT]: ...
@overload
@classmethod
def from_updates(
cls: type[ChatResponse[Any]],
updates: Sequence[ChatResponseUpdate],
*,
output_format_type: Mapping[str, Any],
) -> ChatResponse[Any]: ...
@overload
@classmethod
def from_updates(
@@ -2101,7 +2129,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
cls: type[ChatResponseT],
updates: Sequence[ChatResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
output_format_type: StructuredResponseFormat = None,
) -> ChatResponseT:
"""Joins multiple updates into a single ChatResponse.
@@ -2124,10 +2152,10 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
updates: A sequence of ChatResponseUpdate objects to combine.
Keyword Args:
output_format_type: Optional Pydantic model type to parse the response text into structured data.
output_format_type: Optional Pydantic model type or JSON schema mapping used to parse the
response text into structured data.
"""
response_format = output_format_type if isinstance(output_format_type, type) else None
msg = cls(messages=[], response_format=response_format)
msg = cls(messages=[], response_format=output_format_type)
for update in updates:
_process_update(msg, update)
_finalize_response(msg)
@@ -2142,6 +2170,15 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
output_format_type: type[ResponseModelBoundT],
) -> ChatResponse[ResponseModelBoundT]: ...
@overload
@classmethod
async def from_update_generator(
cls: type[ChatResponse[Any]],
updates: AsyncIterable[ChatResponseUpdate],
*,
output_format_type: Mapping[str, Any],
) -> ChatResponse[Any]: ...
@overload
@classmethod
async def from_update_generator(
@@ -2156,7 +2193,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
cls: type[ChatResponseT],
updates: AsyncIterable[ChatResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
output_format_type: StructuredResponseFormat = None,
) -> ChatResponseT:
"""Joins multiple updates into a single ChatResponse.
@@ -2175,10 +2212,10 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
updates: An async iterable of ChatResponseUpdate objects to combine.
Keyword Args:
output_format_type: Optional Pydantic model type to parse the response text into structured data.
output_format_type: Optional Pydantic model type or JSON schema mapping used to parse the
response text into structured data.
"""
response_format = output_format_type if isinstance(output_format_type, type) else None
msg = cls(messages=[], response_format=response_format)
msg = cls(messages=[], response_format=output_format_type)
async for update in updates:
_process_update(msg, update)
_finalize_response(msg)
@@ -2198,15 +2235,12 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
Raises:
ValidationError: If the response text doesn't match the expected schema.
ValueError: If the response text is not valid JSON for a non-Pydantic structured format.
"""
if self._value_parsed:
return self._value
if (
self._response_format is not None
and isinstance(self._response_format, type)
and issubclass(self._response_format, BaseModel)
):
self._value = cast(ResponseModelT, self._response_format.model_validate_json(self.text))
if self._response_format is not None:
self._value = cast(ResponseModelT, _parse_structured_response_value(self.text, self._response_format))
self._value_parsed = True
return self._value
@@ -2397,7 +2431,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
created_at: CreatedAtT | None = None,
usage_details: UsageDetails | None = None,
value: ResponseModelT | None = None,
response_format: type[BaseModel] | None = None,
response_format: StructuredResponseFormat = None,
continuation_token: ContinuationToken | None = None,
raw_representation: Any | None = None,
additional_properties: dict[str, Any] | None = None,
@@ -2438,7 +2472,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
self.created_at = created_at
self.usage_details = usage_details
self._value: ResponseModelT | None = value
self._response_format: type[BaseModel] | None = response_format
self._response_format: type[BaseModel] | Mapping[str, Any] | None = response_format
self._value_parsed: bool = value is not None
self.additional_properties = (
_restore_compaction_annotation_in_additional_properties(additional_properties) or {}
@@ -2460,15 +2494,12 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
Raises:
ValidationError: If the response text doesn't match the expected schema.
ValueError: If the response text is not valid JSON for a non-Pydantic structured format.
"""
if self._value_parsed:
return self._value
if (
self._response_format is not None
and isinstance(self._response_format, type)
and issubclass(self._response_format, BaseModel)
):
self._value = cast(ResponseModelT, self._response_format.model_validate_json(self.text))
if self._response_format is not None:
self._value = cast(ResponseModelT, _parse_structured_response_value(self.text, self._response_format))
self._value_parsed = True
return self._value
@@ -2492,6 +2523,16 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
value: Any | None = None,
) -> AgentResponse[ResponseModelBoundT]: ...
@overload
@classmethod
def from_updates(
cls: type[AgentResponse[Any]],
updates: Sequence[AgentResponseUpdate],
*,
output_format_type: Mapping[str, Any],
value: Any | None = None,
) -> AgentResponse[Any]: ...
@overload
@classmethod
def from_updates(
@@ -2507,7 +2548,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
cls: type[AgentResponseT],
updates: Sequence[AgentResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
output_format_type: StructuredResponseFormat = None,
value: Any | None = None,
) -> AgentResponseT:
"""Joins multiple updates into a single AgentResponse.
@@ -2516,7 +2557,8 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
updates: A sequence of AgentResponseUpdate objects to combine.
Keyword Args:
output_format_type: Optional Pydantic model type to parse the response text into structured data.
output_format_type: Optional Pydantic model type or JSON schema mapping used to parse the
response text into structured data.
value: Optional pre-parsed structured output value to set directly on the response.
"""
msg = cls(messages=[], response_format=output_format_type, value=value)
@@ -2534,6 +2576,15 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
output_format_type: type[ResponseModelBoundT],
) -> AgentResponse[ResponseModelBoundT]: ...
@overload
@classmethod
async def from_update_generator(
cls: type[AgentResponse[Any]],
updates: AsyncIterable[AgentResponseUpdate],
*,
output_format_type: Mapping[str, Any],
) -> AgentResponse[Any]: ...
@overload
@classmethod
async def from_update_generator(
@@ -2548,7 +2599,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
cls: type[AgentResponseT],
updates: AsyncIterable[AgentResponseUpdate],
*,
output_format_type: type[BaseModel] | None = None,
output_format_type: StructuredResponseFormat = None,
) -> AgentResponseT:
"""Joins multiple updates into a single AgentResponse.
@@ -2556,7 +2607,8 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
updates: An async iterable of AgentResponseUpdate objects to combine.
Keyword Args:
output_format_type: Optional Pydantic model type to parse the response text into structured data
output_format_type: Optional Pydantic model type or JSON schema mapping used to parse the
response text into structured data.
"""
msg = cls(messages=[], response_format=output_format_type)
async for update in updates:
+2 -6
View File
@@ -127,9 +127,7 @@ class MockChatClient:
yield ChatResponseUpdate(contents=[Content.from_text("another update")], role="assistant")
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
response_format = options.get("response_format")
output_format_type = response_format if isinstance(response_format, type) else None
return ChatResponse.from_updates(updates, output_format_type=output_format_type)
return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
return ResponseStream(_stream(), finalizer=_finalize)
@@ -233,9 +231,7 @@ class MockBaseChatClient(
await asyncio.sleep(0)
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
response_format = options.get("response_format")
output_format_type = response_format if isinstance(response_format, type) else None
return ChatResponse.from_updates(updates, output_format_type=output_format_type)
return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
return ResponseStream(_stream(), finalizer=_finalize)
@@ -301,6 +301,56 @@ async def test_chat_client_agent_streaming_response_format_from_run_options(
assert result.value.greeting == "Hi"
async def test_chat_client_agent_response_format_dict_from_default_options(
client: SupportsChatGetResponse,
) -> None:
"""AgentResponse.value should parse JSON dicts from default_options response_format."""
json_text = json.dumps({"greeting": "Hello"})
client.responses.append(ChatResponse(messages=Message(role="assistant", text=json_text))) # type: ignore[attr-defined]
agent = Agent(
client=client,
default_options={"response_format": {"type": "object", "properties": {"greeting": {"type": "string"}}}},
)
result = await agent.run("Hello")
assert result.text == json_text
assert result.value is not None
assert isinstance(result.value, dict)
assert result.value["greeting"] == "Hello"
async def test_chat_client_agent_streaming_response_format_dict_from_run_options(
client: SupportsChatGetResponse,
) -> None:
"""Agent streaming should preserve mapping response_format and parse the final value as a dict."""
json_text = json.dumps({"greeting": "Hi"})
client.streaming_responses.append( # type: ignore[attr-defined]
[
ChatResponseUpdate(
contents=[Content.from_text(json_text)],
role="assistant",
finish_reason="stop",
)
]
)
agent = Agent(client=client)
stream = agent.run(
"Hello",
stream=True,
options={"response_format": {"type": "object", "properties": {"greeting": {"type": "string"}}}},
)
async for _ in stream:
pass
result = await stream.get_final_response()
assert result.text == json_text
assert result.value is not None
assert isinstance(result.value, dict)
assert result.value["greeting"] == "Hi"
async def test_chat_client_agent_create_session(
client: SupportsChatGetResponse,
) -> None:
@@ -191,9 +191,7 @@ def mock_chat_client():
yield ChatResponseUpdate(contents=[Content.from_text(" world")], role="assistant", finish_reason="stop")
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
response_format = options.get("response_format")
output_format_type = response_format if isinstance(response_format, type) else None
return ChatResponse.from_updates(updates, output_format_type=output_format_type)
return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
return ResponseStream(_stream(), finalizer=_finalize)
@@ -800,6 +800,19 @@ def test_chat_response_with_format_init():
assert response.value.response == "Hello"
def test_chat_response_with_mapping_response_format() -> None:
"""ChatResponse.value should parse JSON when response_format is a mapping."""
message = Message(role="assistant", text='{"response": "Hello"}')
response = ChatResponse(
messages=message,
response_format={"type": "object", "properties": {"response": {"type": "string"}}},
)
assert response.value is not None
assert isinstance(response.value, dict)
assert response.value["response"] == "Hello"
def test_chat_response_value_raises_on_invalid_schema():
"""Test that value property raises ValidationError with field constraint details."""
@@ -1004,6 +1017,22 @@ async def test_chat_response_from_async_generator_output_format_in_method():
assert resp.value.response == "Hello"
async def test_chat_response_from_async_generator_mapping_response_format() -> None:
async def gen() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text('{ "respon')], message_id="1")
yield ChatResponseUpdate(contents=[Content.from_text('se": "Hello" }')], message_id="1")
resp = await ChatResponse.from_update_generator(
gen(),
output_format_type={"type": "object", "properties": {"response": {"type": "string"}}},
)
assert resp.text == '{ "response": "Hello" }'
assert resp.value is not None
assert isinstance(resp.value, dict)
assert resp.value["response"] == "Hello"
# region ToolMode