test/response/test_schema.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
import warnings
from dataclasses import dataclass
from enum import Enum
from typing import Annotated, Any, NamedTuple, Union
import pytest
from dirty_equals import IsPartialDict
from pydantic import BaseModel, Field, ValidationError
from typing_extensions import TypedDict
from ag2 import Context, MemoryStream
from ag2.response import ResponseSchema
from ag2.response.schema import RawSchema
from ag2.types import ClassInfo
class TestEmbeddedTypes:
@pytest.mark.parametrize(
("type_", "schema"),
[
pytest.param(int, {"type": "integer"}, id="int"),
pytest.param(float, {"type": "number"}, id="float"),
pytest.param(bool, {"type": "boolean"}, id="bool"),
pytest.param(
[int, float],
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="[T1, T2]",
),
pytest.param(
(int, float),
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="(T1, T2)",
),
pytest.param(
int | float,
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="typing.Union[T1, T2]",
),
pytest.param(
Union[int, float], # noqa: UP007
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="T1 | T2",
),
pytest.param(
list[int],
{"type": "array", "items": {"type": "integer"}},
id="list[T]",
),
],
)
def test_primitive_types_not_embedded(
self,
type_: ClassInfo,
schema: dict[str, Any],
) -> None:
response_schema = ResponseSchema(type_, embed=False)
assert not response_schema._embedded_type
assert response_schema.json_schema == schema
@pytest.mark.parametrize(
("type_", "schema"),
[
pytest.param(int, {"type": "integer"}, id="int"),
pytest.param(float, {"type": "number"}, id="float"),
pytest.param(bool, {"type": "boolean"}, id="bool"),
pytest.param(
[int, float],
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="[T1, T2]",
),
pytest.param(
(int, float),
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="(T1, T2)",
),
pytest.param(
int | float,
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="typing.Union[T1, T2]",
),
pytest.param(
Union[int, float], # noqa: UP007
{"anyOf": [{"type": "integer"}, {"type": "number"}]},
id="T1 | T2",
),
pytest.param(
list[int],
{"type": "array", "items": {"type": "integer"}},
id="list[T]",
),
],
)
def test_primitive_types_embedded(
self,
type_: ClassInfo,
schema: dict[str, Any],
) -> None:
response_schema = ResponseSchema(type_)
assert response_schema._embedded_type
assert response_schema.json_schema == {
"properties": {
"data": {
"description": 'Response with a one-field JSON `"{"data":...}"`',
"title": "Data",
**schema,
}
},
"required": ["data"],
"type": "object",
}
def test_str_has_no_schema(self) -> None:
response_schema = ResponseSchema(str)
assert response_schema._embedded_type is True
assert response_schema.json_schema is None
def test_dict_not_embedded(self) -> None:
response_schema = ResponseSchema(dict[str, int])
assert response_schema._embedded_type is False
assert response_schema.json_schema == {
"type": "object",
"additionalProperties": {"type": "integer"},
}
def test_dataclass_not_embedded(self) -> None:
@dataclass
class Model:
field: str
schema = ResponseSchema(Model)
assert not schema._embedded_type
def test_pydantic_model_not_embedded(self) -> None:
class Model(BaseModel):
field: str
schema = ResponseSchema(Model)
assert not schema._embedded_type
def test_typed_dict_not_embedded(self) -> None:
class Model(TypedDict):
field: str
schema = ResponseSchema(Model)
assert not schema._embedded_type
class TestDataclassSchemas:
def test_simple_dataclass(self) -> None:
@dataclass
class User:
name: str
age: int
schema = ResponseSchema(User)
assert schema.name == "User"
# dataclasses get auto-generated docstrings
assert schema.description == User.__doc__
assert schema.json_schema == {
"properties": {
"name": {"title": "Name", "type": "string"},
"age": {"title": "Age", "type": "integer"},
},
"required": ["name", "age"],
"type": "object",
}
def test_dataclass_with_docstring(self) -> None:
@dataclass
class Response:
"""Structured response from the model."""
value: int
schema = ResponseSchema(Response)
assert schema.name == "Response"
assert schema.description == "Structured response from the model."
def test_dataclass_with_defaults(self) -> None:
@dataclass
class Config:
mode: str
retries: int = 3
schema = ResponseSchema(Config)
assert schema.name == "Config"
assert schema.json_schema == IsPartialDict({
"required": ["mode"],
"properties": IsPartialDict({"retries": IsPartialDict({"default": 3})}),
})
def test_nested_dataclass(self) -> None:
@dataclass
class Address:
city: str
zip_code: str
@dataclass
class Person:
name: str
address: Address
schema = ResponseSchema(Person)
assert schema.name == "Person"
assert schema.json_schema == IsPartialDict({
"$defs": IsPartialDict({"Address": IsPartialDict({"type": "object"})}),
"properties": IsPartialDict({"address": {"$ref": "#/$defs/Address"}}),
})
class TestPydanticModelSchemas:
def test_simple_model(self) -> None:
class Item(BaseModel):
name: str
price: float
schema = ResponseSchema(Item)
assert schema.name == "Item"
assert schema.json_schema == {
"properties": {
"name": {"title": "Name", "type": "string"},
"price": {"title": "Price", "type": "number"},
},
"required": ["name", "price"],
"type": "object",
}
def test_model_with_field_descriptions(self) -> None:
class Person(BaseModel):
name: str = Field(description="Full name")
age: Annotated[int, Field(description="Age in years")]
schema = ResponseSchema(Person)
assert schema.json_schema == IsPartialDict({
"properties": IsPartialDict({
"name": IsPartialDict({"description": "Full name"}),
"age": IsPartialDict({"description": "Age in years"}),
}),
})
def test_model_with_docstring(self) -> None:
class Result(BaseModel):
"""Analysis result."""
score: float
schema = ResponseSchema(Result)
assert schema.description == "Analysis result."
def test_model_with_nested_model(self) -> None:
class Coord(BaseModel):
x: float
y: float
class Shape(BaseModel):
origin: Coord
label: str
schema = ResponseSchema(Shape)
assert schema.name == "Shape"
assert schema.json_schema == IsPartialDict({
"$defs": IsPartialDict({"Coord": IsPartialDict({"type": "object"})}),
})
def test_model_with_enum_field(self) -> None:
class Color(Enum):
RED = "red"
GREEN = "green"
BLUE = "blue"
class Palette(BaseModel):
primary: Color
schema = ResponseSchema(Palette)
assert schema.name == "Palette"
assert schema.json_schema == IsPartialDict({
"$defs": {
"Color": {
"enum": ["red", "green", "blue"],
"title": "Color",
"type": "string",
}
},
})
def test_model_with_annotated_constraints(self) -> None:
class Bounded(BaseModel):
value: Annotated[int, Field(ge=0, le=100)]
schema = ResponseSchema(Bounded)
assert schema.json_schema == IsPartialDict({
"properties": IsPartialDict({
"value": IsPartialDict({"minimum": 0, "maximum": 100}),
}),
})
class TestUnionSchemas:
@pytest.mark.parametrize(
("type_", "expected_any_of"),
[
pytest.param(
int | str,
[{"type": "integer"}, {"type": "string"}],
id="union_operator",
),
pytest.param(
(int, str, float),
[{"type": "integer"}, {"type": "string"}, {"type": "number"}],
id="tuple_of_types",
),
],
)
def test_union_schema(self, type_: ClassInfo, expected_any_of: list) -> None: # type: ignore[type-arg]
schema = ResponseSchema(type_, name="Union", embed=False)
assert schema.name == "Union"
assert schema.json_schema == {"anyOf": expected_any_of}
def test_union_with_dataclass(self) -> None:
@dataclass
class Error:
message: str
@dataclass
class Success:
value: int
schema = ResponseSchema((Error, Success), name="Result", embed=False)
assert schema.name == "Result"
assert schema.json_schema == IsPartialDict({
"$defs": IsPartialDict({
"Error": IsPartialDict({}),
"Success": IsPartialDict({}),
}),
})
@pytest.mark.asyncio
async def test_named_tuple_class_is_a_type_not_a_union_of_members() -> None:
class Point(NamedTuple):
x: int
y: int
schema = ResponseSchema(Point)
assert await schema.validate('{"data": [1, 2]}', Context(MemoryStream())) == Point(1, 2)
class TestNameDescription:
def test_explicit_name_overrides_title(self) -> None:
class MyModel(BaseModel):
x: int
schema = ResponseSchema(MyModel, name="CustomName")
assert schema.name == "CustomName"
def test_name_from_class_title(self) -> None:
class MyModel(BaseModel):
x: int
schema = ResponseSchema(MyModel)
assert schema.name == "MyModel"
def test_fallback_name(self) -> None:
schema = ResponseSchema(int)
# int schema has no "title" key, so falls back to "ResponseSchema"
assert schema.name == "ResponseSchema"
def test_explicit_description(self) -> None:
class MyModel(BaseModel):
x: int
schema = ResponseSchema(MyModel, description="Custom description")
assert schema.description == "Custom description"
def test_schema_description_overrides_docstring(self) -> None:
class MyModel(BaseModel):
"""Model docstring."""
x: int
# docstring takes priority over explicit description
schema = ResponseSchema(MyModel, description="Custom description")
assert schema.description == "Custom description"
def test_metadata_popped_from_schema(self) -> None:
class MyModel(BaseModel):
"""Some description."""
x: int
schema = ResponseSchema(MyModel)
assert "title" not in schema.json_schema
assert "description" not in schema.json_schema
class TestEnsureSchema:
def test_none_returns_none(self) -> None:
assert ResponseSchema.ensure_schema(None) is None
def test_response_proto_returned_as_is(self) -> None:
original = ResponseSchema(int, name="Test")
result = ResponseSchema.ensure_schema(original)
assert result is original
def test_type_wrapped_in_response_schema(self) -> None:
result = ResponseSchema.ensure_schema(int)
assert isinstance(result, ResponseSchema)
assert result.json_schema == {
"properties": {
"data": {
"description": 'Response with a one-field JSON `"{"data":...}"`',
"title": "Data",
"type": "integer",
}
},
"required": ["data"],
"type": "object",
}
def test_union_type_wrapped(self) -> None:
result = ResponseSchema.ensure_schema(int | str)
assert isinstance(result, ResponseSchema)
assert result.json_schema == {
"properties": {
"data": {
"description": 'Response with a one-field JSON `"{"data":...}"`',
"title": "Data",
"anyOf": [
{"type": "integer"},
{"type": "string"},
],
}
},
"required": ["data"],
"type": "object",
}
class TestRawSchema:
def test_creation(self) -> None:
raw = ResponseSchema.from_schema(
{"type": "object", "properties": {"x": {"type": "integer"}}},
name="Custom",
description="A custom schema",
)
assert isinstance(raw, RawSchema)
assert raw.name == "Custom"
assert raw.description == "A custom schema"
assert raw.json_schema == {"type": "object", "properties": {"x": {"type": "integer"}}}
def test_creation_without_description(self) -> None:
raw = ResponseSchema.from_schema({"type": "string"}, name="Simple")
assert raw.description is None
@pytest.mark.asyncio()
async def test_validate_returns_raw_string_with_warning(self) -> None:
raw = ResponseSchema.from_schema({"type": "string"}, name="Raw")
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
result = await raw.validate("hello", context=None) # type: ignore[arg-type]
assert result == "hello"
assert len(w) == 1
assert issubclass(w[0].category, RuntimeWarning)
assert "can't validate" in str(w[0].message)
@pytest.mark.asyncio()
class TestValidation:
@pytest.mark.parametrize(
("type_", "json_input", "expected"),
[
pytest.param(int, '{"data": 42}', 42, id="int"),
pytest.param(float, '{"data": 3.14}', 3.14, id="float"),
pytest.param(bool, '{"data": true}', True, id="bool"),
],
)
async def test_validate_primitive(self, type_: type, json_input: str, expected: object) -> None:
schema = ResponseSchema(type_)
result = await schema.validate(json_input, context=None) # type: ignore[arg-type]
assert result == expected
@pytest.mark.parametrize(
("type_", "json_input", "expected"),
[
pytest.param(int, "42", 42, id="int"),
pytest.param(float, "3.14", 3.14, id="float"),
pytest.param(bool, "true", True, id="bool"),
],
)
async def test_validate_not_embedded_primitive(self, type_: type, json_input: str, expected: object) -> None:
schema = ResponseSchema(type_, embed=False)
result = await schema.validate(json_input, context=None) # type: ignore[arg-type]
assert result == expected
async def test_validate_union(self) -> None:
schema = ResponseSchema(int | str)
assert await schema.validate('{"data": 42}', context=None) == 42 # type: ignore[arg-type]
assert await schema.validate('{"data": "hello"}', context=None) == "hello" # type: ignore[arg-type]
async def test_validate_not_embedded_union(self) -> None:
schema = ResponseSchema(int | str, embed=False)
assert await schema.validate("42", context=None) == 42 # type: ignore[arg-type]
assert await schema.validate('"hello"', context=None) == "hello" # type: ignore[arg-type]
async def test_validate_dataclass(self) -> None:
@dataclass
class Point:
x: int
y: int
schema = ResponseSchema(Point)
result = await schema.validate('{"x": 1, "y": 2}', context=None) # type: ignore[arg-type]
assert result == Point(x=1, y=2)
async def test_validate_pydantic_model(self) -> None:
class Item(BaseModel):
name: str
price: float
schema = ResponseSchema(Item)
result = await schema.validate('{"name": "Widget", "price": 9.99}', context=None) # type: ignore[arg-type]
assert result == Item(name="Widget", price=9.99)
async def test_validate_invalid_json_raises(self) -> None:
schema = ResponseSchema(int, embed=False)
with pytest.raises(Exception):
await schema.validate("not a number", context=None) # type: ignore[arg-type]
@pytest.mark.asyncio()
class TestAnnotatedValidation:
@pytest.mark.parametrize(
("type_", "raw", "expected"),
[
pytest.param(Annotated[bool, "note"], '{"data": true}', True, id="bool"),
pytest.param(Annotated[int, "note"], '{"data": 3}', 3, id="int"),
pytest.param(Annotated[list[int], "note"], '{"data": [1, 2]}', [1, 2], id="list"),
pytest.param(Annotated[dict[str, int], "note"], '{"a": 1}', {"a": 1}, id="dict"),
pytest.param(Annotated[int | str, "note"], '{"data": "x"}', "x", id="union"),
],
)
async def test_validates(self, type_: ClassInfo, raw: str, expected: Any) -> None:
schema = ResponseSchema(type_)
assert await schema.validate(raw, context=None) == expected # type: ignore[arg-type]
async def test_annotated_str_keeps_its_constraints(self) -> None:
schema = ResponseSchema(Annotated[str, Field(min_length=3)], embed=False)
with pytest.raises(ValidationError):
await schema.validate('"x"', context=None) # type: ignore[arg-type]
assert await schema.validate('"xyz"', context=None) == "xyz" # type: ignore[arg-type]
async def test_enum(self) -> None:
class Color(Enum):
RED = "red"
schema = ResponseSchema(Annotated[Color, "note"])
assert await schema.validate('{"data": "red"}', context=None) is Color.RED # type: ignore[arg-type]
async def test_dataclass(self) -> None:
@dataclass
class Point:
x: int
schema = ResponseSchema(Annotated[Point, "note"])
assert await schema.validate('{"x": 1}', context=None) == Point(1) # type: ignore[arg-type]
async def test_model(self) -> None:
class Point(BaseModel):
x: int
schema = ResponseSchema(Annotated[Point, "note"])
assert await schema.validate('{"x": 1}', context=None) == Point(x=1) # type: ignore[arg-type]
class TestAnnotatedSchemas:
def test_dataclass_schema_matches_bare_type(self) -> None:
@dataclass
class Point:
x: int
assert ResponseSchema(Annotated[Point, "note"]).json_schema == ResponseSchema(Point).json_schema
def test_primitive_schema_matches_bare_type(self) -> None:
assert ResponseSchema(Annotated[int, "note"]).json_schema == ResponseSchema(int).json_schema
def test_description_matches_bare_type(self) -> None:
assert ResponseSchema(Annotated[bool, "note"]).description == ResponseSchema(bool).description