diff --git a/README.md b/README.md index 215877b..13fd3f9 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,10 @@ python3 -m venv .venv source .venv/bin/activate pip install -r requirements.txt ``` -Set up the `OPENAI_API_KEY` environment variable (for OpenAI) or the `AZURE_OPENAI_API_KEY`, `AZURE_API_VERSION` and `AZURE_API_BASE` environment variables (for Azure OpenAI). +Set up the `OPENAI_API_KEY` environment variable (for OpenAI) +or the `AZURE_OPENAI_API_KEY`, `AZURE_API_VERSION` and `AZURE_API_BASE` environment variables (for Azure OpenAI) +or `AZURE_API_VERSION` and `AZURE_API_BASE` environment variables + `az login` for Azure OpenAI with Managed Entra. + You can then run the optillm proxy as follows. ```bash diff --git a/optillm.py b/optillm.py index 1044b7a..f3b4eb7 100644 --- a/optillm.py +++ b/optillm.py @@ -31,11 +31,24 @@ API_KEY = os.environ.get("OPENAI_API_KEY") default_client = OpenAI(api_key=API_KEY) else: - default_client = AzureOpenAI( - api_key=os.environ.get("AZURE_OPENAI_API_KEY"), - api_version=os.environ.get("AZURE_API_VERSION"), - azure_endpoint=os.environ.get("AZURE_API_BASE"), -) + API_KEY = os.environ.get("AZURE_OPENAI_API_KEY") + API_VERSION = os.environ.get("AZURE_API_VERSION") + AZURE_ENDPOINT = os.environ.get("AZURE_API_BASE") + if API_KEY is not None: + default_client = AzureOpenAI( + api_key=API_KEY, + api_version=API_VERSION, + azure_endpoint=AZURE_ENDPOINT, + ) + else: + from azure.identity import DefaultAzureCredential, get_bearer_token_provider + azure_credential = DefaultAzureCredential() + token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default") + default_client = AzureOpenAI( + api_version=API_VERSION, + azure_endpoint=AZURE_ENDPOINT, + azure_ad_token_provider=token_provider + ) # Server configuration server_config = { diff --git a/requirements.txt b/requirements.txt index 2bbd104..834a056 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ networkx openai z3-solver aiohttp -flask \ No newline at end of file +flask +azure.identity \ No newline at end of file