Merge remote-tracking branch 'origin/main' into openai/heartbeat

Author: tconley1428

tests/contrib/openai_agents/test_openai.py

@@ -17,9 +17,11 @@
     ItemHelpers,
     MessageOutputItem,
     Model,
+    ModelProvider,
     ModelResponse,
     ModelSettings,
     ModelTracing,
+    OpenAIChatCompletionsModel,
     OpenAIResponsesModel,
     OutputGuardrailTripwireTriggered,
     RunContextWrapper,
@@ -1844,6 +1846,40 @@ async def test_exception_handling(client: Client):
     await assert_status_retry_behavior(404, client, should_retry=False)
 
 
+class CustomModelProvider(ModelProvider):
+    def get_model(self, model_name: Optional[str]) -> Model:
+        client = AsyncOpenAI(base_url="https://api.openai.com/v1")
+        return OpenAIChatCompletionsModel(model="gpt-4o", openai_client=client)
+
+
+async def test_chat_completions_model(client: Client):
+    if not os.environ.get("OPENAI_API_KEY"):
+        pytest.skip("No openai API key")
+
+    new_config = client.config()
+    new_config["plugins"] = [
+        openai_agents.OpenAIAgentsPlugin(
+            model_params=ModelActivityParameters(
+                start_to_close_timeout=timedelta(seconds=30)
+            ),
+            model_provider=CustomModelProvider(),
+        )
+    ]
+    client = Client(**new_config)
+
+    async with new_worker(
+        client,
+        WorkflowToolWorkflow,
+    ) as worker:
+        workflow_handle = await client.start_workflow(
+            WorkflowToolWorkflow.run,
+            id=f"workflow-tool-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            execution_timeout=timedelta(seconds=10),
+        )
+        await workflow_handle.result()
+
+
 class WaitModel(Model):
     async def get_response(
         self,