Behavior of Union Types as Output type

[vikigenius]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

Following the current docs in https://ai.pydantic.dev/output/#structured-output

from pydantic_ai import Agent

CustomType = list[str] | list[int]

agent1 = Agent[None, list[str] | list[int]](
    'openai:gpt-4o-mini',
    output_type=list[str] | list[int],  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

agent2 = Agent[None, CustomType](
    'openai:gpt-4o-mini',
    output_type=CustomType,  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

agent1 is fine. But pyright still complains about agent2

/home/void/Projects/myproj/test.py
  /home/void/Projects/myproj/test.py:13:10 - error: No overloads for "__init__" match the provided arguments (reportCallIssue)
  /home/void/Projects/myproj/test.py:15:17 - error: Argument of type "CustomType" cannot be assigned to parameter "output_type" of type "OutputSpec[CustomType]" in function "__init__"
    Type "UnionType" is not assignable to type "OutputSpec[CustomType]"
      Type "UnionType" is not assignable to type "type[list[str]]"
      Type "UnionType" is not assignable to type "type[list[int]]"
      Type "UnionType" is not assignable to type "(...) -> (CustomType | Awaitable[CustomType])"
      "UnionType" is not assignable to "ToolOutput[CustomType]"
      "UnionType" is not assignable to "NativeOutput[CustomType]"
      "UnionType" is not assignable to "PromptedOutput[CustomType]"
      "UnionType" is not assignable to "TextOutput[CustomType]"
    ... (reportArgumentType)```

This used to work around in the 0.3.5 versions of pydantic-ai. Did something change about the type signature of init

Example Code

from pydantic_ai import Agent

CustomType = list[str] | list[int]

agent1 = Agent[None, list[str] | list[int]](
    'openai:gpt-4o-mini',
    output_type=list[str] | list[int],  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

agent2 = Agent[None, CustomType](
    'openai:gpt-4o-mini',
    output_type=CustomType,  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

Python, Pydantic AI & LLM client version

pydantic-ai 0.6.2
python 3.12

[Kludex]

It seems the first didn't even work in <0.6.0.

I think I fixed it in the PR I dropped deprecation warnings. 🤔

---

@Viicos Can you help here, please?

[Viicos]

It seems to me that:

agent1 = Agent[None, list[str] | list[int]](
    'openai:gpt-4o-mini',
    output_type=list[str] | list[int],  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)
reveal_type(agent1)  # Agent[None, list[str] | list[int]]

# or

agent1 = Agent(
    'openai:gpt-4o-mini',
    output_type=list[str] | list[int],  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)
reveal_type(agent1)  Agent[None, list[str] | list[int]]

Shouldn't be accepted by Pyright. Actually, trying to make a MRE of this with a different class results in a type checker error, so I wasn't able to reproduce this weird behavior. It might be an issue with Pyright that results in this false negative.

Anyway, this is what PEP 747 is meant for, so until then a type ignore comment can be used, or by following the Type checking considerations note in the linked documentation.

[vikigenius]

@Viicos @Kludex sorry i missed it, But the main difference in behaviour from older versions is actually this

from pydantic_ai import Agent

CustomType = list[str] | list[int]

agent1 = Agent[None, list[str] | list[int]](
    'openai:gpt-4o-mini',
    output_type=list[str] | list[int],  
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

agent2 = Agent[None, CustomType](
    'openai:gpt-4o-mini',
    output_type=CustomType, # pyright: ignore[reportArgumentType]
    system_prompt='Extract either colors or sizes from the shapes provided.',
)

Previously just the pyright: ignore[reportArgumentType] was enough. But now there is an additional issue raised which is the no overloads for __init__ match the provided arguments

It could just be an issue with pyright, because I don't see a reason why agent1 gets no complaints at all.

[Viicos]

Overloads were added on the __init__() method not long ago I believe, so you'll most likely have to ignore the new error code as well (pyright: ignore[reportArgumentType, reportCallIssue], or pyright: ignore).