Update OpenAI integration: flush session at end of runner, allow setting model retry opts, set parallel_tool_calls to false

Author: gvdongen

python/restate/ext/openai/__init__.py

@@ -16,7 +16,6 @@
 
 from .runner_wrapper import (
     DurableRunner,
-    DurableModelCalls,
     continue_on_terminal_errors,
     raise_terminal_errors,
     RestateSession,
@@ -43,12 +42,11 @@ def restate_context() -> Context:
 
 
 __all__ = [
-    "DurableModelCalls",
-    "continue_on_terminal_errors",
-    "raise_terminal_errors",
-    "RestateSession",
     "DurableRunner",
+    "RestateSession",
     "LlmRetryOpts",
     "restate_object_context",
     "restate_context",
+    "continue_on_terminal_errors",
+    "raise_terminal_errors",
 ]

python/restate/ext/openai/runner_wrapper.py

@@ -93,7 +93,7 @@ class DurableModelCalls(MultiProvider):
     A Restate model provider that wraps the OpenAI SDK's default MultiProvider.
     """
 
-    def __init__(self, llm_retry_opts: LlmRetryOpts):
+    def __init__(self, llm_retry_opts: LlmRetryOpts | None = None):
         super().__init__()
         self.llm_retry_opts = llm_retry_opts
 
@@ -106,7 +106,7 @@ 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):
+    def __init__(self, model: Model, llm_retry_opts: LlmRetryOpts | None = LlmRetryOpts()):
         self.model = model
         self.model_name = "RestateModelWrapper"
         self.llm_retry_opts = llm_retry_opts
@@ -124,7 +124,6 @@ async def call_llm() -> RestateModelResponse:
         ctx = current_context()
         if ctx is None:
             raise RuntimeError("No current Restate context found, make sure to run inside a Restate handler")
-        print("Calling LLM with retry options:", self.llm_retry_opts)
         result = await ctx.run_typed(
             "call LLM",
             call_llm,
@@ -255,7 +254,7 @@ async def run(
         """
 
         # Set persisting model calls
-        llm_retry_opts = kwargs.get("llm_retry_opts", LlmRetryOpts())
+        llm_retry_opts = kwargs.get("llm_retry_opts", None)
         run_config = kwargs.pop("run_config", RunConfig())
         run_config = dataclasses.replace(run_config, model_provider=DurableModelCalls(llm_retry_opts))