Pydantic AI

Pydantic AI adapter for Flyte.

Bring your own Pydantic AI Agent and run it durably on Flyte. The adapter provides:

• flyteplugins.agents.pydantic_ai.tool — turn a Flyte @env.task into a Pydantic AI tool that executes as a durable child action (own container/GPU, retries, caching). This is the shared flyteplugins.agents.core.tool: Pydantic AI accepts plain (async) callables in Agent(tools=[...]) and infers each tool’s schema from the callable’s signature, which the core wrapper preserves.
• flyteplugins.agents.pydantic_ai.run_agent — run the Pydantic AI agent loop inside your task and return the final answer.

Each tool call runs as a durable Flyte child action, and the run timeline is rendered into the Flyte task report.

Directory

Classes

Methods

Methods

run_agent()

def run_agent(
    input: str,
    tools: typing.Sequence[typing.Any] = (),
    model: typing.Any = None,
    instructions: str | None = None,
    agent: typing.Any = None,
    name: str = 'pydantic-ai-agent',
    durable: bool = True,
    observability: bool = True,
    memory_key: str | None = None,
    **run_kwargs: typing.Any,
) -> str

Run a Pydantic AI agent with the given tools and prompt; return the final text.

Await this from an async task as await run_agent(...); from a sync task use flyteplugins.agents.pydantic_ai.run_agent_sync instead.

Call this from inside an @env.task — that task is the durable parent. Within it, each tool call runs as a durable Flyte child action. Give the enclosing task retries=... for self-healing and report=True to see the agent timeline.

Provide either a pre-built agent (with its tools already attached) or tools + model to have one built for you — not both.

Returns: The agent’s final output as a string.

run_agent_sync()

def run_agent_sync(
    input: str,
    tools: typing.Sequence[typing.Any] = (),
    model: typing.Any = None,
    instructions: str | None = None,
    agent: typing.Any = None,
    name: str = 'pydantic-ai-agent',
    durable: bool = True,
    observability: bool = True,
    memory_key: str | None = None,
    **run_kwargs: typing.Any,
) -> str

Synchronous variant of run_agent for use in sync tasks; runs the async implementation on a dedicated event loop.

Run a Pydantic AI agent with the given tools and prompt; return the final text.

Await this from an async task as await run_agent(...); from a sync task use flyteplugins.agents.pydantic_ai.run_agent_sync instead.

Call this from inside an @env.task — that task is the durable parent. Within it, each tool call runs as a durable Flyte child action. Give the enclosing task retries=... for self-healing and report=True to see the agent timeline.

Provide either a pre-built agent (with its tools already attached) or tools + model to have one built for you — not both.

Returns

The agent’s final output as a string.

tool()

def tool(
    func: AsyncFunctionTaskTemplate | typing.Callable | None = None,
    name: str | None = None,
    description: str | None = None,
) -> typing.Callable

Wrap a Flyte @env.task as a plain async tool function — the generic default.

For SDKs that accept plain Python callables as tools (deriving the schema from the signature + docstring), this is the whole adapter tool: the returned function carries the task’s signature (functools.wraps), dispatches to task.aio() (so each call is a durable Flyte child action), exposes __wrapped_task__, and wires the backing task to flyteplugins.agents.core.ToolTaskResolver. Adapters whose SDK needs a native tool type (e.g. OpenAI’s FunctionTool, Claude’s MCP SdkMcpTool) provide their own instead.

Also accepts any other callable — a plain function or an instance of a callable class defining __call__ — and returns it usable as a tool as-is, since the plain-callable SDKs derive the schema by inspecting the callable (a class instance is inspected through its __call__). A name or description override is applied to the callable best-effort.

Usable bare, parametrized or as a direct call:

@tool
@env.task
async def get_weather(city: str) -> str: ...