pydantic-ai-slim 1.0.18__py3-none-any.whl → 1.1.0__py3-none-any.whl

pydantic_ai/durable_exec/prefect/_cache_policies.py

@@ -0,0 +1,102 @@
+from dataclasses import fields, is_dataclass
+from typing import Any, TypeGuard
+
+from prefect.cache_policies import INPUTS, RUN_ID, TASK_SOURCE, CachePolicy
+from prefect.context import TaskRunContext
+
+from pydantic_ai import ToolsetTool
+from pydantic_ai.tools import RunContext
+
+
+def _is_dict(obj: Any) -> TypeGuard[dict[str, Any]]:
+    return isinstance(obj, dict)
+
+
+def _is_list(obj: Any) -> TypeGuard[list[Any]]:
+    return isinstance(obj, list)
+
+
+def _is_tuple(obj: Any) -> TypeGuard[tuple[Any, ...]]:
+    return isinstance(obj, tuple)
+
+
+def _is_toolset_tool(obj: Any) -> TypeGuard[ToolsetTool]:
+    return isinstance(obj, ToolsetTool)
+
+
+def _replace_run_context(
+    inputs: dict[str, Any],
+) -> Any:
+    """Replace RunContext objects with a dict containing only hashable fields."""
+    for key, value in inputs.items():
+        if isinstance(value, RunContext):
+            inputs[key] = {
+                'retries': value.retries,
+                'tool_call_id': value.tool_call_id,
+                'tool_name': value.tool_name,
+                'tool_call_approved': value.tool_call_approved,
+                'retry': value.retry,
+                'max_retries': value.max_retries,
+                'run_step': value.run_step,
+            }
+
+    return inputs
+
+
+def _strip_timestamps(
+    obj: Any | dict[str, Any] | list[Any] | tuple[Any, ...],
+) -> Any:
+    """Recursively convert dataclasses to dicts, excluding timestamp fields."""
+    if is_dataclass(obj) and not isinstance(obj, type):
+        result: dict[str, Any] = {}
+        for f in fields(obj):
+            if f.name != 'timestamp':
+                value = getattr(obj, f.name)
+                result[f.name] = _strip_timestamps(value)
+        return result
+    elif _is_dict(obj):
+        return {k: _strip_timestamps(v) for k, v in obj.items() if k != 'timestamp'}
+    elif _is_list(obj):
+        return [_strip_timestamps(item) for item in obj]
+    elif _is_tuple(obj):
+        return tuple(_strip_timestamps(item) for item in obj)
+    return obj
+
+
+def _replace_toolsets(
+    inputs: dict[str, Any],
+) -> Any:
+    """Replace Toolset objects with a dict containing only hashable fields."""
+    inputs = inputs.copy()
+    for key, value in inputs.items():
+        if _is_toolset_tool(value):
+            inputs[key] = {field.name: getattr(value, field.name) for field in fields(value) if field.name != 'toolset'}
+    return inputs
+
+
+class PrefectAgentInputs(CachePolicy):
+    """Cache policy designed to handle input hashing for PrefectAgent cache keys.
+
+    Computes a cache key based on inputs, ignoring nested 'timestamp' fields
+    and serializing RunContext objects to only include hashable fields.
+    """
+
+    def compute_key(
+        self,
+        task_ctx: TaskRunContext,
+        inputs: dict[str, Any],
+        flow_parameters: dict[str, Any],
+        **kwargs: Any,
+    ) -> str | None:
+        """Compute cache key from inputs with timestamps removed and RunContext serialized."""
+        if not inputs:
+            return None
+
+        inputs_without_toolsets = _replace_toolsets(inputs)
+        inputs_with_hashable_context = _replace_run_context(inputs_without_toolsets)
+        filtered_inputs = _strip_timestamps(inputs_with_hashable_context)
+
+        return INPUTS.compute_key(task_ctx, filtered_inputs, flow_parameters, **kwargs)
+
+
+DEFAULT_PYDANTIC_AI_CACHE_POLICY = PrefectAgentInputs() + TASK_SOURCE + RUN_ID

pydantic_ai/durable_exec/prefect/_function_toolset.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import Any
+
+from prefect import task
+
+from pydantic_ai import FunctionToolset, ToolsetTool
+from pydantic_ai.tools import AgentDepsT, RunContext
+
+from ._toolset import PrefectWrapperToolset
+from ._types import TaskConfig, default_task_config
+
+
+class PrefectFunctionToolset(PrefectWrapperToolset[AgentDepsT]):
+    """A wrapper for FunctionToolset that integrates with Prefect, turning tool calls into Prefect tasks."""
+
+    def __init__(
+        self,
+        wrapped: FunctionToolset[AgentDepsT],
+        *,
+        task_config: TaskConfig,
+        tool_task_config: dict[str, TaskConfig | None],
+    ):
+        super().__init__(wrapped)
+        self._task_config = default_task_config | (task_config or {})
+        self._tool_task_config = tool_task_config or {}
+
+        @task
+        async def _call_tool_task(
+            tool_name: str,
+            tool_args: dict[str, Any],
+            ctx: RunContext[AgentDepsT],
+            tool: ToolsetTool[AgentDepsT],
+        ) -> Any:
+            return await super(PrefectFunctionToolset, self).call_tool(tool_name, tool_args, ctx, tool)
+
+        self._call_tool_task = _call_tool_task
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> Any:
+        """Call a tool, wrapped as a Prefect task with a descriptive name."""
+        # Check if this specific tool has custom config or is disabled
+        tool_specific_config = self._tool_task_config.get(name, default_task_config)
+        if tool_specific_config is None:
+            # None means this tool should not be wrapped as a task
+            return await super().call_tool(name, tool_args, ctx, tool)
+
+        # Merge tool-specific config with default config
+        merged_config = self._task_config | tool_specific_config
+
+        return await self._call_tool_task.with_options(name=f'Call Tool: {name}', **merged_config)(
+            name, tool_args, ctx, tool
+        )

pydantic_ai/durable_exec/prefect/_mcp_server.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import TYPE_CHECKING, Any
+
+from prefect import task
+from typing_extensions import Self
+
+from pydantic_ai import ToolsetTool
+from pydantic_ai.tools import AgentDepsT, RunContext
+
+from ._toolset import PrefectWrapperToolset
+from ._types import TaskConfig, default_task_config
+
+if TYPE_CHECKING:
+    from pydantic_ai.mcp import MCPServer, ToolResult
+
+
+class PrefectMCPServer(PrefectWrapperToolset[AgentDepsT], ABC):
+    """A wrapper for MCPServer that integrates with Prefect, turning call_tool and get_tools into Prefect tasks."""
+
+    def __init__(
+        self,
+        wrapped: MCPServer,
+        *,
+        task_config: TaskConfig,
+    ):
+        super().__init__(wrapped)
+        self._task_config = default_task_config | (task_config or {})
+        self._mcp_id = wrapped.id
+
+        @task
+        async def _call_tool_task(
+            tool_name: str,
+            tool_args: dict[str, Any],
+            ctx: RunContext[AgentDepsT],
+            tool: ToolsetTool[AgentDepsT],
+        ) -> ToolResult:
+            return await super(PrefectMCPServer, self).call_tool(tool_name, tool_args, ctx, tool)
+
+        self._call_tool_task = _call_tool_task
+
+    async def __aenter__(self) -> Self:
+        await self.wrapped.__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        return await self.wrapped.__aexit__(*args)
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> ToolResult:
+        """Call an MCP tool, wrapped as a Prefect task with a descriptive name."""
+        return await self._call_tool_task.with_options(name=f'Call MCP Tool: {name}', **self._task_config)(
+            name, tool_args, ctx, tool
+        )

pydantic_ai/durable_exec/prefect/_model.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import Any
+
+from prefect import task
+from prefect.context import FlowRunContext
+
+from pydantic_ai import (
+    ModelMessage,
+    ModelResponse,
+    ModelResponseStreamEvent,
+)
+from pydantic_ai.agent import EventStreamHandler
+from pydantic_ai.models import ModelRequestParameters, StreamedResponse
+from pydantic_ai.models.wrapper import WrapperModel
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import RunContext
+from pydantic_ai.usage import RequestUsage
+
+from ._types import TaskConfig, default_task_config
+
+
+class PrefectStreamedResponse(StreamedResponse):
+    """A non-streaming response wrapper for Prefect tasks.
+
+    When a model request is executed inside a Prefect flow, the entire stream
+    is consumed within the task, and this wrapper is returned containing the
+    final response.
+    """
+
+    def __init__(self, model_request_parameters: ModelRequestParameters, response: ModelResponse):
+        super().__init__(model_request_parameters)
+        self.response = response
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Return an empty iterator since the stream has already been consumed."""
+        return
+        # noinspection PyUnreachableCode
+        yield
+
+    def get(self) -> ModelResponse:
+        return self.response
+
+    def usage(self) -> RequestUsage:
+        return self.response.usage  # pragma: no cover
+
+    @property
+    def model_name(self) -> str:
+        return self.response.model_name or ''  # pragma: no cover
+
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
+    @property
+    def timestamp(self) -> datetime:
+        return self.response.timestamp  # pragma: no cover
+
+
+class PrefectModel(WrapperModel):
+    """A wrapper for Model that integrates with Prefect, turning request and request_stream into Prefect tasks."""
+
+    def __init__(
+        self,
+        model: Any,
+        *,
+        task_config: TaskConfig,
+        event_stream_handler: EventStreamHandler[Any] | None = None,
+    ):
+        super().__init__(model)
+        self.task_config = default_task_config | (task_config or {})
+        self.event_stream_handler = event_stream_handler
+
+        @task
+        async def wrapped_request(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+        ) -> ModelResponse:
+            response = await super(PrefectModel, self).request(messages, model_settings, model_request_parameters)
+            return response
+
+        self._wrapped_request = wrapped_request
+
+        @task
+        async def request_stream_task(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+            ctx: RunContext[Any] | None,
+        ) -> ModelResponse:
+            async with super(PrefectModel, self).request_stream(
+                messages, model_settings, model_request_parameters, ctx
+            ) as streamed_response:
+                if self.event_stream_handler is not None:
+                    assert ctx is not None, (
+                        'A Prefect model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. '
+                        'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+                    )
+                    await self.event_stream_handler(ctx, streamed_response)
+
+                # Consume the entire stream
+                async for _ in streamed_response:
+                    pass
+            response = streamed_response.get()
+            return response
+
+        self._wrapped_request_stream = request_stream_task
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a model request, wrapped as a Prefect task when in a flow."""
+        return await self._wrapped_request.with_options(
+            name=f'Model Request: {self.wrapped.model_name}', **self.task_config
+        )(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming model request.
+
+        When inside a Prefect flow, the stream is consumed within a task and
+        a non-streaming response is returned. When not in a flow, behaves normally.
+        """
+        # Check if we're in a flow context
+        flow_run_context = FlowRunContext.get()
+
+        # If not in a flow, just call the wrapped request_stream method
+        if flow_run_context is None:
+            async with super().request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                yield streamed_response
+                return
+
+        # If in a flow, consume the stream in a task and return the final response
+        response = await self._wrapped_request_stream.with_options(
+            name=f'Model Request (Streaming): {self.wrapped.model_name}', **self.task_config
+        )(messages, model_settings, model_request_parameters, run_context)
+        yield PrefectStreamedResponse(model_request_parameters, response)

pydantic_ai/durable_exec/prefect/_toolset.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from pydantic_ai import AbstractToolset, FunctionToolset, WrapperToolset
+from pydantic_ai.tools import AgentDepsT
+
+from ._types import TaskConfig
+
+if TYPE_CHECKING:
+    pass
+
+
+class PrefectWrapperToolset(WrapperToolset[AgentDepsT], ABC):
+    """Base class for Prefect-wrapped toolsets."""
+
+    @property
+    def id(self) -> str | None:
+        # Prefect toolsets should have IDs for better task naming
+        return self.wrapped.id
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        # Prefect-ified toolsets cannot be swapped out after the fact.
+        return self
+
+
+def prefectify_toolset(
+    toolset: AbstractToolset[AgentDepsT],
+    mcp_task_config: TaskConfig,
+    tool_task_config: TaskConfig,
+    tool_task_config_by_name: dict[str, TaskConfig | None],
+) -> AbstractToolset[AgentDepsT]:
+    """Wrap a toolset to integrate it with Prefect.
+
+    Args:
+        toolset: The toolset to wrap.
+        mcp_task_config: The Prefect task config to use for MCP server tasks.
+        tool_task_config: The default Prefect task config to use for tool calls.
+        tool_task_config_by_name: Per-tool task configuration. Keys are tool names, values are TaskConfig or None.
+    """
+    if isinstance(toolset, FunctionToolset):
+        from ._function_toolset import PrefectFunctionToolset
+
+        return PrefectFunctionToolset(
+            wrapped=toolset,
+            task_config=tool_task_config,
+            tool_task_config=tool_task_config_by_name,
+        )
+
+    try:
+        from pydantic_ai.mcp import MCPServer
+
+        from ._mcp_server import PrefectMCPServer
+    except ImportError:
+        pass
+    else:
+        if isinstance(toolset, MCPServer):
+            return PrefectMCPServer(
+                wrapped=toolset,
+                task_config=mcp_task_config,
+            )
+
+    return toolset

pydantic_ai/durable_exec/prefect/_types.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from prefect.cache_policies import CachePolicy
+from prefect.results import ResultStorage
+from typing_extensions import TypedDict
+
+from pydantic_ai.durable_exec.prefect._cache_policies import DEFAULT_PYDANTIC_AI_CACHE_POLICY
+
+
+class TaskConfig(TypedDict, total=False):
+    """Configuration for a task in Prefect.
+
+    These options are passed to the `@task` decorator.
+    """
+
+    retries: int
+    """Maximum number of retries for the task."""
+
+    retry_delay_seconds: float | list[float]
+    """Delay between retries in seconds. Can be a single value or a list for custom backoff."""
+
+    timeout_seconds: float
+    """Maximum time in seconds for the task to complete."""
+
+    cache_policy: CachePolicy
+    """Prefect cache policy for the task."""
+
+    persist_result: bool
+    """Whether to persist the task result."""
+
+    result_storage: ResultStorage
+    """Prefect result storage for the task. Should be a storage block or a block slug like `s3-bucket/my-storage`."""
+
+    log_prints: bool
+    """Whether to log print statements from the task."""
+
+
+default_task_config = TaskConfig(
+    retries=0,
+    retry_delay_seconds=1.0,
+    persist_result=True,
+    log_prints=False,
+    cache_policy=DEFAULT_PYDANTIC_AI_CACHE_POLICY,
+)

pydantic_ai/models/__init__.py

@@ -685,6 +685,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'grok',
         'heroku',
         'moonshotai',
+        'ollama',
         'openai',
         'openai-chat',
         'openrouter',

pydantic_ai/models/google.py

@@ -73,6 +73,7 @@ try:
         GroundingMetadata,
         HttpOptionsDict,
         MediaResolution,
+        Modality,
         Part,
         PartDict,
         SafetySettingDict,
@@ -415,6 +416,10 @@ class GoogleModel(Model):
         tool_config = self._get_tool_config(model_request_parameters, tools)
         system_instruction, contents = await self._map_messages(messages)
 
+        modalities = [Modality.TEXT.value]
+        if self.profile.supports_image_output:
+            modalities.append(Modality.IMAGE.value)
+
         http_options: HttpOptionsDict = {
             'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
         }
@@ -443,6 +448,7 @@ class GoogleModel(Model):
             tool_config=tool_config,
             response_mime_type=response_mime_type,
             response_schema=response_schema,
+            response_modalities=modalities,
         )
         return contents, config
 

pydantic_ai/providers/gateway.py

@@ -4,7 +4,6 @@ from __future__ import annotations as _annotations
 
 import os
 from typing import TYPE_CHECKING, Any, Literal, overload
-from urllib.parse import urljoin
 
 import httpx
 
@@ -84,22 +83,22 @@ def gateway_provider(
             ' to use the Pydantic AI Gateway provider.'
         )
 
-    base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', 'http://localhost:8787')
+    base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', 'http://localhost:8787/proxy')
     http_client = http_client or cached_async_http_client(provider=f'gateway-{upstream_provider}')
     http_client.event_hooks = {'request': [_request_hook]}
 
     if upstream_provider in ('openai', 'openai-chat'):
         from .openai import OpenAIProvider
 
-        return OpenAIProvider(api_key=api_key, base_url=urljoin(base_url, 'openai'), http_client=http_client)
+        return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'openai'), http_client=http_client)
     elif upstream_provider == 'openai-responses':
         from .openai import OpenAIProvider
 
-        return OpenAIProvider(api_key=api_key, base_url=urljoin(base_url, 'openai'), http_client=http_client)
+        return OpenAIProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'openai'), http_client=http_client)
     elif upstream_provider == 'groq':
         from .groq import GroqProvider
 
-        return GroqProvider(api_key=api_key, base_url=urljoin(base_url, 'groq'), http_client=http_client)
+        return GroqProvider(api_key=api_key, base_url=_merge_url_path(base_url, 'groq'), http_client=http_client)
     elif upstream_provider == 'anthropic':
         from anthropic import AsyncAnthropic
 
@@ -108,7 +107,7 @@ def gateway_provider(
         return AnthropicProvider(
             anthropic_client=AsyncAnthropic(
                 auth_token=api_key,
-                base_url=urljoin(base_url, 'anthropic'),
+                base_url=_merge_url_path(base_url, 'anthropic'),
                 http_client=http_client,
             )
         )
@@ -122,7 +121,7 @@ def gateway_provider(
                 vertexai=True,
                 api_key='unset',
                 http_options={
-                    'base_url': f'{base_url}/google-vertex',
+                    'base_url': _merge_url_path(base_url, 'google-vertex'),
                     'headers': {'User-Agent': get_user_agent(), 'Authorization': api_key},
                     # TODO(Marcelo): Until https://github.com/googleapis/python-genai/issues/1357 is solved.
                     'async_client_args': {
@@ -185,3 +184,13 @@ async def _request_hook(request: httpx.Request) -> httpx.Request:
     request.headers.update(headers)
 
     return request
+
+
+def _merge_url_path(base_url: str, path: str) -> str:
+    """Merge a base URL and a path.
+
+    Args:
+        base_url: The base URL to merge.
+        path: The path to merge.
+    """
+    return base_url.rstrip('/') + '/' + path.lstrip('/')

pydantic_ai/toolsets/function.py

@@ -109,6 +109,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         /,
         *,
         name: str | None = None,
+        description: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat | None = None,
@@ -126,6 +127,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         /,
         *,
         name: str | None = None,
+        description: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat | None = None,
@@ -169,6 +171,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         Args:
             func: The tool function to register.
             name: The name of the tool, defaults to the function name.
+            description: The description of the tool,defaults to the function docstring.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
@@ -197,18 +200,19 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         ) -> ToolFuncEither[AgentDepsT, ToolParams]:
             # noinspection PyTypeChecker
             self.add_function(
-                func_,
-                None,
-                name,
-                retries,
-                prepare,
-                docstring_format,
-                require_parameter_descriptions,
-                schema_generator,
-                strict,
-                sequential,
-                requires_approval,
-                metadata,
+                func=func_,
+                takes_ctx=None,
+                name=name,
+                description=description,
+                retries=retries,
+                prepare=prepare,
+                docstring_format=docstring_format,
+                require_parameter_descriptions=require_parameter_descriptions,
+                schema_generator=schema_generator,
+                strict=strict,
+                sequential=sequential,
+                requires_approval=requires_approval,
+                metadata=metadata,
             )
             return func_
 
@@ -219,6 +223,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
         func: ToolFuncEither[AgentDepsT, ToolParams],
         takes_ctx: bool | None = None,
         name: str | None = None,
+        description: str | None = None,
         retries: int | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat | None = None,
@@ -240,6 +245,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
             func: The tool function to register.
             takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] as its first argument. If `None`, this is inferred from the function signature.
             name: The name of the tool, defaults to the function name.
+            description: The description of the tool, defaults to the function docstring.
             retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
             prepare: custom method to prepare the tool definition for each step, return `None` to omit this
@@ -279,6 +285,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
             func,
             takes_ctx=takes_ctx,
             name=name,
+            description=description,
             max_retries=retries,
             prepare=prepare,
             docstring_format=docstring_format,