mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python (fix:gemini): make Gemini honor declarative outputSchema, not just JSON mode (#5893)
* fix(gemini): preserve schema response_format * fix(gemini): satisfy pyright strict in response schema extraction Cast Any-narrowed mappings to Mapping[str, Any] in the structured-output schema helpers so pyright strict no longer reports partially-unknown member, argument, and variable types. Pass response_format["format"] straight into the recursive extractor, which already guards non-mapping inputs. No behavior change. * fix(gemini): use Sequence[object] cast to satisfy both mypy and pyright The Sequence[Any] cast pyright strict needs to know the loop element type is reported as a redundant-cast by mypy, which already narrows the isinstance branch to Sequence[Any]. Cast to Sequence[object] instead: pyright gets a fully known element type and mypy no longer sees an identical-type cast. No behavior change. --------- Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
This commit is contained in:
cooleryu
committed by
GitHub
Unverified
GitHub
parent
bf4ad48cf2
commit
d5335fbeae
@@ -9,7 +9,7 @@ from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import Content, FunctionTool, Message
|
||||
from agent_framework import Agent, Content, FunctionTool, Message
|
||||
from google.genai import types
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -915,6 +915,20 @@ async def test_response_format_populates_value_on_chat_response() -> None:
|
||||
assert response.value == Reply(text="hello")
|
||||
|
||||
|
||||
async def test_response_format_mapping_populates_value_on_chat_response() -> None:
|
||||
"""When response_format is a JSON schema mapping, ChatResponse.value must parse the response text."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"text": "hello"}')]))
|
||||
schema = {"type": "object", "properties": {"text": {"type": "string"}}}
|
||||
|
||||
response = await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
)
|
||||
|
||||
assert response.value == {"text": "hello"}
|
||||
|
||||
|
||||
async def test_response_schema_added_to_config() -> None:
|
||||
"""Sets both response_mime_type and the raw schema on the config when response_schema is given."""
|
||||
client, mock = _make_gemini_client()
|
||||
@@ -931,6 +945,284 @@ async def test_response_schema_added_to_config() -> None:
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_raw_json_schema_added_to_config() -> None:
|
||||
"""For declarative outputSchema, response_format may already be a raw JSON schema mapping."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"answer": {"type": "string", "description": "The answer."}},
|
||||
"required": ["answer"],
|
||||
}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_agent_default_options_response_format_raw_schema_added_to_config() -> None:
|
||||
"""Agent default_options is the path used by declarative outputSchema."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}
|
||||
agent = Agent(client=client, default_options={"response_format": schema})
|
||||
|
||||
await agent.run("Hi")
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_complex_raw_json_schema_preserved() -> None:
|
||||
"""Nested declarative schemas should be forwarded without losing shape or constraints."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "ok"}')]))
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"answer": {"type": "string"},
|
||||
"citations": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {"type": "string"},
|
||||
"confidence": {"type": "number"},
|
||||
},
|
||||
"required": ["source"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["answer"],
|
||||
"additionalProperties": False,
|
||||
}
|
||||
|
||||
await client.get_response(
|
||||
messages=[
|
||||
Message(
|
||||
role="user",
|
||||
contents=[
|
||||
Content.from_text(
|
||||
"Summarize a long document while preserving citation metadata.\n" + ("context\n" * 128)
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
options={"response_format": schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_json_schema_envelope_added_to_config() -> None:
|
||||
"""OpenAI-style json_schema envelopes should still provide Gemini with the inner schema."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"type": "json_schema", "json_schema": {"name": "Answer", "schema": schema}}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_format_envelope_added_to_config() -> None:
|
||||
"""Responses-style format envelopes should also provide Gemini with the nested schema."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"format": {"type": "json_schema", "name": "Answer", "schema": schema}}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_direct_schema_key_added_to_config() -> None:
|
||||
"""Provider-normalized mappings with a direct schema key should be accepted."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"schema": schema}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_json_schema_envelope_preserves_empty_schema() -> None:
|
||||
"""An explicitly empty JSON schema is still a schema and should not be dropped as falsy."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
|
||||
schema: dict[str, Any] = {}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"type": "json_schema", "json_schema": {"name": "AnyJson", "schema": schema}}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_anyof_raw_schema_added_to_config() -> None:
|
||||
"""Raw schemas without a type should still be recognized when they use JSON Schema keywords."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='"ok"')]))
|
||||
schema = {"anyOf": [{"type": "string"}, {"type": "number"}]}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_union_type_raw_schema_added_to_config() -> None:
|
||||
"""JSON Schema union type arrays should be treated as raw schemas."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
|
||||
schema = {"type": ["object", "null"], "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_response_format_json_object_does_not_set_schema() -> None:
|
||||
"""A JSON-object response_format requests JSON output but is not itself a Gemini response schema."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"type": "json_object"}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema is None
|
||||
|
||||
|
||||
async def test_response_format_json_schema_without_inner_schema_does_not_set_schema() -> None:
|
||||
"""A json_schema envelope without a schema should not be mistaken for a raw JSON schema."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": {"type": "json_schema", "json_schema": {"name": "MissingSchema"}}},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema is None
|
||||
|
||||
|
||||
async def test_response_schema_takes_precedence_over_response_format_schema() -> None:
|
||||
"""An explicit Gemini response_schema should win when both schema options are present."""
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
|
||||
response_format_schema = {"type": "object", "properties": {"name": {"type": "string"}}}
|
||||
response_schema = {"type": "object", "properties": {"id": {"type": "integer"}}}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": response_format_schema, "response_schema": response_schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == response_schema
|
||||
|
||||
|
||||
async def test_response_format_raw_schema_kept_with_tools() -> None:
|
||||
"""Structured output must still reach Gemini when function tools are present."""
|
||||
|
||||
def calculator(expression: str) -> str:
|
||||
"""Evaluate a simple expression."""
|
||||
return expression
|
||||
|
||||
tool = FunctionTool(name="calculator", func=calculator)
|
||||
client, mock = _make_gemini_client()
|
||||
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "4"}')]))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}
|
||||
|
||||
await client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("What is 2 + 2?")])],
|
||||
options={"tools": [tool], "response_format": schema},
|
||||
)
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
|
||||
assert config.response_schema == schema
|
||||
assert config.tools is not None
|
||||
assert config.tools[0].function_declarations[0].name == "calculator"
|
||||
|
||||
|
||||
async def test_streaming_response_format_raw_schema_added_to_config() -> None:
|
||||
"""Streaming requests use the same config path and should also forward raw schema mappings."""
|
||||
client, mock = _make_gemini_client()
|
||||
chunks = [_make_response([_make_part(text='{"answer": "hello"}')], finish_reason="STOP")]
|
||||
mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
stream = client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
stream=True,
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
config: types.GenerateContentConfig = mock.aio.models.generate_content_stream.call_args.kwargs["config"]
|
||||
assert config.response_mime_type == "application/json"
|
||||
assert config.response_schema == schema
|
||||
|
||||
|
||||
async def test_streaming_response_format_mapping_populates_final_value() -> None:
|
||||
"""Streaming responses should preserve mapping response_format for final value parsing."""
|
||||
client, mock = _make_gemini_client()
|
||||
chunks = [_make_response([_make_part(text='{"answer": "hello"}')], finish_reason="STOP")]
|
||||
mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks))
|
||||
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
stream = client.get_response(
|
||||
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
|
||||
options={"response_format": schema},
|
||||
stream=True,
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
final = await stream.get_final_response()
|
||||
assert final.value == {"answer": "hello"}
|
||||
|
||||
|
||||
async def test_streaming_response_format_passed_to_build_response_stream() -> None:
|
||||
"""Verifies that response_format is forwarded to _build_response_stream when streaming
|
||||
so that structured output parsing works correctly on the final assembled response.
|
||||
|
||||
Reference in New Issue
Block a user