Integrate turnstile into OpenAI ext

Author: igalshilman

python/restate/ext/openai/models.py

@@ -0,0 +1,79 @@
+#
+#  Copyright (c) 2023-2024 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""
+This module contains the optional OpenAI integration for Restate.
+"""
+
+import dataclasses
+
+from agents import (
+    Usage,
+)
+from agents.items import TResponseOutputItem
+from agents.items import TResponseInputItem
+from datetime import timedelta
+from typing import Optional
+from pydantic import BaseModel
+
+from restate.ext.turnstile import Turnstile
+
+
+class State:
+    __slots__ = ("turnstile",)
+
+    def __init__(self) -> None:
+        self.turnstile = Turnstile([])
+
+
+@dataclasses.dataclass
+class LlmRetryOpts:
+    max_attempts: Optional[int] = 10
+    """Max number of attempts (including the initial), before giving up.
+
+    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
+    max_duration: Optional[timedelta] = None
+    """Max duration of retries, before giving up.
+
+    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
+    initial_retry_interval: Optional[timedelta] = timedelta(seconds=1)
+    """Initial interval for the first retry attempt.
+    Retry interval will grow by a factor specified in `retry_interval_factor`.
+
+    If any of the other retry related fields is specified, the default for this field is 50 milliseconds, otherwise restate will fallback to the overall invocation retry policy."""
+    max_retry_interval: Optional[timedelta] = None
+    """Max interval between retries.
+    Retry interval will grow by a factor specified in `retry_interval_factor`.
+
+    The default is 10 seconds."""
+    retry_interval_factor: Optional[float] = None
+    """Exponentiation factor to use when computing the next retry delay.
+
+    If any of the other retry related fields is specified, the default for this field is `2`, meaning retry interval will double at each attempt, otherwise restate will fallback to the overall invocation retry policy."""
+
+
+# The OpenAI ModelResponse class is a dataclass with Pydantic fields.
+# The Restate SDK cannot serialize this. So we turn the ModelResponse int a Pydantic model.
+class RestateModelResponse(BaseModel):
+    output: list[TResponseOutputItem]
+    """A list of outputs (messages, tool calls, etc) generated by the model"""
+
+    usage: Usage
+    """The usage information for the response."""
+
+    response_id: str | None
+    """An ID for the response which can be used to refer to the response in subsequent calls to the
+    model. Not supported by all model providers.
+    If using OpenAI models via the Responses API, this is the `response_id` parameter, and it can
+    be passed to `Runner.run`.
+    """
+
+    def to_input_items(self) -> list[TResponseInputItem]:
+        return [it.model_dump(exclude_unset=True) for it in self.output]  # type: ignore

python/restate/ext/openai/runner_wrapper.py

@@ -15,7 +15,6 @@
 import dataclasses
 
 from agents import (
-    Usage,
     Model,
     RunContextWrapper,
     AgentsException,
@@ -28,84 +27,43 @@
     Runner,
 )
 from agents.models.multi_provider import MultiProvider
-from agents.items import TResponseStreamEvent, TResponseOutputItem, ModelResponse
-from agents.memory.session import SessionABC
+from agents.items import TResponseStreamEvent, ModelResponse
 from agents.items import TResponseInputItem
-from datetime import timedelta
-from typing import List, Any, AsyncIterator, Optional, cast
-from pydantic import BaseModel
+from typing import Any, AsyncIterator
 
 from restate.exceptions import SdkInternalBaseException
+from restate.ext.turnstile import Turnstile
 from restate.extensions import current_context
-from restate import RunOptions, ObjectContext, TerminalError
+from restate import RunOptions, TerminalError
 
-
-@dataclasses.dataclass
-class LlmRetryOpts:
-    max_attempts: Optional[int] = 10
-    """Max number of attempts (including the initial), before giving up.
-
-    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
-    max_duration: Optional[timedelta] = None
-    """Max duration of retries, before giving up.
-
-    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
-    initial_retry_interval: Optional[timedelta] = timedelta(seconds=1)
-    """Initial interval for the first retry attempt.
-    Retry interval will grow by a factor specified in `retry_interval_factor`.
-
-    If any of the other retry related fields is specified, the default for this field is 50 milliseconds, otherwise restate will fallback to the overall invocation retry policy."""
-    max_retry_interval: Optional[timedelta] = None
-    """Max interval between retries.
-    Retry interval will grow by a factor specified in `retry_interval_factor`.
-
-    The default is 10 seconds."""
-    retry_interval_factor: Optional[float] = None
-    """Exponentiation factor to use when computing the next retry delay.
-
-    If any of the other retry related fields is specified, the default for this field is `2`, meaning retry interval will double at each attempt, otherwise restate will fallback to the overall invocation retry policy."""
-
-
-# The OpenAI ModelResponse class is a dataclass with Pydantic fields.
-# The Restate SDK cannot serialize this. So we turn the ModelResponse int a Pydantic model.
-class RestateModelResponse(BaseModel):
-    output: list[TResponseOutputItem]
-    """A list of outputs (messages, tool calls, etc) generated by the model"""
-
-    usage: Usage
-    """The usage information for the response."""
-
-    response_id: str | None
-    """An ID for the response which can be used to refer to the response in subsequent calls to the
-    model. Not supported by all model providers.
-    If using OpenAI models via the Responses API, this is the `response_id` parameter, and it can
-    be passed to `Runner.run`.
-    """
-
-    def to_input_items(self) -> list[TResponseInputItem]:
-        return [it.model_dump(exclude_unset=True) for it in self.output]  # type: ignore
+from .utils import get_function_call_ids, wrap_agent_tools
+from .models import LlmRetryOpts, RestateModelResponse, State
+from .session import RestateSession
 
 
 class DurableModelCalls(MultiProvider):
     """
     A Restate model provider that wraps the OpenAI SDK's default MultiProvider.
     """
 
-    def __init__(self, llm_retry_opts: LlmRetryOpts | None = None):
+    def __init__(self, state: State, llm_retry_opts: LlmRetryOpts | None = None):
         super().__init__()
         self.llm_retry_opts = llm_retry_opts
+        self.state = state
 
     def get_model(self, model_name: str | None) -> Model:
-        return RestateModelWrapper(super().get_model(model_name or None), self.llm_retry_opts)
+        model = super().get_model(model_name or None)
+        return RestateModelWrapper(model, self.state, self.llm_retry_opts)
 
 
 class RestateModelWrapper(Model):
     """
     A wrapper around the OpenAI SDK's Model that persists LLM calls in the Restate journal.
     """
 
-    def __init__(self, model: Model, llm_retry_opts: LlmRetryOpts | None = None):
+    def __init__(self, model: Model, state: State, llm_retry_opts: LlmRetryOpts | None = None):
         self.model = model
+        self.state = state
         self.model_name = "RestateModelWrapper"
         self.llm_retry_opts = llm_retry_opts if llm_retry_opts is not None else LlmRetryOpts()
 
@@ -133,6 +91,10 @@ async def call_llm() -> RestateModelResponse:
                 retry_interval_factor=self.llm_retry_opts.retry_interval_factor,
             ),
         )
+        # collect function call IDs, to
+        ids = get_function_call_ids(result.output)
+        self.state.turnstile = Turnstile(ids)
+
         # convert back to original ModelResponse
         return ModelResponse(
             output=result.output,
@@ -144,47 +106,6 @@ def stream_response(self, *args, **kwargs) -> AsyncIterator[TResponseStreamEvent
         raise TerminalError("Streaming is not supported in Restate. Use `get_response` instead.")
 
 
-class RestateSession(SessionABC):
-    """Restate session implementation following the Session protocol."""
-
-    def __init__(self):
-        self._items: List[TResponseInputItem] | None = None
-
-    def _ctx(self) -> ObjectContext:
-        return cast(ObjectContext, current_context())
-
-    async def get_items(self, limit: int | None = None) -> List[TResponseInputItem]:
-        """Retrieve conversation history for this session."""
-        if self._items is None:
-            self._items = await self._ctx().get("items") or []
-        if limit is not None:
-            return self._items[-limit:]
-        return self._items.copy()
-
-    async def add_items(self, items: List[TResponseInputItem]) -> None:
-        """Store new items for this session."""
-        if self._items is None:
-            self._items = await self._ctx().get("items") or []
-        self._items.extend(items)
-
-    async def pop_item(self) -> TResponseInputItem | None:
-        """Remove and return the most recent item from this session."""
-        if self._items is None:
-            self._items = await self._ctx().get("items") or []
-        if self._items:
-            return self._items.pop()
-        return None
-
-    def flush(self) -> None:
-        """Flush the session items to the context."""
-        self._ctx().set("items", self._items)
-
-    async def clear_session(self) -> None:
-        """Clear all items for this session."""
-        self._items = []
-        self._ctx().clear("items")
-
-
 class AgentsTerminalException(AgentsException, TerminalError):
     """Exception that is both an AgentsException and a restate.TerminalError."""
 
@@ -255,10 +176,13 @@ async def run(
             The result from Runner.run
         """
 
+        # execution state
+        state = State()
+
         # Set persisting model calls
         llm_retry_opts = kwargs.pop("llm_retry_opts", None)
         run_config = kwargs.pop("run_config", RunConfig())
-        run_config = dataclasses.replace(run_config, model_provider=DurableModelCalls(llm_retry_opts))
+        run_config = dataclasses.replace(run_config, model_provider=DurableModelCalls(state, llm_retry_opts))
 
         # Disable parallel tool calls
         model_settings = run_config.model_settings
@@ -281,9 +205,10 @@ async def run(
                 raise TerminalError("When use_restate_session is True, session config cannot be provided.")
             session = RestateSession()
 
+        agent = wrap_agent_tools(starting_agent, state)
         try:
             result = await Runner.run(
-                starting_agent=starting_agent, input=input, run_config=run_config, session=session, **kwargs
+                starting_agent=agent, input=input, run_config=run_config, session=session, **kwargs
             )
         finally:
             # Flush session items to Restate

python/restate/ext/openai/session.py

@@ -0,0 +1,61 @@
+#
+#  Copyright (c) 2023-2025 - Restate Software, Inc., Restate GmbH
+#
+#  This file is part of the Restate SDK for Python,
+#  which is released under the MIT license.
+#
+#  You can find a copy of the license in file LICENSE in the root
+#  directory of this repository or package, or at
+#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
+#
+"""
+This module contains the optional OpenAI integration for Restate.
+"""
+
+from agents.memory.session import SessionABC
+from agents.items import TResponseInputItem
+from typing import List, cast
+
+from restate.extensions import current_context
+from restate import ObjectContext
+
+
+class RestateSession(SessionABC):
+    """Restate session implementation following the Session protocol."""
+
+    def __init__(self):
+        self._items: List[TResponseInputItem] | None = None
+
+    def _ctx(self) -> ObjectContext:
+        return cast(ObjectContext, current_context())
+
+    async def get_items(self, limit: int | None = None) -> List[TResponseInputItem]:
+        """Retrieve conversation history for this session."""
+        if self._items is None:
+            self._items = await self._ctx().get("items") or []
+        if limit is not None:
+            return self._items[-limit:]
+        return self._items.copy()
+
+    async def add_items(self, items: List[TResponseInputItem]) -> None:
+        """Store new items for this session."""
+        if self._items is None:
+            self._items = await self._ctx().get("items") or []
+        self._items.extend(items)
+
+    async def pop_item(self) -> TResponseInputItem | None:
+        """Remove and return the most recent item from this session."""
+        if self._items is None:
+            self._items = await self._ctx().get("items") or []
+        if self._items:
+            return self._items.pop()
+        return None
+
+    def flush(self) -> None:
+        """Flush the session items to the context."""
+        self._ctx().set("items", self._items)
+
+    async def clear_session(self) -> None:
+        """Clear all items for this session."""
+        self._items = []
+        self._ctx().clear("items")