diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index b7febaf4256..f17204bcf49 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3067,6 +3067,18 @@ class ContinuationToken(TypedDict): # endregion +def _validate_json_by_alias(model: Any, text: str) -> Any: + """Validate JSON against ``model`` using only field aliases, as in the schema sent to providers. + + The by_alias/by_name overrides need pydantic 2.11; earlier versions use the model's own settings. + """ + from pydantic import VERSION + + if tuple(int(part) for part in VERSION.split(".")[:2]) >= (2, 11): + return model.model_validate_json(text, by_alias=True, by_name=False) + return model.model_validate_json(text) + + def _parse_structured_response_value(text: str, response_format: Any | None) -> Any | None: if response_format is None: return None @@ -3076,7 +3088,7 @@ def _parse_structured_response_value(text: str, response_format: Any | None) -> from pydantic import BaseModel if issubclass(response_format, BaseModel): - return response_format.model_validate_json(text) + return _validate_json_by_alias(response_format, text) if isinstance(response_format, Mapping): try: return json.loads(text) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 626353676ed..e26177fdf9d 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -7721,3 +7721,23 @@ def test_data_uri_readers_agree_on_the_base64_marker(): payload_marker = f"data:image/png,x;base64,{png_b64}" with pytest.raises(ValueError, match="Data URI must use base64 encoding."): detect_media_type_from_base64(data_uri=payload_marker) + + +def test_agent_response_value_validates_by_alias_only() -> None: + """Structured output is validated with aliases only, matching the schema sent to providers.""" + from pydantic import VERSION, ConfigDict + + if tuple(int(part) for part in VERSION.split(".")[:2]) < (2, 11): + pytest.skip("validate_by_alias/validate_by_name require pydantic 2.11") + + class Aliased(BaseModel): + model_config = ConfigDict(validate_by_alias=False, validate_by_name=True) + item_count: int = Field(alias="itemCount") + + response = AgentResponse( + messages=[Message(role="assistant", contents=['{"itemCount": 3}'])], + response_format=Aliased, + ) + + assert response.value is not None + assert response.value.item_count == 3 diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 4dfdd70894a..8d209097868 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -53,6 +53,7 @@ ) from agent_framework.exceptions import AgentException, ContentError from agent_framework.observability import AgentTelemetryLayer +from pydantic import BaseModel from ._feature_usage import FeatureIndex @@ -193,6 +194,25 @@ async def _resolve_function_approval( logger = logging.getLogger("agent_framework.github_copilot") +def _to_response_schema(response_format: Any) -> Any: + """Return a value the Copilot SDK can JSON-serialize as ``response_schema``. + + Pydantic model classes are passed through unchanged. Any other mapping (for example a + ``MappingProxyType`` or ``UserDict``) is converted, recursively, to plain ``dict``/``list`` objects. + """ + if response_format is None or isinstance(response_format, type): + return response_format + + def _plain(value: Any) -> Any: + if isinstance(value, Mapping): + return {key: _plain(item) for key, item in cast(Mapping[Any, Any], value).items()} + if isinstance(value, (list, tuple)): + return [_plain(item) for item in cast(Sequence[Any], value)] + return value + + return _plain(response_format) + + def _provider_service_session_state_keys(agent: object, client: object | None) -> frozenset[str]: """Return provider-owned session-state keys declared by the Copilot agent or client.""" keys: set[str] = set() @@ -575,6 +595,13 @@ class GitHubCopilotOptions(TypedDict, total=False): telemetry: TelemetryConfig """OpenTelemetry configuration for the Copilot CLI process.""" + response_format: type[BaseModel] | Mapping[str, Any] + """Structured output format for the run: a Pydantic model type or a JSON Schema mapping. + + Forwarded to the Copilot SDK as ``response_schema`` so the provider constrains the + final answer to the schema, and the parsed result is exposed on ``AgentResponse.value``. + Streaming still yields text; the value is parsed once the stream completes.""" + on_pre_tool_use: PreToolUseHandler """Pre-tool-use hook handler for the Copilot SDK. @@ -916,7 +943,14 @@ async def _after_run_hook(response: AgentResponse) -> None: logger.exception("Error running after_run providers in streaming result hook") def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: - return AgentResponse.from_updates(updates) + response_format: Any = cast(Mapping[str, Any], options or {}).get( + "response_format" + ) or self._default_options.get("response_format") + if response_format is not None: + # The structured value comes from the root agent's answer, so sub-agent updates + # (which still stream to the caller) are left out of the final response. + updates = [update for update in updates if not getattr(update.raw_representation, "agent_id", None)] + return AgentResponse.from_updates(updates, output_format_type=response_format) return ResponseStream( self._stream_updates( @@ -986,6 +1020,8 @@ async def _run_impl( "together with on_permission_request instead." ) timeout = opts.get("timeout") or self._settings.get("timeout") or DEFAULT_TIMEOUT_SECONDS + response_format = opts.get("response_format") or self._default_options.get("response_format") + response_schema = _to_response_schema(response_format) input_messages = normalize_messages(messages) @@ -1038,7 +1074,12 @@ def usage_event_handler(event: SessionEvent) -> None: unsubscribe = copilot_session.on(usage_event_handler) try: mark_feature_used(FeatureIndex.GITHUB_COPILOT) - response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) + response_event = await copilot_session.send_and_wait( + prompt, + attachments=attachments, + response_schema=response_schema, + timeout=timeout, + ) except Exception as ex: raise AgentException(f"GitHub Copilot request failed: {ex}") from ex finally: @@ -1070,6 +1111,7 @@ def usage_event_handler(event: SessionEvent) -> None: finish_reason=cast(Any, finish_reason), usage_details=usage_details, additional_properties={"model": model} if model else None, + response_format=response_format, ) session_context._response = response # type: ignore[assignment] await self._run_after_providers(session=session, context=session_context) @@ -1122,6 +1164,8 @@ async def _stream_updates( "(set via default_options). Remove on_function_approval and use on_pre_tool_use " "together with on_permission_request instead." ) + response_format = opts.get("response_format") or self._default_options.get("response_format") + response_schema = _to_response_schema(response_format) input_messages = normalize_messages(messages) @@ -1219,7 +1263,9 @@ def event_handler(event: SessionEvent) -> None: ) queue.put_nowait(update) elif event.type == SessionEventType.SESSION_IDLE: - queue.put_nowait(None) + # A sub-agent going idle must not end the run; only the root agent's idle does. + if not getattr(event, "agent_id", None): + queue.put_nowait(None) elif event.type == SessionEventType.SESSION_ERROR: error_data: Any = event.data error_msg = error_data.message or "Unknown error" @@ -1239,7 +1285,7 @@ def event_handler(event: SessionEvent) -> None: try: mark_feature_used(FeatureIndex.GITHUB_COPILOT) - await copilot_session.send(prompt, attachments=attachments) + await copilot_session.send(prompt, attachments=attachments, response_schema=response_schema) while (item := await queue.get()) is not None: if isinstance(item, Exception): @@ -1794,6 +1840,7 @@ def _build_session_kwargs( "on_pre_tool_use", "on_function_approval", "timeout", + "response_format", "cli_path", "log_level", "base_directory", diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index d6da1598b6c..5e3dd2e8bac 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -11,12 +11,13 @@ from collections.abc import Sequence from datetime import datetime, timezone from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import pytest +from pydantic import BaseModel, ConfigDict, Field, ValidationError copilot = pytest.importorskip("copilot") @@ -1213,6 +1214,217 @@ def mock_on(handler: Any) -> Any: assert responses[3].contents[0].type == "text" +class TestGitHubCopilotAgentStructuredOutput: + """Test cases for response_format (structured output) support.""" + + class Answer(BaseModel): + value: int + + async def test_run_passes_response_format_as_response_schema( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """response_format is forwarded to send_and_wait and parsed into AgentResponse.value.""" + mock_session.send_and_wait.return_value = create_session_event( + SessionEventType.ASSISTANT_MESSAGE, content='{"value": 42}', message_id="msg-1" + ) + + agent = GitHubCopilotAgent(client=mock_client) + response = await agent.run("What is 19 + 23?", options=cast(Any, {"response_format": self.Answer})) + + assert mock_session.send_and_wait.call_args.kwargs["response_schema"] is self.Answer + assert response.value == self.Answer(value=42) + + async def test_run_uses_default_options_response_format( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """response_format from default_options is used and not leaked into create_session.""" + mock_session.send_and_wait.return_value = create_session_event( + SessionEventType.ASSISTANT_MESSAGE, content='{"value": 7}', message_id="msg-1" + ) + + agent = GitHubCopilotAgent(client=mock_client, default_options=cast(Any, {"response_format": self.Answer})) + response = await agent.run("Hi") + + assert mock_session.send_and_wait.call_args.kwargs["response_schema"] is self.Answer + assert "response_format" not in mock_client.create_session.call_args.kwargs + assert response.value == self.Answer(value=7) + + async def test_run_accepts_json_schema_mapping( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """A JSON Schema mapping is forwarded unchanged.""" + schema = {"type": "object", "properties": {"value": {"type": "integer"}}, "required": ["value"]} + mock_session.send_and_wait.return_value = create_session_event( + SessionEventType.ASSISTANT_MESSAGE, content='{"value": 1}', message_id="msg-1" + ) + + agent = GitHubCopilotAgent(client=mock_client) + response = await agent.run("Hi", options=cast(Any, {"response_format": schema})) + + assert mock_session.send_and_wait.call_args.kwargs["response_schema"] == schema + assert response.value == {"value": 1} + + async def test_run_without_response_format_sends_no_schema( + self, + mock_client: MagicMock, + mock_session: MagicMock, + assistant_message_event: SessionEvent, + ) -> None: + """No schema is sent when response_format is not set.""" + mock_session.send_and_wait.return_value = assistant_message_event + + agent = GitHubCopilotAgent(client=mock_client) + await agent.run("Hi") + + assert mock_session.send_and_wait.call_args.kwargs["response_schema"] is None + + async def test_run_converts_mapping_schema_to_plain_dict( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """Non-dict mappings are converted to plain dicts so the SDK can serialize them.""" + schema = MappingProxyType({ + "type": "object", + "properties": MappingProxyType({"value": MappingProxyType({"type": "integer"})}), + }) + mock_session.send_and_wait.return_value = create_session_event( + SessionEventType.ASSISTANT_MESSAGE, content='{"value": 1}', message_id="msg-1" + ) + + agent = GitHubCopilotAgent(client=mock_client) + await agent.run("Hi", options=cast(Any, {"response_format": schema})) + + sent = mock_session.send_and_wait.call_args.kwargs["response_schema"] + assert type(sent) is dict + assert type(sent["properties"]["value"]) is dict + json.dumps(sent) + + async def test_run_parses_aliased_model_by_alias( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """Aliased models are validated by alias only, matching the schema sent to the SDK.""" + + class Aliased(BaseModel): + model_config = ConfigDict(validate_by_alias=False, validate_by_name=True) + item_count: int = Field(alias="itemCount") + + mock_session.send_and_wait.return_value = create_session_event( + SessionEventType.ASSISTANT_MESSAGE, content='{"itemCount": 3}', message_id="msg-1" + ) + # Core validation with the model's own settings would fail here. + with pytest.raises(ValidationError): + Aliased.model_validate_json('{"itemCount": 3}') + + agent = GitHubCopilotAgent(client=mock_client) + response = await agent.run("Hi", options=cast(Any, {"response_format": Aliased})) + + assert response.value is not None + assert response.value.item_count == 3 + + async def test_streaming_value_ignores_trailing_subagent_text( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """A sub-agent delta after the root answer does not replace the structured value.""" + child = create_session_event( + SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content="child chatter", message_id="child" + ) + child.agent_id = "sub-agent" + events = [ + create_session_event( + SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content='{"value": 9}', message_id="root" + ), + child, + create_session_event(SessionEventType.SESSION_IDLE), + ] + + def mock_on(handler: Any) -> Any: + for event in events: + handler(event) + return lambda: None + + mock_session.on = mock_on + + agent = GitHubCopilotAgent(client=mock_client) + stream = agent.run("Hi", stream=True, options=cast(Any, {"response_format": self.Answer})) + texts = [u.text async for u in stream] + response = await stream.get_final_response() + + assert "child chatter" in "".join(texts) + assert response.value == self.Answer(value=9) + + async def test_streaming_ignores_subagent_idle_before_root_answer( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """A sub-agent idle event does not end the stream or drop the root structured answer.""" + child_idle = create_session_event(SessionEventType.SESSION_IDLE) + child_idle.agent_id = "sub-agent" + events = [ + create_session_event(SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content="preamble", message_id="pre"), + child_idle, + create_session_event( + SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content='{"value": 11}', message_id="root" + ), + create_session_event(SessionEventType.SESSION_IDLE), + ] + + def mock_on(handler: Any) -> Any: + for event in events: + handler(event) + return lambda: None + + mock_session.on = mock_on + + agent = GitHubCopilotAgent(client=mock_client) + stream = agent.run("Hi", stream=True, options=cast(Any, {"response_format": self.Answer})) + texts = [u.text async for u in stream] + response = await stream.get_final_response() + + assert '{"value": 11}' in "".join(texts) + assert response.value is not None + assert response.value.value == 11 + + async def test_streaming_passes_response_format_and_parses_value( + self, + mock_client: MagicMock, + mock_session: MagicMock, + ) -> None: + """Streaming forwards the schema to send and parses the final value after the stream.""" + events = [ + create_session_event(SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content='{"value": ', message_id="m"), + create_session_event(SessionEventType.ASSISTANT_MESSAGE_DELTA, delta_content="5}", message_id="m"), + create_session_event(SessionEventType.SESSION_IDLE), + ] + + def mock_on(handler: Any) -> Any: + for event in events: + handler(event) + return lambda: None + + mock_session.on = mock_on + + agent = GitHubCopilotAgent(client=mock_client) + stream = agent.run("Hi", stream=True, options=cast(Any, {"response_format": self.Answer})) + async for _ in stream: + pass + response = await stream.get_final_response() + + assert mock_session.send.call_args.kwargs["response_schema"] is self.Answer + assert response.value == self.Answer(value=5) + + class TestGitHubCopilotAgentSessionManagement: """Test cases for session management."""