Add support for schema_generator to ToolOutput

[dmontagu]

Inspired by this thread https://pydanticlogfire.slack.com/archives/C083V7PMHHA/p1744850394661209

It would be nice to be able to specify the schema_generator for a ToolOutput same as we do for normal tools. This PR makes that possible.

Needs tests and docs still.

import os
from pydantic import BaseModel, Field
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema

from pydantic_ai import Agent, ToolOutput
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider
import logfire

from pydantic_ai.tools import GenerateToolJsonSchema

logfire.configure(console=logfire.ConsoleOptions(verbose=True))
logfire.instrument_pydantic_ai()
logfire.instrument_httpx(capture_all=True)

model = OpenAIModel(
    model_name="Qwen/Qwen2.5-72B-Instruct-Turbo",
    provider=OpenAIProvider(
        base_url="https://api.together.xyz/v1",
        api_key=os.getenv("TOGETHER_API_KEY"),
    ),
)

class Yuk(BaseModel):
    x: str | None =  Field(
        default=None,
        description="The word that follows x",
    )
    y: str | None = Field(
        default=None,
        description="The word that follows y",
    )

class MyGenerateToolJsonSchema(GenerateToolJsonSchema):
    def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
        result = super().nullable_schema(schema)
        if 'anyOf' in result:
            result['oneOf'] = result.pop('anyOf')
        return result

def main():
    agent = Agent(model=model, output_type=ToolOutput(type_=Yuk, schema_generator=MyGenerateToolJsonSchema), system_prompt="output the word that follows x or y, output None if no word is found")
    result = agent.run_sync("x: sky, y: ")
    print(result)

if __name__ == "__main__":
    main()

[github-actions]

Docs Preview

commit:	f70bffe
Preview URL:	https://9c7c0f5b-pydantic-ai-previews.pydantic.workers.dev

[DouweM]

Implementation looks good