test/response/test_callable.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
from dataclasses import dataclass
from typing import Annotated, Any
import pytest
from dirty_equals import IsPartialDict
from pydantic import Field
from ag2 import Context, Depends, Variable
from ag2.response import response_schema
from ag2.stream import MemoryStream
class TestNameDescription:
def test_name_from_function(self) -> None:
@response_schema
def my_parser(content: str) -> int:
return int(content)
assert my_parser.name == "my_parser"
def test_explicit_name(self) -> None:
@response_schema(name="custom")
def parse(content: str) -> int:
return int(content)
assert parse.name == "custom"
def test_description_from_docstring(self) -> None:
@response_schema
def parse(content: str) -> int:
"""Parse an integer."""
return int(content)
assert parse.description == "Parse an integer."
def test_explicit_description(self) -> None:
@response_schema(description="Custom description")
def parse(content: str) -> int:
return int(content)
assert parse.description == "Custom description"
def test_no_docstring_gives_empty_string(self) -> None:
@response_schema
async def parse(content: str) -> int:
return int(content)
assert parse.description is None
class TestSchemaGeneration:
def test_single_str_param_no_schema(self) -> None:
@response_schema
def parse(content: str) -> int:
return int(content)
assert parse.json_schema is None
def test_single_int_param_generates_schema(self) -> None:
@response_schema
def double(value: int) -> int:
return value * 2
assert double.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"data": IsPartialDict({
"description": 'Response with a one-field JSON `"{"data":...}"`',
"title": "Data",
"type": "integer",
}),
}),
"required": ["data"],
})
def test_single_dataclass_param_generates_schema(self) -> None:
@dataclass
class Point:
x: int
y: int
@response_schema
def process(point: Point) -> str:
return f"{point.x},{point.y}"
assert process.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"x": IsPartialDict({"type": "integer"}),
"y": IsPartialDict({"type": "integer"}),
}),
})
def test_single_union_param_generates_schema(self) -> None:
@response_schema
def process(value: int | str) -> str:
return str(value)
assert process.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"data": IsPartialDict({
"description": 'Response with a one-field JSON `"{"data":...}"`',
"title": "Data",
"anyOf": [
{"type": "integer"},
{"type": "string"},
],
}),
}),
"required": ["data"],
})
def test_multi_params_generates_object_schema(self) -> None:
@response_schema
def combine(name: str, age: int) -> str:
return f"{name} is {age}"
assert combine.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"name": IsPartialDict({"type": "string"}),
"age": IsPartialDict({"type": "integer"}),
}),
"required": ["name", "age"],
})
def test_multi_params_with_field_descriptions(self) -> None:
@response_schema
def process(
name: Annotated[str, Field(description="User name")],
score: float = Field(1.0, description="Test score"),
) -> str:
return f"{name}: {score}"
assert process.json_schema == IsPartialDict({
"properties": IsPartialDict({
"name": IsPartialDict({"description": "User name"}),
"score": IsPartialDict({"description": "Test score", "default": 1.0}),
}),
})
def test_multi_params_with_defaults(self) -> None:
@response_schema
def process(
name: str,
retries: int = 3,
) -> str:
return f"{name} x{retries}"
assert process.json_schema == IsPartialDict({
"required": ["name"],
"properties": IsPartialDict({
"retries": IsPartialDict({"default": 3}),
}),
})
def test_custom_fields_has_no_effect(self) -> None:
@response_schema
def combine(
name: str,
age: int,
# has no effect
ctx: Context,
var: Annotated[int, Variable()],
) -> str:
return f"{name} is {age}"
assert combine.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"name": IsPartialDict({"type": "string"}),
"age": IsPartialDict({"type": "integer"}),
}),
"required": ["name", "age"],
})
def test_respect_dependencies(self) -> None:
def get_age(age: int) -> int:
return age
@response_schema
def combine(
name: str,
age: Annotated[int, Depends(get_age)],
) -> str:
return f"{name} is {age}"
assert combine.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"name": IsPartialDict({"type": "string"}),
"age": IsPartialDict({"type": "integer"}),
}),
"required": ["name", "age"],
})
def test_custom_schema_overrides_generated(self) -> None:
custom = {"type": "object", "properties": {"x": {"type": "number"}}}
@response_schema(schema=custom)
def parse(content: str) -> float:
return float(content)
assert parse.json_schema == custom
class TestValidation:
@pytest.mark.asyncio()
async def test_sync_str_param(self) -> None:
@response_schema
def parse(content: str) -> int:
return int(content)
result = await parse.validate("42", context=None) # type: ignore[arg-type]
assert result == 42
@pytest.mark.asyncio()
async def test_async_str_param(self) -> None:
@response_schema
async def parse(content: str) -> int:
return int(content)
result = await parse.validate("42", context=None) # type: ignore[arg-type]
assert result == 42
@pytest.mark.asyncio()
async def test_single_int_param_deserializes(self) -> None:
@response_schema
def double(value: int) -> int:
return value * 2
result = await double.validate('{"data": 21}', context=None) # type: ignore[arg-type]
assert result == 42
@pytest.mark.asyncio()
async def test_single_dataclass_param_deserializes(self) -> None:
@dataclass
class Point:
x: int
y: int
@response_schema
def stringify(point: Point) -> str:
return f"{point.x},{point.y}"
result = await stringify.validate(
'{"x": 3, "y": 4}',
context=None, # type: ignore[arg-type]
)
assert result == "3,4"
@pytest.mark.asyncio()
async def test_single_union_param_deserializes(self) -> None:
@response_schema
def to_str(value: int | str) -> str:
return str(value)
assert await to_str.validate('{"data": 42}', context=None) == "42" # type: ignore[arg-type]
assert await to_str.validate('{"data": "hello"}', context=None) == "hello" # type: ignore[arg-type]
@pytest.mark.asyncio()
async def test_multi_params_with_defaults_uses_default(self) -> None:
@response_schema
def process(name: str, retries: int = 3) -> str:
return f"{name} x{retries}"
result = await process.validate(
'{"name": "test"}',
context=None, # type: ignore[arg-type]
)
assert result == "test x3"
@pytest.mark.asyncio()
async def test_multi_params_with_defaults_overridden(self) -> None:
@response_schema
def process(name: str, retries: int = 3) -> str:
return f"{name} x{retries}"
result = await process.validate(
'{"name": "test", "retries": 5}',
context=None, # type: ignore[arg-type]
)
assert result == "test x5"
@pytest.mark.asyncio()
async def test_multi_params_deserializes_json(self) -> None:
@response_schema
def greet(name: str, age: int) -> str:
return f"{name} is {age}"
result = await greet.validate(
'{"name": "Alice", "age": 30}',
context=None, # type: ignore[arg-type]
)
assert result == "Alice is 30"
@pytest.mark.asyncio()
async def test_async_multi_params(self) -> None:
@response_schema
async def greet(name: str, age: int) -> str:
return f"{name} is {age}"
result = await greet.validate(
'{"name": "Bob", "age": 25}',
context=None, # type: ignore[arg-type]
)
assert result == "Bob is 25"
class TestDependencyInjection:
@staticmethod
def _make_context(
*,
variables: dict[str, Any] | None = None,
dependencies: dict[Any, Any] | None = None,
) -> "Context":
return Context(
stream=MemoryStream(),
variables=variables or {},
dependencies=dependencies or {},
)
@pytest.mark.asyncio()
async def test_context_injected(self) -> None:
@response_schema
def parse(content: str, ctx: Context) -> dict[str, Any]:
return {"value": content, "has_vars": "key" in ctx.variables}
context = self._make_context(variables={"key": "val"})
result = await parse.validate("hello", context=context)
assert result == {"value": "hello", "has_vars": True}
@pytest.mark.asyncio()
async def test_variable_injected(self) -> None:
@response_schema
def parse(content: str, lang: Annotated[str, Variable()]) -> str:
return f"{content} ({lang})"
context = self._make_context(variables={"lang": "en"})
result = await parse.validate("hello", context=context)
assert result == "hello (en)"
@pytest.mark.asyncio()
async def test_variable_with_custom_name(self) -> None:
@response_schema
def parse(content: str, lang: Annotated[str, Variable("language")]) -> str:
return f"{content} ({lang})"
context = self._make_context(variables={"language": "fr"})
result = await parse.validate("hello", context=context)
assert result == "hello (fr)"
@pytest.mark.asyncio()
async def test_variable_with_default(self) -> None:
@response_schema
def parse(content: str, lang: Annotated[str, Variable(default="en")]) -> str:
return f"{content} ({lang})"
context = self._make_context()
result = await parse.validate("hello", context=context)
assert result == "hello (en)"
@pytest.mark.asyncio()
async def test_depends_injected(self) -> None:
def get_suffix() -> str:
return "!!!"
@response_schema
def parse(content: str, suffix: Annotated[str, Depends(get_suffix)]) -> str:
return content + suffix
context = self._make_context()
result = await parse.validate("hello", context=context)
assert result == "hello!!!"
@pytest.mark.asyncio()
async def test_all_di_together(self) -> None:
def get_separator() -> str:
return " | "
@response_schema
def parse(
content: str,
ctx: Context,
lang: Annotated[str, Variable()],
sep: Annotated[str, Depends(get_separator)],
) -> str:
prefix = "stream" if ctx.stream is not None else "no-stream"
return sep.join([prefix, content, lang])
context = self._make_context(variables={"lang": "en"})
result = await parse.validate("hello", context=context)
assert result == "stream | hello | en"
@pytest.mark.asyncio()
async def test_di_excluded_from_multi_param_schema(self) -> None:
@response_schema
def parse(
name: str,
age: int,
ctx: Context,
lang: Annotated[str, Variable(default="en")],
) -> str:
return f"{name}({age}) [{lang}]"
# DI params excluded from schema (multi-arg message body is a flat object, not ``data``-wrapped)
assert parse.json_schema == IsPartialDict({
"type": "object",
"properties": IsPartialDict({
"name": IsPartialDict({"type": "string"}),
"age": IsPartialDict({"type": "integer"}),
}),
"required": ["name", "age"],
})
# But resolved at runtime
context = self._make_context(variables={"lang": "fr"})
result = await parse.validate(
'{"name": "Alice", "age": 30}',
context=context,
)
assert result == "Alice(30) [fr]"