wehub-resource-sync
235 lines
8.2 KiB
Python
235 lines
8.2 KiB
Python
import json
|
|
from typing import Any, Literal, cast
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
from typing_extensions import TypedDict
|
|
|
|
from agents import (
|
|
Agent,
|
|
AgentOutputSchema,
|
|
AgentOutputSchemaBase,
|
|
ModelBehaviorError,
|
|
UserError,
|
|
)
|
|
from agents.agent_output import _WRAPPER_DICT_KEY
|
|
from agents.run_internal.run_loop import get_output_schema
|
|
from agents.util import _json
|
|
|
|
|
|
def test_plain_text_output():
|
|
agent = Agent(name="test")
|
|
output_schema = get_output_schema(agent)
|
|
assert not output_schema, "Shouldn't have an output tool config without an output type"
|
|
|
|
agent = Agent(name="test", output_type=str)
|
|
assert not output_schema, "Shouldn't have an output tool config with str output type"
|
|
|
|
|
|
class Foo(BaseModel):
|
|
bar: str
|
|
|
|
|
|
def test_structured_output_pydantic():
|
|
agent = Agent(name="test", output_type=Foo)
|
|
output_schema = get_output_schema(agent)
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
|
|
assert isinstance(output_schema, AgentOutputSchema)
|
|
assert output_schema.output_type == Foo, "Should have the correct output type"
|
|
assert not output_schema._is_wrapped, "Pydantic objects should not be wrapped"
|
|
for key, value in Foo.model_json_schema().items():
|
|
assert output_schema.json_schema()[key] == value
|
|
|
|
json_str = Foo(bar="baz").model_dump_json()
|
|
validated = output_schema.validate_json(json_str)
|
|
assert validated == Foo(bar="baz")
|
|
|
|
|
|
class Bar(TypedDict):
|
|
bar: str
|
|
|
|
|
|
def test_structured_output_typed_dict():
|
|
agent = Agent(name="test", output_type=Bar)
|
|
output_schema = get_output_schema(agent)
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
assert isinstance(output_schema, AgentOutputSchema)
|
|
assert output_schema.output_type == Bar, "Should have the correct output type"
|
|
assert not output_schema._is_wrapped, "TypedDicts should not be wrapped"
|
|
|
|
json_str = json.dumps(Bar(bar="baz"))
|
|
validated = output_schema.validate_json(json_str)
|
|
assert validated == Bar(bar="baz")
|
|
|
|
|
|
def test_structured_output_list():
|
|
agent = Agent(name="test", output_type=list[str])
|
|
output_schema = get_output_schema(agent)
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
assert isinstance(output_schema, AgentOutputSchema)
|
|
assert output_schema.output_type == list[str], "Should have the correct output type"
|
|
assert output_schema._is_wrapped, "Lists should be wrapped"
|
|
|
|
# This is testing implementation details, but it's useful to make sure this doesn't break
|
|
json_str = json.dumps({_WRAPPER_DICT_KEY: ["foo", "bar"]})
|
|
validated = output_schema.validate_json(json_str)
|
|
assert validated == ["foo", "bar"]
|
|
|
|
|
|
def test_structured_output_literal_name_handles_literal_values():
|
|
output_schema = AgentOutputSchema(output_type=cast(type[Any], Literal["ok"]))
|
|
|
|
assert output_schema.name() == "Literal['ok']"
|
|
|
|
|
|
def test_structured_output_nested_literal_name_handles_literal_values():
|
|
output_schema = AgentOutputSchema(output_type=list[Literal["ok", "done"]])
|
|
|
|
assert output_schema.name() == "list[Literal['ok', 'done']]"
|
|
|
|
|
|
def test_structured_output_generic_dict_is_not_wrapped():
|
|
output_schema = AgentOutputSchema(output_type=dict[str, int], strict_json_schema=False)
|
|
assert output_schema.output_type == dict[str, int]
|
|
assert not output_schema._is_wrapped, "Generic dict output should not be wrapped"
|
|
assert "response" not in output_schema.json_schema().get("properties", {})
|
|
|
|
validated = output_schema.validate_json(json.dumps({"foo": 1}))
|
|
assert validated == {"foo": 1}
|
|
|
|
|
|
def test_structured_output_generic_dict_rejects_wrapper_shape():
|
|
output_schema = AgentOutputSchema(output_type=dict[str, int], strict_json_schema=False)
|
|
|
|
with pytest.raises(ModelBehaviorError):
|
|
output_schema.validate_json(json.dumps({"response": {"foo": 1}}))
|
|
|
|
|
|
def test_bad_json_raises_error(mocker):
|
|
agent = Agent(name="test", output_type=Foo)
|
|
output_schema = get_output_schema(agent)
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
|
|
with pytest.raises(ModelBehaviorError):
|
|
output_schema.validate_json("not valid json")
|
|
|
|
agent = Agent(name="test", output_type=list[str])
|
|
output_schema = get_output_schema(agent)
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
|
|
mock_validate_json = mocker.patch.object(_json, "validate_json")
|
|
mock_validate_json.return_value = ["foo"]
|
|
|
|
with pytest.raises(ModelBehaviorError):
|
|
output_schema.validate_json(json.dumps(["foo"]))
|
|
|
|
mock_validate_json.return_value = {"value": "foo"}
|
|
|
|
with pytest.raises(ModelBehaviorError):
|
|
output_schema.validate_json(json.dumps(["foo"]))
|
|
|
|
|
|
def test_plain_text_obj_doesnt_produce_schema():
|
|
output_wrapper = AgentOutputSchema(output_type=str)
|
|
with pytest.raises(UserError):
|
|
output_wrapper.json_schema()
|
|
|
|
|
|
def test_structured_output_is_strict():
|
|
output_wrapper = AgentOutputSchema(output_type=Foo)
|
|
assert output_wrapper.is_strict_json_schema()
|
|
for key, value in Foo.model_json_schema().items():
|
|
assert output_wrapper.json_schema()[key] == value
|
|
|
|
assert (
|
|
"additionalProperties" in output_wrapper.json_schema()
|
|
and not output_wrapper.json_schema()["additionalProperties"]
|
|
)
|
|
|
|
|
|
def test_setting_strict_false_works():
|
|
output_wrapper = AgentOutputSchema(output_type=Foo, strict_json_schema=False)
|
|
assert not output_wrapper.is_strict_json_schema()
|
|
assert output_wrapper.json_schema() == Foo.model_json_schema()
|
|
assert output_wrapper.json_schema() == Foo.model_json_schema()
|
|
|
|
|
|
_CUSTOM_OUTPUT_SCHEMA_JSON_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"foo": {"type": "string"},
|
|
},
|
|
"required": ["foo"],
|
|
}
|
|
|
|
|
|
class CustomOutputSchema(AgentOutputSchemaBase):
|
|
def is_plain_text(self) -> bool:
|
|
return False
|
|
|
|
def name(self) -> str:
|
|
return "FooBarBaz"
|
|
|
|
def json_schema(self) -> dict[str, Any]:
|
|
return _CUSTOM_OUTPUT_SCHEMA_JSON_SCHEMA
|
|
|
|
def is_strict_json_schema(self) -> bool:
|
|
return False
|
|
|
|
def validate_json(self, json_str: str) -> Any:
|
|
return ["some", "output"]
|
|
|
|
|
|
def test_custom_output_schema():
|
|
custom_output_schema = CustomOutputSchema()
|
|
agent = Agent(name="test", output_type=custom_output_schema)
|
|
output_schema = get_output_schema(agent)
|
|
|
|
assert output_schema, "Should have an output tool config with a structured output type"
|
|
assert isinstance(output_schema, CustomOutputSchema)
|
|
assert output_schema.json_schema() == _CUSTOM_OUTPUT_SCHEMA_JSON_SCHEMA
|
|
assert not output_schema.is_strict_json_schema()
|
|
assert not output_schema.is_plain_text()
|
|
|
|
json_str = json.dumps({"foo": "bar"})
|
|
validated = output_schema.validate_json(json_str)
|
|
assert validated == ["some", "output"]
|
|
|
|
|
|
class StrictOutput(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
|
|
def test_agent_output_schema_strict_rejects_type_coercion():
|
|
"""With strict_json_schema=True (default), string input for an int field must raise
|
|
ModelBehaviorError instead of being silently coerced."""
|
|
schema = AgentOutputSchema(output_type=StrictOutput, strict_json_schema=True)
|
|
assert schema.is_strict_json_schema()
|
|
|
|
# age is a string "25" — strict mode should reject this
|
|
malformed_json = '{"name": "Alice", "age": "25"}'
|
|
with pytest.raises(ModelBehaviorError, match="Invalid JSON"):
|
|
schema.validate_json(malformed_json)
|
|
|
|
# Correctly typed input should still be accepted
|
|
valid_json = '{"name": "Alice", "age": 25}'
|
|
result = schema.validate_json(valid_json)
|
|
assert result.name == "Alice"
|
|
assert result.age == 25
|
|
|
|
|
|
def test_agent_output_schema_lenient_allows_type_coercion():
|
|
"""With strict_json_schema=False, Pydantic's default lenient mode silently coerces
|
|
string input for an int field — verifying backward compatibility."""
|
|
schema = AgentOutputSchema(output_type=StrictOutput, strict_json_schema=False)
|
|
assert not schema.is_strict_json_schema()
|
|
|
|
# age is a string "25" — lenient mode should coerce it to int 25
|
|
coerced_json = '{"name": "Alice", "age": "25"}'
|
|
result = schema.validate_json(coerced_json)
|
|
assert result.name == "Alice"
|
|
assert result.age == 25
|
|
assert isinstance(result.age, int)
|