feat: update OpenAIAssistantAgent to support AsyncAzureOpenAI client (#…

…5312)

Resolves #5179

python/packages/autogen-ext/src/autogen_ext/agents/openai/_openai_assistant_agent.py

@@ -37,7 +37,7 @@
 from autogen_core.models._types import FunctionExecutionResult
 from autogen_core.tools import FunctionTool, Tool
 
-from openai import NOT_GIVEN, AsyncClient, NotGiven
+from openai import NOT_GIVEN, AsyncAzureOpenAI, AsyncOpenAI, NotGiven
 from openai.pagination import AsyncCursorPage
 from openai.resources.beta.threads import AsyncMessages, AsyncRuns, AsyncThreads
 from openai.types import FileObject
@@ -78,16 +78,17 @@ def _convert_tool_to_function_param(tool: Tool) -> "FunctionToolParam":
 
 
 class OpenAIAssistantAgent(BaseChatAgent):
-    """An agent implementation that uses the OpenAI Assistant API to generate responses.
+    """An agent implementation that uses the Assistant API to generate responses.
 
     Installation:
 
     .. code-block:: bash
 
         pip install "autogen-ext[openai]"
+        # pip install "autogen-ext[openai,azure]"  # For Azure OpenAI Assistant
 
 
-    This agent leverages the OpenAI Assistant API to create AI assistants with capabilities like:
+    This agent leverages the Assistant API to create AI assistants with capabilities like:
 
     * Code interpretation and execution
     * File handling and search
@@ -111,10 +112,13 @@ class OpenAIAssistantAgent(BaseChatAgent):
 
     You can use an existing thread or assistant by providing the `thread_id` or `assistant_id` parameters.
 
-    Example:
+    Examples:
+
+        Use the assistant to analyze data in a CSV file:
+
         .. code-block:: python
 
-            from openai import AsyncClient
+            from openai import AsyncOpenAI
             from autogen_core import CancellationToken
             import asyncio
             from autogen_ext.agents.openai import OpenAIAssistantAgent
@@ -125,7 +129,7 @@ async def example():
                 cancellation_token = CancellationToken()
 
                 # Create an OpenAI client
-                client = AsyncClient(api_key="your-api-key", base_url="your-base-url")
+                client = AsyncOpenAI(api_key="your-api-key", base_url="your-base-url")
 
                 # Create an assistant with code interpreter
                 assistant = OpenAIAssistantAgent(
@@ -154,10 +158,55 @@ async def example():
 
             asyncio.run(example())
 
+        Use Azure OpenAI Assistant with AAD authentication:
+
+        .. code-block:: python
+
+            from openai import AsyncAzureOpenAI
+            import asyncio
+            from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+            from autogen_core import CancellationToken
+            from autogen_ext.agents.openai import OpenAIAssistantAgent
+            from autogen_agentchat.messages import TextMessage
+
+
+            async def example():
+                cancellation_token = CancellationToken()
+
+                # Create an Azure OpenAI client
+                token_provider = get_bearer_token_provider(DefaultAzureCredential())
+                client = AsyncAzureOpenAI(
+                    azure_deployment="YOUR_AZURE_DEPLOYMENT",
+                    api_version="YOUR_API_VERSION",
+                    azure_endpoint="YOUR_AZURE_ENDPOINT",
+                    azure_ad_token_provider=token_provider,
+                )
+
+                # Create an assistant with code interpreter
+                assistant = OpenAIAssistantAgent(
+                    name="Python Helper",
+                    description="Helps with Python programming",
+                    client=client,
+                    model="gpt-4o",
+                    instructions="You are a helpful Python programming assistant.",
+                    tools=["code_interpreter"],
+                )
+
+                # Get response from the assistant
+                response = await assistant.on_messages([TextMessage(source="user", content="Hello.")], cancellation_token)
+
+                print(response)
+
+                # Clean up resources
+                await assistant.delete_assistant(cancellation_token)
+
+
+            asyncio.run(example())
+
     Args:
         name (str): Name of the assistant
         description (str): Description of the assistant's purpose
-        client (AsyncClient): OpenAI API client instance
+        client (AsyncOpenAI | AsyncAzureOpenAI): OpenAI client or Azure OpenAI client instance
         model (str): Model to use (e.g. "gpt-4")
         instructions (str): System instructions for the assistant
         tools (Optional[Iterable[Union[Literal["code_interpreter", "file_search"], Tool | Callable[..., Any] | Callable[..., Awaitable[Any]]]]]): Tools the assistant can use
@@ -174,7 +223,7 @@ def __init__(
         self,
         name: str,
         description: str,
-        client: "AsyncClient",
+        client: AsyncOpenAI | AsyncAzureOpenAI,
         model: str,
         instructions: str,
         tools: Optional[