From 65613328336478eedd4c656c63ecb65ac875e86b Mon Sep 17 00:00:00 2001 From: Aditya Raut Date: Mon, 28 Sep 2026 01:31:23 +0530 Subject: [PATCH 1/2] feat: add schema-constrained structured outputs and recovery loop to Agent Fixes #12989 --- haystack/components/agents/agent.py | 189 +++++++++++++++- ...gent-response-schema-eb2baa37108e4b44.yaml | 33 +++ test/components/agents/test_agent.py | 210 ++++++++++++++++++ 3 files changed, 428 insertions(+), 4 deletions(-) create mode 100644 releasenotes/notes/agent-response-schema-eb2baa37108e4b44.yaml diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index dc72a9433e1..ff3adbbff94 100644 --- a/haystack/components/agents/agent.py +++ b/haystack/components/agents/agent.py @@ -3,9 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import inspect +import json +import re from dataclasses import dataclass from typing import Any, Literal, cast +import jsonschema +from pydantic import BaseModel, ValidationError + from haystack import component, logging, tracing from haystack.components.agents.state.state import ( State, @@ -62,6 +67,7 @@ from haystack.utils.async_utils import _execute_component_async from haystack.utils.callable_serialization import deserialize_callable, serialize_callable from haystack.utils.deserialization import deserialize_component_inplace +from haystack.utils.type_serialization import deserialize_type, serialize_type logger = logging.getLogger(__name__) @@ -183,6 +189,68 @@ def _get_model_exit_reason(messages: list[ChatMessage]) -> str | None: return None +def _extract_json_text(text: str) -> str: + """Extract JSON string from text, stripping markdown code blocks if present.""" + match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text) + if match: + return match.group(1).strip() + return text.strip() + + +def _validate_against_schema( + raw_text: str, response_schema: type[BaseModel] | dict[str, Any] +) -> tuple[Any, str | None]: + """ + Validate raw text against a response schema (Pydantic model or JSON Schema dict). + + :returns: A tuple of (parsed_output, error_message). On success, error_message is None. + """ + cleaned_text = _extract_json_text(raw_text) + try: + data = json.loads(cleaned_text) + except (json.JSONDecodeError, TypeError) as e: + return None, f"Invalid JSON format: {e}" + + if isinstance(response_schema, type) and issubclass(response_schema, BaseModel): + try: + if hasattr(response_schema, "model_validate"): + parsed = response_schema.model_validate(data) + else: + parsed = response_schema.parse_obj(data) + return parsed, None + except ValidationError as e: + return None, f"JSON does not conform to schema: {e}" + elif isinstance(response_schema, dict): + try: + jsonschema.validate(instance=data, schema=response_schema) + return data, None + except jsonschema.ValidationError as e: + return None, f"JSON does not conform to schema: {e.message}" + except jsonschema.exceptions.SchemaError as e: + return None, f"Invalid JSON schema: {e.message}" + return None, f"Unsupported response_schema type: {type(response_schema)}" + + +def _schema_correction_prompt(error_msg: str, response_schema: type[BaseModel] | dict[str, Any]) -> str: + """Construct a user prompt instructing the model to fix its invalid output.""" + if isinstance(response_schema, type) and issubclass(response_schema, BaseModel): + if hasattr(response_schema, "model_json_schema"): + schema_str = json.dumps(response_schema.model_json_schema(), indent=2) + else: + schema_str = json.dumps(response_schema.schema(), indent=2) + elif isinstance(response_schema, dict): + schema_str = json.dumps(response_schema, indent=2) + else: + schema_str = str(response_schema) + + return ( + f"Your previous response did not match the expected schema:\n" + f"{error_msg}\n\n" + f"Please provide only a valid JSON response adhering to the following schema:\n" + f"{schema_str}" + ) + + def _pending_tool_call_messages_from_state(state: State) -> list[ChatMessage]: """ Return the pending tool-call message after `before_tool` hooks have run. @@ -219,6 +287,7 @@ class _ExecutionContext: chat_generator_inputs: dict tool_execution_inputs: dict counter: int = 0 + schema_retry_count: int = 0 @component @@ -377,7 +446,7 @@ def require_save(state: State) -> None: """ - def __init__( # noqa: PLR0913 + def __init__( # noqa: PLR0913, PLR0915 self, *, chat_generator: ChatGenerator, @@ -393,6 +462,8 @@ def __init__( # noqa: PLR0913 tool_concurrency_limit: int = 4, tool_streaming_callback_passthrough: bool = False, hooks: dict[HookPoint, list[Hook]] | None = None, + response_schema: type[BaseModel] | dict[str, Any] | None = None, + max_schema_retries: int = 3, ) -> None: """ Initialize the agent component. @@ -451,6 +522,11 @@ def __init__( # noqa: PLR0913 reached (unlike "on_exit"). Mutations to the state (e.g. appending a final message) are reflected in the returned `messages` / `last_message` and `state_schema` outputs. Setting `continue_run` here has no effect. + :param response_schema: Optional Pydantic BaseModel subclass or JSON schema dictionary to validate + and parse the Agent's final text response against. When set, the validated output is made available + under the `"structured_output"` key in the run output and state. + :param max_schema_retries: Maximum number of correction attempts to make if the model's text output + fails schema validation. Defaults to 3. :raises TypeError: If the chat_generator does not support tools parameter in its run method. :raises ValueError: If any `user_prompt` variable overlaps with the `state_schema` or `run` method parameters, if a hook is registered under an unknown hook point, or if a hook is registered under a hook point it does @@ -469,11 +545,29 @@ def __init__( # noqa: PLR0913 "The Agent component requires a chat generator that supports tools when tools are provided." ) + if response_schema is not None: + is_pydantic = isinstance(response_schema, type) and issubclass(response_schema, BaseModel) + is_json_schema = isinstance(response_schema, dict) + if not (is_pydantic or is_json_schema): + raise TypeError( + f"response_schema must be a Pydantic BaseModel subclass or a JSON schema dict, " + f"got {type(response_schema)}." + ) + if is_json_schema: + try: + jsonschema.Draft202012Validator.check_schema(response_schema) + except jsonschema.exceptions.SchemaError as e: + raise ValueError(f"Invalid JSON schema provided in response_schema: {e.message}") from e + if max_schema_retries < 0: + raise ValueError("max_schema_retries must be greater than or equal to 0.") + if exit_conditions is None: exit_conditions = ["text"] if state_schema is not None: reserved_keys = _RUN_METADATA_STATE_KEYS.keys() | _INTERNAL_STATE_KEYS.keys() + if response_schema is not None: + reserved_keys = reserved_keys | {"structured_output"} reserved_used = sorted(set(state_schema) & reserved_keys) if reserved_used: raise ValueError( @@ -503,6 +597,8 @@ def __init__( # noqa: PLR0913 self.tool_concurrency_limit = tool_concurrency_limit self.tool_streaming_callback_passthrough = tool_streaming_callback_passthrough self.hooks = hooks + self.response_schema = response_schema + self.max_schema_retries = max_schema_retries # --- State schema --- # shallow copy is sufficient: we only add a top-level "messages" key, never mutate nested values @@ -512,6 +608,13 @@ def __init__( # noqa: PLR0913 self.resolved_state_schema["messages"] = {"type": list[ChatMessage], "handler": merge_lists} for key, config in {**_RUN_METADATA_STATE_KEYS, **_INTERNAL_STATE_KEYS}.items(): self.resolved_state_schema[key] = dict(config) + if self.response_schema is not None: + out_type = ( + self.response_schema | None + if isinstance(self.response_schema, type) and issubclass(self.response_schema, BaseModel) + else dict[str, Any] | None + ) + self.resolved_state_schema["structured_output"] = {"type": out_type, "handler": replace_values} # --- Component I/O --- self._run_method_params = _get_run_method_params(self) @@ -522,7 +625,11 @@ def __init__( # noqa: PLR0913 continue output_types[param] = config["type"] # Run-metadata keys are populated by the Agent itself and exposed as outputs only, not inputs. - if param not in self._run_method_params and param not in _RUN_METADATA_STATE_KEYS: + if ( + param not in self._run_method_params + and param not in _RUN_METADATA_STATE_KEYS + and param != "structured_output" + ): component.set_input_type(self, name=param, type=config["type"], default=None) component.set_output_types(self, **output_types) @@ -639,6 +746,13 @@ def to_dict(self) -> dict[str, Any]: :returns: Dictionary with serialized data. """ + serialized_response_schema: str | dict[str, Any] | None = None + if self.response_schema is not None: + if isinstance(self.response_schema, type) and issubclass(self.response_schema, BaseModel): + serialized_response_schema = serialize_type(self.response_schema) + elif isinstance(self.response_schema, dict): + serialized_response_schema = self.response_schema + return default_to_dict( self, chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"), @@ -654,6 +768,8 @@ def to_dict(self) -> dict[str, Any]: tool_concurrency_limit=self.tool_concurrency_limit, tool_streaming_callback_passthrough=self.tool_streaming_callback_passthrough, hooks=_serialize_hooks_dictionary(self.hooks) if self.hooks else None, + response_schema=serialized_response_schema, + max_schema_retries=self.max_schema_retries, ) @classmethod @@ -679,6 +795,9 @@ def from_dict(cls, data: dict[str, Any]) -> "Agent": if init_params.get("hooks") is not None: init_params["hooks"] = _deserialize_hooks_dictionary(init_params["hooks"]) + if (resp_schema := init_params.get("response_schema")) is not None and isinstance(resp_schema, str): + init_params["response_schema"] = deserialize_type(resp_schema) + return default_from_dict(cls, data) def _create_agent_span(self, tools: ToolsType) -> Any: @@ -772,6 +891,8 @@ def _initialize_fresh_execution( state.set("continue_run", False) state.set("tools", flat_tools) state.set("hook_context", hook_context or {}) + if self.response_schema is not None: + state.set("structured_output", None) streaming_callback = select_streaming_callback( # type: ignore[call-overload] init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async @@ -870,6 +991,9 @@ def run( tool's result), or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition), or a custom reason a hook supplied through the `stop_run` state key. - Any additional keys defined in the `state_schema`. + - "structured_output": The parsed and validated structured object (Pydantic model instance + or dict) if `response_schema` was configured and validation succeeded, or None if validation + failed or retries were exhausted. """ agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs} self.warm_up() @@ -956,6 +1080,9 @@ async def run_async( tool's result), or `"max_agent_steps"` (the Agent hit `max_agent_steps` before meeting an exit condition), or a custom reason a hook supplied through the `stop_run` state key. - Any additional keys defined in the `state_schema`. + - "structured_output": The parsed and validated structured object (Pydantic model instance + or dict) if `response_schema` was configured and validation succeeded, or None if validation + failed or retries were exhausted. """ agent_inputs = {"messages": messages, "streaming_callback": streaming_callback, **kwargs} await self.warm_up_async() @@ -1032,7 +1159,14 @@ def _run_step(self, exe_context: _ExecutionContext, agent_span: tracing.Span) -> if not current_tools or model_exit_reason is not None: exe_context.counter += 1 exe_context.state.set("step_count", exe_context.counter) - exe_context.state.set("exit_reason", model_exit_reason or _EXIT_REASON_TEXT) + exit_reason = model_exit_reason or _EXIT_REASON_TEXT + if ( + self.response_schema is not None + and model_exit_reason == _EXIT_REASON_TEXT + and self._handle_response_schema(exe_context=exe_context, llm_messages=llm_messages) + ): + return True + exe_context.state.set("exit_reason", exit_reason) return self._continue_after_exit_hooks(exe_context=exe_context) _run_hooks(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state) @@ -1102,7 +1236,14 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac if not current_tools or model_exit_reason is not None: exe_context.counter += 1 exe_context.state.set("step_count", exe_context.counter) - exe_context.state.set("exit_reason", model_exit_reason or _EXIT_REASON_TEXT) + exit_reason = model_exit_reason or _EXIT_REASON_TEXT + if ( + self.response_schema is not None + and model_exit_reason == _EXIT_REASON_TEXT + and self._handle_response_schema(exe_context=exe_context, llm_messages=llm_messages) + ): + return True + exe_context.state.set("exit_reason", exit_reason) return await self._continue_after_exit_hooks_async(exe_context=exe_context) await _run_hooks_async(hooks=self.hooks, hook_point=BEFORE_TOOL, state=exe_context.state) @@ -1133,6 +1274,46 @@ async def _run_step_async(self, exe_context: _ExecutionContext, agent_span: trac return await self._continue_after_exit_hooks_async(exe_context=exe_context) return True + def _handle_response_schema(self, exe_context: _ExecutionContext, llm_messages: list[ChatMessage]) -> bool: + """ + Validate the assistant reply against response_schema. + + If validation succeeds, writes structured_output to state and returns False (no retry). + If validation fails and retries are available, appends a correction message to state, + increments retry count, and returns True (retry). + If validation fails and retries are exhausted, sets structured_output to None, logs a + warning, and returns False (no retry). + """ + if self.response_schema is None or not llm_messages: + return False + + last_msg = llm_messages[-1] + raw_text = last_msg.text or "" + parsed, error = _validate_against_schema(raw_text, self.response_schema) + if error is None: + exe_context.state.set("structured_output", parsed) + return False + + if exe_context.schema_retry_count < self.max_schema_retries and exe_context.counter < self.max_agent_steps: + exe_context.schema_retry_count += 1 + logger.info( + "Model output failed response_schema validation ({count}/{max_retries}). " + "Prompting for correction. Error: {error}", + count=exe_context.schema_retry_count, + max_retries=self.max_schema_retries, + error=error, + ) + correction_text = _schema_correction_prompt(error, self.response_schema) + correction_msg = ChatMessage.from_user(correction_text) + exe_context.state.set("messages", [correction_msg], handler_override=merge_lists) + return True + + logger.warning( + "Model output failed response_schema validation and retry limit was reached: {error}", error=error + ) + exe_context.state.set("structured_output", None) + return False + def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> str | None: """ Decide whether the agent should stop looping and, if so, on which tool. diff --git a/releasenotes/notes/agent-response-schema-eb2baa37108e4b44.yaml b/releasenotes/notes/agent-response-schema-eb2baa37108e4b44.yaml new file mode 100644 index 00000000000..5d24fe614f7 --- /dev/null +++ b/releasenotes/notes/agent-response-schema-eb2baa37108e4b44.yaml @@ -0,0 +1,33 @@ +--- +features: + - | + Added ``response_schema`` and ``max_schema_retries`` parameters to the ``Agent`` component, + enabling schema-constrained structured output validation and automatic recovery. + + Users can provide a Pydantic ``BaseModel`` subclass or a JSON Schema dictionary. When the Agent + reaches a terminal text response, it validates the generated output against the schema. + If valid, the parsed model or dictionary is emitted under the ``structured_output`` state and + output socket. If validation fails, the Agent automatically appends a targeted correction prompt + to the conversation history and retries up to ``max_schema_retries`` times within ``max_agent_steps``. + + Example usage: + .. code:: python + + from pydantic import BaseModel + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIChatGenerator + from haystack.dataclasses import ChatMessage + + class UserProfile(BaseModel): + name: str + age: int + interests: list[str] + + agent = Agent( + chat_generator=OpenAIChatGenerator(), + response_schema=UserProfile, + max_schema_retries=3, + ) + + result = agent.run(messages=[ChatMessage.from_user("Extract profile for Alice, 30, loves cycling.")]) + profile: UserProfile | None = result["structured_output"] diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 08e568fd3ac..f5b42a04560 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -15,6 +15,7 @@ from jinja2 import TemplateSyntaxError from openai import OpenAI, Stream from openai.types.chat import ChatCompletionChunk, chat_completion_chunk +from pydantic import BaseModel from haystack import Document, Pipeline, component from haystack.components.agents.agent import Agent, _get_model_exit_reason @@ -29,6 +30,7 @@ from haystack.components.retrievers.in_memory import InMemoryBM25Retriever from haystack.components.routers.conditional_router import ConditionalRouter from haystack.core.component.types import OutputSocket +from haystack.core.serialization import allow_deserialization_module from haystack.dataclasses import ChatMessage, ToolCall from haystack.dataclasses.chat_message import ChatRole, TextContent from haystack.dataclasses.streaming_chunk import StreamingChunk @@ -390,6 +392,8 @@ def test_to_dict(self, weather_tool, component_tool, monkeypatch): "tool_concurrency_limit": 5, "tool_streaming_callback_passthrough": True, "hooks": None, + "response_schema": None, + "max_schema_retries": 3, }, } assert serialized_agent == expected_structure @@ -2279,3 +2283,209 @@ def run(self) -> dict: assert "agent" not in result run_mock.assert_not_called() + + +class CityInfo(BaseModel): + city: str + temperature: float + + +class TestAgentStructuredOutput: + def test_init_validation(self, weather_tool): + # Invalid response_schema type + with pytest.raises( + TypeError, match="response_schema must be a Pydantic BaseModel subclass or a JSON schema dict" + ): + Agent(chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema=123) # type: ignore[arg-type] + + # Negative max_schema_retries + with pytest.raises(ValueError, match="max_schema_retries must be greater than or equal to 0"): + Agent( + chat_generator=MockChatGenerator("hi"), + tools=[weather_tool], + response_schema=CityInfo, + max_schema_retries=-1, + ) + + # Reserved structured_output in state_schema when response_schema is set + with pytest.raises(ValueError, match="reserved for Agent internal state"): + Agent( + chat_generator=MockChatGenerator("hi"), + tools=[weather_tool], + response_schema=CityInfo, + state_schema={"structured_output": {"type": str}}, + ) + + # Malformed JSON schema dict + with pytest.raises(ValueError, match="Invalid JSON schema provided in response_schema"): + Agent( + chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema={"type": "invalid_type"} + ) + + # Output types + agent = Agent(chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema=CityInfo) + output_sockets = agent.__haystack_output__._sockets_dict + input_sockets = agent.__haystack_input__._sockets_dict + assert "structured_output" in output_sockets + assert output_sockets["structured_output"].type == (CityInfo | None) + assert "structured_output" not in input_sockets + + # Dict schema output type + dict_schema = {"type": "object", "properties": {"city": {"type": "string"}}} + agent_dict = Agent(chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema=dict_schema) + output_sockets_dict = agent_dict.__haystack_output__._sockets_dict + assert "structured_output" in output_sockets_dict + assert output_sockets_dict["structured_output"].type == (dict[str, Any] | None) + + def test_no_tools_with_tool_call_does_not_trigger_schema_validation(self): + reply_with_tool_call = ChatMessage.from_assistant( + tool_calls=[ToolCall(tool_name="unknown_tool", arguments={"foo": "bar"})] + ) + generator = MockChatGenerator(reply_with_tool_call) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=2) + result = agent.run(messages=[ChatMessage.from_user("Run unknown tool")]) + + assert result["exit_reason"] == "text" + assert result["step_count"] == 1 + assert result["structured_output"] is None + + def test_pydantic_schema_success(self): + generator = MockChatGenerator('{"city": "Paris", "temperature": 21.5}') + agent = Agent(chat_generator=generator, response_schema=CityInfo) + result = agent.run(messages=[ChatMessage.from_user("Weather in Paris?")]) + + assert isinstance(result["structured_output"], CityInfo) + assert result["structured_output"].city == "Paris" + assert result["structured_output"].temperature == 21.5 + assert result["exit_reason"] == "text" + assert result["step_count"] == 1 + + def test_json_schema_dict_success(self): + schema = { + "type": "object", + "properties": {"city": {"type": "string"}, "count": {"type": "integer"}}, + "required": ["city", "count"], + } + generator = MockChatGenerator('{"city": "Berlin", "count": 42}') + agent = Agent(chat_generator=generator, response_schema=schema) + result = agent.run(messages=[ChatMessage.from_user("Info?")]) + + assert result["structured_output"] == {"city": "Berlin", "count": 42} + assert result["exit_reason"] == "text" + assert result["step_count"] == 1 + + def test_markdown_fenced_json(self): + response_text = '```json\n{"city": "Tokyo", "temperature": 18.0}\n```' + generator = MockChatGenerator(response_text) + agent = Agent(chat_generator=generator, response_schema=CityInfo) + result = agent.run(messages=[ChatMessage.from_user("Weather in Tokyo?")]) + + assert isinstance(result["structured_output"], CityInfo) + assert result["structured_output"].city == "Tokyo" + assert result["step_count"] == 1 + + def test_retry_on_invalid_json_then_success(self): + replies = [ + ChatMessage.from_assistant("I cannot provide JSON right now."), + ChatMessage.from_assistant('{"city": "Rome", "temperature": 25.0}'), + ] + generator = MockChatGenerator(replies) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=2) + result = agent.run(messages=[ChatMessage.from_user("Weather in Rome?")]) + + assert isinstance(result["structured_output"], CityInfo) + assert result["structured_output"].city == "Rome" + assert result["step_count"] == 2 + # Verify correction message was injected as a user turn + assert len(result["messages"]) == 4 + assert result["messages"][0].is_from(ChatRole.USER) + assert result["messages"][1].is_from(ChatRole.ASSISTANT) + assert result["messages"][2].is_from(ChatRole.USER) + assert "Your previous response did not match the expected schema" in result["messages"][2].text + assert result["messages"][3].is_from(ChatRole.ASSISTANT) + + def test_retry_on_schema_mismatch_then_success(self): + replies = [ + ChatMessage.from_assistant('{"city": "Madrid", "temperature": "twenty"}'), # string instead of float + ChatMessage.from_assistant('{"city": "Madrid", "temperature": 20.0}'), + ] + generator = MockChatGenerator(replies) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=2) + result = agent.run(messages=[ChatMessage.from_user("Weather in Madrid?")]) + + assert isinstance(result["structured_output"], CityInfo) + assert result["structured_output"].temperature == 20.0 + assert result["step_count"] == 2 + + def test_retries_exhausted_returns_none(self): + replies = [ + ChatMessage.from_assistant("Invalid JSON 1"), + ChatMessage.from_assistant("Invalid JSON 2"), + ChatMessage.from_assistant("Invalid JSON 3"), + ] + generator = MockChatGenerator(replies) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=2) + result = agent.run(messages=[ChatMessage.from_user("Weather?")]) + + assert result["structured_output"] is None + assert result["exit_reason"] == "text" + assert result["step_count"] == 3 + assert result["last_message"].text == "Invalid JSON 3" + + def test_max_agent_steps_stops_retries(self): + replies = [ChatMessage.from_assistant("Invalid JSON 1"), ChatMessage.from_assistant("Invalid JSON 2")] + generator = MockChatGenerator(replies) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=5, max_agent_steps=1) + result = agent.run(messages=[ChatMessage.from_user("Weather?")]) + + assert result["structured_output"] is None + assert result["step_count"] == 1 + + async def test_run_async_mirrors_sync(self): + replies = [ + ChatMessage.from_assistant("Not JSON"), + ChatMessage.from_assistant('{"city": "London", "temperature": 15.0}'), + ] + generator = MockChatGenerator(replies) + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=2) + result = await agent.run_async(messages=[ChatMessage.from_user("Weather in London?")]) + + assert isinstance(result["structured_output"], CityInfo) + assert result["structured_output"].city == "London" + assert result["step_count"] == 2 + + def test_response_schema_none_backward_compatibility(self): + generator = MockChatGenerator("Plain text response") + agent = Agent(chat_generator=generator) + result = agent.run(messages=[ChatMessage.from_user("Hello")]) + + assert "structured_output" not in result + assert "structured_output" not in agent.resolved_state_schema + assert result["last_message"].text == "Plain text response" + + def test_serialization_and_deserialization(self): + allow_deserialization_module(CityInfo.__module__) + generator = MockChatGenerator("test") + agent = Agent(chat_generator=generator, response_schema=CityInfo, max_schema_retries=5) + serialized = agent.to_dict() + + assert serialized["init_parameters"]["response_schema"] == f"{CityInfo.__module__}.CityInfo" + assert serialized["init_parameters"]["max_schema_retries"] == 5 + + deserialized = Agent.from_dict(serialized) + assert deserialized.response_schema == CityInfo + assert deserialized.max_schema_retries == 5 + + # Dict schema serialization + dict_schema = {"type": "object", "properties": {"x": {"type": "number"}}} + agent_dict = Agent(chat_generator=generator, response_schema=dict_schema) + serialized_dict = agent_dict.to_dict() + assert serialized_dict["init_parameters"]["response_schema"] == dict_schema + + deserialized_dict = Agent.from_dict(serialized_dict) + assert deserialized_dict.response_schema == dict_schema + + # Clone + cloned = agent.clone(max_schema_retries=1) + assert cloned.response_schema == CityInfo + assert cloned.max_schema_retries == 1 From 11b80bc9152c6b931445ccccf5dc9e4f09025c3b Mon Sep 17 00:00:00 2001 From: Aditya Raut Date: Mon, 28 Sep 2026 01:45:19 +0530 Subject: [PATCH 2/2] test: update test_to_dict assertions for Agent response_schema and add mypy type ignores --- test/components/agents/test_agent.py | 6 +++--- test/components/agents/test_agent_hitl.py | 2 ++ test/tools/test_agent_tool.py | 2 ++ 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index f5b42a04560..45c4fd5f562 100644 --- a/test/components/agents/test_agent.py +++ b/test/components/agents/test_agent.py @@ -2324,8 +2324,8 @@ def test_init_validation(self, weather_tool): # Output types agent = Agent(chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema=CityInfo) - output_sockets = agent.__haystack_output__._sockets_dict - input_sockets = agent.__haystack_input__._sockets_dict + output_sockets = agent.__haystack_output__._sockets_dict # type: ignore[attr-defined] + input_sockets = agent.__haystack_input__._sockets_dict # type: ignore[attr-defined] assert "structured_output" in output_sockets assert output_sockets["structured_output"].type == (CityInfo | None) assert "structured_output" not in input_sockets @@ -2333,7 +2333,7 @@ def test_init_validation(self, weather_tool): # Dict schema output type dict_schema = {"type": "object", "properties": {"city": {"type": "string"}}} agent_dict = Agent(chat_generator=MockChatGenerator("hi"), tools=[weather_tool], response_schema=dict_schema) - output_sockets_dict = agent_dict.__haystack_output__._sockets_dict + output_sockets_dict = agent_dict.__haystack_output__._sockets_dict # type: ignore[attr-defined] assert "structured_output" in output_sockets_dict assert output_sockets_dict["structured_output"].type == (dict[str, Any] | None) diff --git a/test/components/agents/test_agent_hitl.py b/test/components/agents/test_agent_hitl.py index 94ef3e96923..cbc1a81c13c 100644 --- a/test/components/agents/test_agent_hitl.py +++ b/test/components/agents/test_agent_hitl.py @@ -153,6 +153,8 @@ def test_to_dict(self, tools, confirmation_hook, monkeypatch): } ] }, + "response_schema": None, + "max_schema_retries": 3, }, } diff --git a/test/tools/test_agent_tool.py b/test/tools/test_agent_tool.py index 884666b1667..eb215019ff7 100644 --- a/test/tools/test_agent_tool.py +++ b/test/tools/test_agent_tool.py @@ -185,6 +185,8 @@ def test_to_dict(self): "tool_concurrency_limit": 4, "tool_streaming_callback_passthrough": False, "hooks": None, + "response_schema": None, + "max_schema_retries": 3, }, }, },