test_agent_model_settings.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. import pytest
  2. from agents import ModelSettings
  3. from pydantic import BaseModel, Field
  4. from agency_swarm import Agent
  5. @pytest.mark.asyncio
  6. async def test_agent_structured_response_output_type():
  7. """Agent should follow a Pydantic response schema (output_type) without mocking Runner.
  8. This uses the real agent execution path and asserts the final_output can be parsed to the
  9. declared Pydantic model, verifying schema-conformant responses.
  10. """
  11. class GreetingSchema(BaseModel):
  12. greeting: str = Field(..., description="Greeting text")
  13. recipient: str = Field(..., description="Who is greeted")
  14. num_messages: int = Field(..., description="Number of messages in your conversation history")
  15. agent = Agent(
  16. name="SchemaAgent",
  17. instructions=(
  18. "When asked to greet someone, respond ONLY as a strict JSON object matching the schema: "
  19. "{greeting: string, recipient: string}. Do not include any extra text."
  20. ),
  21. output_type=GreetingSchema,
  22. model_settings=ModelSettings(temperature=0.0),
  23. )
  24. # Ask the agent to greet a recipient; rely on the model+instructions to produce structured JSON
  25. result = await agent.get_response("Hello, my name is John")
  26. print(f"result: {result.final_output}")
  27. print(f"result type: {type(result.final_output)}")
  28. assert result is not None and not isinstance(result.final_output, str)
  29. # Parse the output as JSON and validate with the schema
  30. assert (
  31. isinstance(result.final_output.greeting, str)
  32. and isinstance(result.final_output.recipient, str)
  33. and isinstance(result.final_output.num_messages, int)
  34. )
  35. @pytest.mark.asyncio
  36. async def test_max_tokens_limits_output_length():
  37. """Agent should respect max_tokens by producing a very short response.
  38. We request a ~500-word poem but set max_tokens=16 and verify the output
  39. is significantly shorter than the requested length.
  40. """
  41. agent = Agent(
  42. name="TokenLimitAgent",
  43. instructions=("Respond to the user's request. Keep your answer within the model's limits."),
  44. model_settings=ModelSettings(temperature=0.0, max_tokens=16),
  45. )
  46. prompt = (
  47. "Please write a 500-word poem about the changing seasons, rich imagery, varied meter, "
  48. "and vivid emotions. Avoid bullet points; produce continuous verse."
  49. )
  50. result = await agent.get_response(prompt)
  51. assert result is not None and isinstance(result.final_output, str)
  52. text = result.final_output.strip()
  53. # Ensure we got something back
  54. assert len(text) > 0
  55. # Heuristic: with max_tokens=16, response should be very short compared to 500 words
  56. word_count = len(text.split())
  57. assert word_count < 80, f"Expected a truncated response due to low max_tokens; got ~{word_count} words"