Add run options

Author: igalshilman

python/restate/ext/pydantic/_agent.py

@@ -4,7 +4,7 @@
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from typing import Any, overload
 
-from restate import TerminalError
+from restate import RunOptions, TerminalError
 from restate.ext.pydantic._utils import state_context
 from restate.extensions import current_context
 
@@ -93,6 +93,7 @@ def __init__(
         *,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
         disable_auto_wrapping_tools: bool = False,
+        run_options: RunOptions | None = None,
     ):
         super().__init__(wrapped)
         if not isinstance(wrapped.model, Model):
@@ -102,12 +103,16 @@ def __init__(
 
         self._event_stream_handler = event_stream_handler
         self._disable_auto_wrapping_tools = disable_auto_wrapping_tools
-        self._model = RestateModelWrapper(wrapped.model, event_stream_handler=event_stream_handler, max_attempts=3)
+
+        if run_options is None:
+            run_options = RunOptions(max_attempts=3)
+
+        self._model = RestateModelWrapper(wrapped.model, run_options, event_stream_handler=event_stream_handler)
 
         def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
             """Set the Restate context for the toolset, wrapping tools if needed."""
             if isinstance(toolset, FunctionToolset) and not disable_auto_wrapping_tools:
-                return RestateContextRunToolSet(toolset)
+                return RestateContextRunToolSet(toolset, run_options)
             try:
                 from pydantic_ai.mcp import MCPServer
 
@@ -116,7 +121,7 @@ def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDe
                 pass
             else:
                 if isinstance(toolset, MCPServer):
-                    return RestateMCPServer(toolset)
+                    return RestateMCPServer(toolset, run_options)
 
             return toolset
 

python/restate/ext/pydantic/_model.py

@@ -2,6 +2,7 @@
 from contextlib import asynccontextmanager
 from datetime import datetime
 from typing import Any
+import dataclasses
 
 from restate import RunOptions, SdkInternalBaseException
 from restate.ext.pydantic._utils import current_state
@@ -60,11 +61,12 @@ class RestateModelWrapper(WrapperModel):
     def __init__(
         self,
         wrapped: Model,
+        run_options: RunOptions,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
-        max_attempts: int | None = None,
     ):
         super().__init__(wrapped)
-        self._options = RunOptions(serde=MODEL_RESPONSE_SERDE, max_attempts=max_attempts)
+
+        self._options = dataclasses.replace(run_options, serde=MODEL_RESPONSE_SERDE)
         self._event_stream_handler = event_stream_handler
 
     async def request(self, *args: Any, **kwargs: Any) -> ModelResponse:

python/restate/ext/pydantic/_toolset.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any, Literal
 
 from restate import RunOptions, SdkInternalBaseException, TerminalError
@@ -55,9 +55,9 @@ class RestateMCPToolRunResult:
 class RestateContextRunToolSet(WrapperToolset[AgentDepsT]):
     """A toolset that automatically wraps tool calls with restate's `ctx.run_typed()`."""
 
-    def __init__(self, wrapped: AbstractToolset[AgentDepsT]):
+    def __init__(self, wrapped: AbstractToolset[AgentDepsT], run_options: RunOptions):
         super().__init__(wrapped)
-        self.options = RunOptions[RestateContextRunResult](serde=CONTEXT_RUN_SERDE)
+        self.options = replace(run_options, serde=CONTEXT_RUN_SERDE)
 
     async def call_tool(
         self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
@@ -119,9 +119,11 @@ def visit_and_replace(
 class RestateMCPServer(WrapperToolset[AgentDepsT]):
     """A wrapper for MCPServer that integrates with restate."""
 
-    def __init__(self, wrapped: MCPServer):
+    def __init__(self, wrapped: MCPServer, run_options: RunOptions):
         super().__init__(wrapped)
         self._wrapped = wrapped
+        self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE)
+        self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE)
 
     def visit_and_replace(
         self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
@@ -136,16 +138,14 @@ async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult:
             # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity.
             return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()})
 
-        options = RunOptions(serde=MCP_GET_TOOLS_SERDE)
-
         context = current_context()
         if context is None:
             raise UserError(
                 "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
             )
 
         try:
-            tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, options)
+            tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options)
             return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.output.items()}
         except SdkInternalBaseException as e:
             raise Exception("Internal error during get_tools call") from e
@@ -165,14 +165,13 @@ async def call_tool_in_context() -> RestateMCPToolRunResult:
             res = await self._wrapped.call_tool(name, tool_args, ctx, tool)
             return RestateMCPToolRunResult(output=res)
 
-        options = RunOptions(serde=MCP_RUN_SERDE)
         context = current_context()
         if context is None:
             raise UserError(
                 "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context."
             )
         try:
-            res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, options)
+            res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options)
         except SdkInternalBaseException as e:
             raise Exception("Internal error during tool call") from e