Skip to content

Conversation

@devin-ai-integration
Copy link
Contributor

Fix Union type handling in Pydantic output generation

Summary

Fixed a bug in generate_model_description where all Union types were incorrectly wrapped in Optional[], even when None was not part of the Union. This caused imprecise schema descriptions to be sent to LLMs when generating structured outputs.

Before:

  • Union[str, int]Optional[Union[str, int]]
  • Union[str, int, None]Optional[Union[str, int]]

After:

  • Union[str, int]Union[str, int]
  • Union[str, int, None]Optional[Union[str, int]]

Key Changes:

  • Added logic to check if None is actually present in the Union before wrapping with Optional
  • Added support for Python 3.10+ pipe syntax (int | str | None)
  • Added comprehensive test coverage for Union type scenarios
  • Updated existing test expectation to match corrected behavior

Fixes #3735

Review & Testing Checklist for Human

  • Verify schema generation: Run pytest tests/test_union_types.py -vv to confirm all Union type tests pass (13 tests should pass)
  • Check format consistency: Review the changes in converter.py to ensure the Union/Optional logic is correct and handles edge cases properly
  • VCR cassette updates: Two tests (test_converter_with_llama3_1_model and test_converter_with_llama3_2_model) will fail with cassette mismatches because the schema format changed. These cassettes should be regenerated or the tests marked as expected to change
  • End-to-end validation: Test with a real crew/task using output_pydantic with Union types to ensure LLMs correctly understand the new schema format

Recommended test plan:

from typing import Union
from pydantic import BaseModel
from crewai import Agent, Task, Crew

class SuccessResponse(BaseModel):
    status: str = "success"
    data: str

class ErrorResponse(BaseModel):
    status: str = "error"
    message: str

class Response(BaseModel):
    result: Union[SuccessResponse, ErrorResponse]

# Test that the agent can correctly return either variant
task = Task(
    description="Return a success response with data='test'",
    expected_output="A structured response",
    output_pydantic=Response,
    agent=your_agent
)

Notes

- Fixed generate_model_description to correctly handle Union types
- Union types without None are now properly formatted as Union[type1, type2]
- Union types with None are correctly wrapped in Optional[Union[...]]
- Added support for Python 3.10+ pipe syntax (int | str | None)
- Added comprehensive tests for Union type support
- Updated existing test expectations to match corrected behavior

This fixes issue #3735 where Union types were incorrectly wrapped in Optional
even when None was not part of the Union.

Co-Authored-By: João <joao@crewai.com>
@devin-ai-integration
Copy link
Contributor Author

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment and CI monitoring
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

1 participant