test/response/test_decision.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 json
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum, IntEnum
from typing import Annotated, Any
import pytest
from pydantic import Field, RootModel
pytest.importorskip("openai")
pytest.importorskip("typesafe_sdk")
import httpx2
from ag2 import Agent, ResponseSchema
from ag2.config import OpenAIDecisionsConfig, TypeSafeConfig
from ag2.config.openai import UnsupportedResponseSchemaError as OpenAIUnsupported
from ag2.config.typesafe import UnsupportedResponseSchemaError as TypeSafeUnsupported
from ag2.response import Question
USAGE = {
"input_tokens": 1,
"output_tokens": 0,
"total_tokens": 1,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
}
class Department(Enum):
"""Which team should handle this ticket?"""
BILLING = "billing"
"""Payments, invoicing, refunds."""
TECHNICAL = "technical"
"""Bugs, outages, integrations."""
class Plain(Enum):
A = "a"
B = "b"
class Severity(IntEnum):
"""How severe is this?"""
LOW = 0
"""Cosmetic."""
HIGH = 1
"""Outage."""
class Bare(IntEnum):
LOW = 0
HIGH = 1
class Reversed(IntEnum):
HIGH = 1
LOW = 0
class Refund(RootModel[Annotated[bool, Field(description="Is this a refund?")]]):
pass
UnreadableChoice = Enum("UnreadableChoice", {"A": "a", "B": "b"})
UnreadableScore = IntEnum("UnreadableScore", {"LOW": 0, "HIGH": 1})
@dataclass
class Asked:
"""What the model was asked, the same shape for both providers."""
instructions: str | None
options: dict[str, str | None]
class Provider(ABC):
"""A decision provider behind a fake HTTP API; subclasses supply the wire format."""
unsupported: type[Exception]
describe: str # config field that overrides option descriptions
def __init__(self) -> None:
self.questions: list[Any] = []
self.answer: dict[str, Any] = {}
self.client = httpx2.AsyncClient(transport=httpx2.MockTransport(self._handle))
@abstractmethod
def _handle(self, request: httpx2.Request) -> httpx2.Response: ...
@abstractmethod
def _config(self, **config: Any) -> Any: ...
@abstractmethod
def predicate(self, probability: float) -> dict[str, Any]: ...
@abstractmethod
def choice(self, value: str, options: tuple[str, ...]) -> dict[str, Any]: ...
@abstractmethod
def score(self, value: float, levels: int) -> dict[str, Any]: ...
@abstractmethod
def asked(self) -> Asked: ...
async def ask(self, schema: Any, answer: dict[str, Any], *, prompt: str | None = None, **config: Any) -> Any:
self.answer = answer
agent = Agent("a", prompt=prompt or [], config=self._config(**config), response_schema=schema)
return await (await agent.ask("hi")).content()
class OpenAIProvider(Provider):
unsupported = OpenAIUnsupported
describe = "descriptions"
def _handle(self, request: httpx2.Request) -> httpx2.Response:
self.questions.extend(json.loads(request.content)["questions"])
body = {"model": "m", "answers": [{"name": "answer", **self.answer}], "usage": USAGE}
return httpx2.Response(200, json=body)
def _config(self, **config: Any) -> Any:
return OpenAIDecisionsConfig(api_key="k", http_client=self.client, **config)
def predicate(self, probability: float) -> dict[str, Any]:
return {"type": "predicate", "probability": probability}
def choice(self, value: str, options: tuple[str, ...]) -> dict[str, Any]:
probabilities = [{"value": o, "probability": 1 / len(options)} for o in options]
return {"type": "choice", "choice": value, "confidence": 0.9, "probabilities": probabilities}
def score(self, value: float, levels: int) -> dict[str, Any]:
probabilities = [{"value": i, "label": str(i), "probability": 1 / levels} for i in range(levels)]
return {"type": "score", "score": value, "confidence": 0.5, "probabilities": probabilities}
def asked(self) -> Asked:
[question] = self.questions
entries = question.get("choices") or question.get("levels") or []
return Asked(
question["instructions"],
{str(o.get("value", i)): o.get("description") for i, o in enumerate(entries)},
)
class TypeSafeProvider(Provider):
unsupported = TypeSafeUnsupported
describe = "criteria"
def _handle(self, request: httpx2.Request) -> httpx2.Response:
self.questions.append(json.loads(request.content)["questions"]["answer"])
body = {"model": "m", "answers": {"answer": self.answer}, "usage": {"input_tokens": 1, "output_tokens": 0}}
return httpx2.Response(200, json=body)
def _config(self, **config: Any) -> Any:
return TypeSafeConfig(api_key="k", http_client=self.client, **config)
def predicate(self, probability: float) -> dict[str, Any]:
return {"type": "noul", "noul": probability}
def choice(self, value: str, options: tuple[str, ...]) -> dict[str, Any]:
probabilities = {o: 1 / len(options) for o in options}
return {"type": "choice", "choice": value, "confidence": 0.9, "probabilities": probabilities}
def score(self, value: float, levels: int) -> dict[str, Any]:
legend = {str(i): "x" for i in range(levels)}
probabilities = {str(i): 1 / levels for i in range(levels)}
return {"type": "score", "score": value, "confidence": 0.5, "legend": legend, "probabilities": probabilities}
def asked(self) -> Asked:
[question] = self.questions
criteria = question.get("criteria") or {}
if isinstance(criteria, list):
criteria = {str(i): c for i, c in enumerate(criteria)}
return Asked(question.get("instructions"), criteria)
@pytest.fixture(params=[OpenAIProvider, TypeSafeProvider], ids=["openai", "typesafe"])
def provider(request: pytest.FixtureRequest) -> Provider:
provider_type: type[Provider] = request.param
return provider_type()
@pytest.mark.asyncio
class TestQuestionMarker:
async def test_question_on_plain_bool(self, provider: Provider) -> None:
schema = Annotated[bool, Question("Is this a refund request?")]
assert await provider.ask(schema, provider.predicate(0.9)) is True
assert provider.asked() == Asked("Is this a refund request?", {})
async def test_options_on_third_party_enum(self, provider: Provider) -> None:
schema = Annotated[Plain, Question("Pick one", options={"a": "First", "b": "Second"})]
assert await provider.ask(schema, provider.choice("b", ("a", "b"))) is Plain.B
assert provider.asked() == Asked("Pick one", {"a": "First", "b": "Second"})
async def test_score_levels_keyed_by_int(self, provider: Provider) -> None:
schema = Annotated[Bare, Question("Rate it", options={0: "Fine", 1: "Broken"})]
assert await provider.ask(schema, provider.score(0.8, 2)) is Bare.HIGH
assert provider.asked() == Asked("Rate it", {"0": "Fine", "1": "Broken"})
async def test_question_without_options_keeps_docstring_options(self, provider: Provider) -> None:
schema = Annotated[Department, Question("Route it")]
await provider.ask(schema, provider.choice("billing", ("billing", "technical")))
assert provider.asked() == Asked(
"Route it", {"billing": "Payments, invoicing, refunds.", "technical": "Bugs, outages, integrations."}
)
@pytest.mark.asyncio
class TestPrecedence:
async def test_docstring_is_the_question(self, provider: Provider) -> None:
await provider.ask(Department, provider.choice("billing", ("billing", "technical")))
assert provider.asked().instructions == "Which team should handle this ticket?"
async def test_marker_beats_docstring(self, provider: Provider) -> None:
await provider.ask(
Annotated[Department, Question("From marker")], provider.choice("billing", ("billing", "technical"))
)
assert provider.asked().instructions == "From marker"
async def test_description_beats_marker(self, provider: Provider) -> None:
schema = ResponseSchema(Annotated[Department, Question("From marker")], description="From description")
await provider.ask(schema, provider.choice("billing", ("billing", "technical")))
assert provider.asked().instructions == "From description"
async def test_marker_options_beat_docstrings(self, provider: Provider) -> None:
schema = Annotated[Department, Question(options={"billing": "Money"})]
await provider.ask(schema, provider.choice("billing", ("billing", "technical")))
assert provider.asked().options == {"billing": "Money", "technical": "Bugs, outages, integrations."}
async def test_config_beats_marker_options(self, provider: Provider) -> None:
schema = Annotated[Department, Question(options={"billing": "Money", "technical": "Code"})]
await provider.ask(
schema, provider.choice("billing", ("billing", "technical")), **{provider.describe: {"billing": "Cash"}}
)
assert provider.asked().options == {"billing": "Cash", "technical": "Code"}
async def test_prompt_and_question_are_joined(self, provider: Provider) -> None:
await provider.ask(Annotated[bool, Question("Is it?")], provider.predicate(0.1), prompt="Be strict.")
assert provider.asked().instructions == "Be strict.\n\nIs it?"
@pytest.mark.asyncio
class TestAnswers:
@pytest.mark.parametrize(("probability", "expected"), [(0.7, True), (0.2, False)])
async def test_bool_is_thresholded(self, provider: Provider, probability: float, expected: bool) -> None:
assert await provider.ask(Annotated[bool, Question("?")], provider.predicate(probability)) is expected
async def test_probability_is_returned_raw(self, provider: Provider) -> None:
schema = ResponseSchema.from_schema({"type": "number", "minimum": 0, "maximum": 1}, name="p", description="?")
assert await provider.ask(schema, provider.predicate(0.37)) == "0.37"
@pytest.mark.parametrize(("raw", "expected"), [(0.4, Severity.LOW), (0.6, Severity.HIGH), (9.0, Severity.HIGH)])
async def test_score_snaps_and_clamps(self, provider: Provider, raw: float, expected: Severity) -> None:
assert await provider.ask(Severity, provider.score(raw, 2)) is expected
async def test_envelope_does_not_change_the_question(self, provider: Provider) -> None:
wrapped = ResponseSchema(bool, description="Is it?")
bare = ResponseSchema(bool, description="Is it?", embed=False)
assert await provider.ask(wrapped, provider.predicate(0.9)) is True
first = provider.asked()
provider.questions.clear()
assert await provider.ask(bare, provider.predicate(0.9)) is True
assert provider.asked() == first
@pytest.mark.asyncio
class TestRejections:
@pytest.mark.parametrize("schema", [str, int, ResponseSchema(str, description="?")])
async def test_non_decision_schema(self, provider: Provider, schema: Any) -> None:
with pytest.raises(provider.unsupported, match="decision-only"):
await provider.ask(schema, provider.predicate(0.5), prompt="Q")
assert provider.questions == []
async def test_bare_bool_without_prompt(self, provider: Provider) -> None:
with pytest.raises(ValueError, match="question|asking"):
await provider.ask(bool, provider.predicate(0.5))
assert provider.questions == []
@pytest.mark.asyncio
class TestUnreadableSource:
async def test_choice_proceeds_without_descriptions(self, provider: Provider) -> None:
assert await provider.ask(UnreadableChoice, provider.choice("a", ("a", "b")), prompt="Q") is (
UnreadableChoice.A
)
assert provider.asked().options == {"a": None, "b": None}
async def test_marker_describes_unreadable_choice(self, provider: Provider) -> None:
schema = Annotated[UnreadableChoice, Question(options={"a": "First"})]
await provider.ask(schema, provider.choice("a", ("a", "b")), prompt="Q")
assert provider.asked().options == {"a": "First", "b": None}
async def test_score_without_descriptions(self) -> None:
provider = OpenAIProvider()
assert await provider.ask(UnreadableScore, provider.score(0.9, 2), prompt="Q") is UnreadableScore.HIGH
assert provider.asked().options == {"0": None, "1": None}
async def test_typesafe_score_requires_descriptions_and_names_the_way_out(self) -> None:
provider = TypeSafeProvider()
with pytest.raises(TypeSafeUnsupported, match=r"Question\(options=") as e:
await provider.ask(UnreadableScore, provider.score(0.9, 2), prompt="Q")
assert "criteria" in str(e.value)
assert provider.questions == []
async def test_marker_makes_typesafe_score_succeed(self) -> None:
provider = TypeSafeProvider()
schema = Annotated[UnreadableScore, Question(options={0: "Fine", 1: "Broken"})]
assert await provider.ask(schema, provider.score(0.9, 2), prompt="Q") is UnreadableScore.HIGH
assert provider.asked().options == {"0": "Fine", "1": "Broken"}
@pytest.mark.asyncio
class TestEdges:
async def test_description_equal_to_docstring_still_beats_marker(self, provider: Provider) -> None:
schema = ResponseSchema(Annotated[Department, Question("From marker")], description=Department.__doc__)
await provider.ask(schema, provider.choice("billing", ("billing", "technical")))
assert provider.asked().instructions == "Which team should handle this ticket?"
@pytest.mark.parametrize(
"schema",
[
pytest.param(Annotated[Department, Question(options={"nope": "x"})], id="unknown-choice"),
pytest.param(Annotated[Severity, Question(options={"0": "zero"})], id="score-as-string"),
pytest.param(Annotated[bool, Question("?", options={"true": "x"})], id="predicate"),
],
)
async def test_option_keys_must_name_options(self, provider: Provider, schema: Any) -> None:
with pytest.raises(ValueError, match=r"Question\(options="):
await provider.ask(schema, provider.predicate(0.5), prompt="Q")
assert provider.questions == []
async def test_score_declared_out_of_order_is_sent_by_value(self, provider: Provider) -> None:
schema = Annotated[Reversed, Question("Rate it", options={0: "Fine", 1: "Broken"})]
assert await provider.ask(schema, provider.score(0.8, 2)) is Reversed.HIGH
assert list(provider.asked().options.items()) == [("0", "Fine"), ("1", "Broken")]
async def test_question_lifted_from_pydantic_schema(self, provider: Provider) -> None:
assert await provider.ask(Refund, provider.predicate(0.9)) == Refund(True)
assert provider.asked() == Asked("Is this a refund?", {})
async def test_bool_key_does_not_name_a_score_level(self, provider: Provider) -> None:
schema = Annotated[Bare, Question(options={True: "x"})]
with pytest.raises(ValueError, match=r"Question\(options="):
await provider.ask(schema, provider.score(0.5, 2), prompt="Q")
assert provider.questions == []
async def test_annotated_str_is_not_a_decision(self, provider: Provider) -> None:
with pytest.raises(provider.unsupported, match="not a decision type"):
await provider.ask(Annotated[str, Question("?")], provider.predicate(0.5), prompt="Q")