Add async context manager to OpenAI model client

Author: mikemikimike

python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

@@ -1142,6 +1142,12 @@ async def _create_stream_chunks_beta_client(
     async def close(self) -> None:
         await self._client.close()
 
+    async def __aenter__(self) -> "BaseOpenAIChatCompletionClient":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        await self.close()
+
     def actual_usage(self) -> RequestUsage:
         return self._actual_usage
 

python/packages/autogen-ext/tests/models/test_openai_model_client.py

@@ -78,6 +78,17 @@ class MyArgs(BaseModel):
     query: str = Field(description="The description.")
 
 
+@pytest.mark.asyncio
+async def test_openai_client_is_async_context_manager() -> None:
+    client = OpenAIChatCompletionClient(model="gpt-4o", api_key="test")
+    client.close = AsyncMock()  # type: ignore[method-assign]
+
+    async with client as entered:
+        assert entered is client
+
+    client.close.assert_awaited_once()
+
+
 class MockChunkDefinition(BaseModel):
     # defining elements for diffentiating mocking chunks
     chunk_choice: ChunkChoice