Merge pull request #48 from microsoft/azure_auth_support

Azure auth support

Author: gaganso

README.md

@@ -113,6 +113,8 @@ echo "OPENAI_API_KEY=<YOUR_OPENAI_API_KEY>" > .env
 python3 clients/gpt.py # you can also change the problem to solve in the main() function
 ```
 
+Our repository comes with a variety of pre-integrated agents, including agents that enable **secure authentication with Azure OpenAI endpoints using identity-based access**. Please check out [Clients](/clients) for a comprehensive list of all implemented clients.
+
 The clients will automatically load API keys from your .env file.
 
 You can check the running status of the cluster using [k9s](https://k9scli.io/) or other cluster monitoring tools conveniently.

clients/README.md

@@ -9,6 +9,7 @@ These clients are some baselines that we have implemented and evaluated to help
 - [DeepSeek](/clients/deepseek.py): A naive DeepSeek series LLM agent with only shell access.
 - [Qwen](/clients/qwen.py): A naive Qwen series LLM agent with only shell access.
 - [vLLM](/clients/vllm.py): A naive vLLM agent with any open source LLM deployed locally and only shell access.
+- [GPT with Azure OpenAI](/clients/gpt_azure_identity.py): A naive GPT4-based LLM agent for Azure OpenAI, using identity-based authentication.
 - [ReAct](/clients/react.py): A naive LLM agent that uses the ReAct framework.
 - [FLASH](/clients/flash.py): A naive LLM agent that uses status supervision and hindsight integration components to ensure the high reliability of workflow execution.
 - [OpenRouter](/clients/openrouter.py): A naive OpenRouter LLM agent with only shell access.
@@ -83,21 +84,36 @@ cp .env.example .env
 - `OPENROUTER_MODEL`: OpenRouter model to use (default: `openai/gpt-4o-mini`)
 - `USE_WANDB`: Enable Weights & Biases logging (default: `false`)
 
-<!--
-Note: The script [GPT-managed-identity](/clients/gpt_managed_identity.py) uses the `DefaultAzureCredential` method from the `azure-identity` package to authenticate. This method simplifies authentication by supporting various credential types, including managed identities.
+### Keyless Authentication
 
-We recommend using a [user-assigned managed identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-manage-user-assigned-managed-identities?pivots=identity-mi-methods-azp) for this setup. Ensure the following steps are completed:
+The script [`gpt_azure_identity.py`](/clients/gpt_azure_identity.py) supports keyless authentication for **securely** accessing Azure OpenAI endpoints. It supports two authentication methods:
 
-1. **Role Assignment**: Assign the managed identity appropriate roles:
-    - A role that provides read access to the VM, such as the built-in **Reader** role.
-    - A role that grants read/write access to the Azure OpenAI Service, such as the **Azure AI Developer** role.
+- [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/?view=azure-cli-latest)
+- [User-assigned managed identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-manage-user-assigned-managed-identities?pivots=identity-mi-methods-azp)
 
-2. **Attach the Managed Identity to the Controller VM**:
-    Follow the steps in the official documentation to add the managed identity to the VM:
-    [Add a user-assigned managed identity to a VM](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-configure-managed-identities?pivots=qs-configure-portal-windows-vm#user-assigned-managed-identity).
+#### 1. Azure CLI Authentication
 
-Please ensure the required Azure configuration is provided using the /configs/example_azure_config.yml file, or use it as a template to create a new configuration file
+**Steps**
+- The user must have the appropriate role assigned (e.g., `Cognitive Services OpenAI User`) on the Azure OpenAI resource.
+- Run the following command to authenticate ([How to install the Azure CLI?](https://learn.microsoft.com/en-us/cli/azure/install-azure-cli?view=azure-cli-latest)):
+
+```bash
+az login --scope https://cognitiveservices.azure.com/.default
+```
+
+#### 2. Managed Identity Authentication
+
+- Follow the official documentation to assign a user-assigned managed identity to the VM where the client script would be run:  
+[Add a user-assigned managed identity to a VM](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-configure-managed-identities?pivots=qs-configure-portal-windows-vm#user-assigned-managed-identity)
+- The managed identity must have the appropriate role assigned (e.g., `Cognitive Services OpenAI User`) on the Azure OpenAI resource.
+- Specify the managed identity to use by setting the following environment variable before running the script:
+
+```bash
+export AZURE_CLIENT_ID=<client-id>
+```
+
+Please ensure the required Azure configuration is provided using the /configs/example_azure_config.yml file, or use it as a template to create a new configuration file.
 
 ### Useful Links
-1. [How to configure Azure OpenAI Service with Microsoft Entra ID authentication](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity)
-2. [Azure Identity client library for Python](https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential) -->
+1. [How to configure Azure OpenAI Service with Microsoft Entra ID authentication](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/managed-identity)  
+2. [Azure Identity client library for Python](https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential)

clients/configs/example_azure_config.yml

@@ -1,5 +1,2 @@
-subscription_id: <subscription_id>
-resource_group_name: <resource_group_name>
-workspace_name: <workspace_name>
 azure_endpoint: <azure_endpoint>
 api_version: <api_version>

clients/gpt_managed_identity.py clients/gpt_azure_identity.pyclients/gpt_managed_identity.py renamed to clients/gpt_azure_identity.py

@@ -9,16 +9,17 @@
 
 import sys
 import asyncio
+import os
 
 from aiopslab.orchestrator import Orchestrator
 from clients.utils.llm import GPTClient
 from clients.utils.templates import DOCS_SHELL_ONLY
 
 
 class Agent:
-    def __init__(self, azure_config_file: str):
+    def __init__(self, auth_type: str, azure_config_file: str):
         self.history = []
-        self.llm = GPTClient(auth_type="managed", azure_config_file=azure_config_file)
+        self.llm = GPTClient(auth_type=auth_type, azure_config_file=azure_config_file)
 
     def init_context(self, problem_desc: str, instructions: str, apis: str):
         """Initialize the context for the agent."""
@@ -59,12 +60,24 @@ def _filter_dict(self, dictionary, filter_func):
 
 
 if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        raise Exception(
-            "Please provide a filename as argument. Usage: python gpt_managed_identity.py <azure_config_file>"
-        )
-
-    agent = Agent(azure_config_file=sys.argv[1])
+    #TODO: use argparse for better argument handling
+    script_name = sys.argv[0]
+    if len(sys.argv) < 3:
+        print(f"Error: Missing required arguments.")
+        print(f"\nUsage: python {script_name} <authentication_type> <azure_config_file>")
+        print("\nArguments:")
+        print("  <authentication_type> : Required. The method to use for authentication.")
+        print("                            Accepted values: 'cli', 'managed_identity'")
+        print("  <azure_config_file>     : Required. Path to the JSON configuration file for Azure OpenAI.")
+        print("\nExamples:")
+        print(f"  python {script_name} cli ./configs/dev_config.json", file=sys.stderr)
+        print(f"  python {script_name} managed-identity ./configs/prod_config.json", file=sys.stderr)
+        
+        # Exit with a non-zero status code to indicate an error
+        sys.exit(1)
+    auth_type = sys.argv[1]
+    azure_config_file = sys.argv[2]
+    agent = Agent(auth_type=auth_type, azure_config_file=azure_config_file)
 
     orchestrator = Orchestrator()
     orchestrator.register_agent(agent, name="gpt-w-shell")

clients/utils/llm.py

@@ -1,21 +1,37 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""An common abstraction for a cached LLM inference setup. Currently supports OpenAI's gpt-4-turbo and other models."""
+"""A common abstraction for a cached LLM inference setup. Currently supports OpenAI's gpt-4-turbo and other models."""
 
 
 import os
-from openai import OpenAI
+import json
+import yaml
 from groq import Groq
 from pathlib import Path
-import json
+from typing import Optional, List, Dict
+from dataclasses import dataclass
+
+from groq import Groq
+from openai import OpenAI, AzureOpenAI
+from azure.identity import get_bearer_token_provider, AzureCliCredential, ManagedIdentityCredential
+
 from dotenv import load_dotenv
 
 # Load environment variables from the .env file
 load_dotenv()
+"""An common abstraction for a cached LLM inference setup. Currently supports OpenAI's gpt-4-turbo and other models."""
+
 
 CACHE_DIR = Path("./cache_dir")
 CACHE_PATH = CACHE_DIR / "cache.json"
+GPT_MODEL = "gpt-4o"
+
+
+@dataclass
+class AzureConfig:
+    azure_endpoint: str
+    api_version: str
 
 
 class Cache:
@@ -53,20 +69,58 @@ def save_cache(self):
 class GPTClient:
     """Abstraction for OpenAI's GPT series model."""
 
-    def __init__(self):
+    def __init__(self, auth_type: str = "key", api_key: Optional[str] = None, azure_config_file: Optional[str] = None, use_cache: bool = True):
         self.cache = Cache()
+        self.client = self._setup_client(auth_type, api_key, azure_config_file)
+
+    def _load_azure_config(self, yaml_file_path: str) -> AzureConfig:
+        with open(yaml_file_path, "r") as file:
+            azure_config_data = yaml.safe_load(file)
+            return AzureConfig(
+                azure_endpoint=azure_config_data.get("azure_endpoint"),
+                api_version=azure_config_data.get("api_version"),
+            )
+
+    def _setup_client(self, auth_type: str, api_key: Optional[str], azure_config_file: Optional[str]):
+        azure_identity_opts = ["cli", "managed_identity"]
+        if auth_type == "key":
+            # TODO: support Azure OpenAI client.
+            api_key = api_key or os.getenv("OPENAI_API_KEY")
+            if not api_key:
+                raise ValueError("API key must be provided or set in OPENAI_API_KEY environment variable")
+            return OpenAI(api_key=api_key)
+        elif auth_type in azure_identity_opts:
+            if not azure_config_file:
+                raise ValueError("Azure configuration file must be provided for access via managed identity.\n Check AIOpsLab/clients/configs/example_azure_config.yml for an example.")
+            azure_config = self._load_azure_config(azure_config_file)
+            if auth_type == "cli":
+                credential = AzureCliCredential()
+            elif auth_type == "managed_identity":
+                client_id = os.getenv("AZURE_CLIENT_ID")
+                if client_id is None:
+                    raise ValueError("Managed identity selected but AZURE_CLIENT_ID is not set.")
+                credential = ManagedIdentityCredential(client_id=client_id)
+            token_provider = get_bearer_token_provider(
+                credential, "https://cognitiveservices.azure.com/.default"
+            )
+            return AzureOpenAI(
+                api_version=azure_config.api_version,
+                azure_endpoint=azure_config.azure_endpoint,
+                azure_ad_token_provider=token_provider
+            )
+        else:
+            raise ValueError("auth_type must be one of 'key', 'cli', or 'managed_identity'")
 
     def inference(self, payload: list[dict[str, str]]) -> list[str]:
         if self.cache is not None:
             cache_result = self.cache.get_from_cache(payload)
             if cache_result is not None:
                 return cache_result
 
-        client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
         try:
-            response = client.chat.completions.create(
+            response = self.client.chat.completions.create(
                 messages=payload,  # type: ignore
-                model="gpt-4-turbo-2024-04-09",
+                model=GPT_MODEL,
                 max_tokens=1024,
                 temperature=0.5,
                 top_p=0.95,