pydantic-ai-slim 1.0.18__py3-none-any.whl → 1.2.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

@@ -55,6 +55,8 @@ KnownModelName = TypeAliasType(
         'anthropic:claude-3-5-sonnet-20240620',
         'anthropic:claude-3-5-sonnet-20241022',
         'anthropic:claude-3-5-sonnet-latest',
+        'anthropic:claude-haiku-4-5',
+        'anthropic:claude-haiku-4-5-20251001',
         'anthropic:claude-3-7-sonnet-20250219',
         'anthropic:claude-3-7-sonnet-latest',
         'anthropic:claude-3-haiku-20240307',
@@ -685,6 +687,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