From 2a3bc8c19029d7f7584a25631f0858a509c9949a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 10:59:18 -0700 Subject: [PATCH 1/6] add azure_ad_token_provider as all litellm params --- litellm/types/utils.py | 1 + litellm/utils.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8efbe5a1149..62a3a5bcd4a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1116,6 +1116,7 @@ all_litellm_params = [ "cooldown_time", "cache_key", "max_retries", + "azure_ad_token_provider", ] diff --git a/litellm/utils.py b/litellm/utils.py index 0e9e531e98a..ea6b9952425 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2323,6 +2323,7 @@ def get_litellm_params( output_cost_per_second=None, cooldown_time=None, text_completion=None, + azure_ad_token_provider=None, ): litellm_params = { "acompletion": acompletion, @@ -2347,6 +2348,7 @@ def get_litellm_params( "output_cost_per_second": output_cost_per_second, "cooldown_time": cooldown_time, "text_completion": text_completion, + "azure_ad_token_provider": azure_ad_token_provider, } return litellm_params From 8f657b40f50bbd8cc6cfca55e910046216e4a052 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 11:03:49 -0700 Subject: [PATCH 2/6] use azure_ad_token_provider to init clients --- litellm/main.py | 2 ++ litellm/proxy/proxy_config.yaml | 8 +++-- .../client_initalization_utils.py | 30 +++++++++++++++++-- 3 files changed, 35 insertions(+), 5 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 80a9a94a34a..cf37d46d159 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -943,6 +943,7 @@ def completion( output_cost_per_token=output_cost_per_token, cooldown_time=cooldown_time, text_completion=kwargs.get("text_completion"), + azure_ad_token_provider=kwargs.get("azure_ad_token_provider"), ) logging.update_environment_variables( model=model, @@ -3229,6 +3230,7 @@ def embedding( "model_config", "cooldown_time", "tags", + "azure_ad_token_provider", ] default_params = openai_params + litellm_params non_default_params = { diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 65c7f70525c..7d9b97246cc 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,9 +1,11 @@ model_list: - model_name: gpt-4 litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ + model: azure/chatgpt-v-2 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + azure_ad_token_provider: True + guardrails: - guardrail_name: "lakera-pre-guard" diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index f396defb51f..b0b13fac7fa 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -1,7 +1,7 @@ import asyncio import os import traceback -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable import httpx import openai @@ -173,6 +173,11 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): organization = litellm.get_secret(organization_env_name) litellm_params["organization"] = organization + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + if azure_ad_token_provider is not None: + verbose_router_logger.debug("Using Azure AD Token Provider for Azure Auth") + azure_ad_token_provider = get_azure_ad_token_from_entrata_id() + if custom_llm_provider == "azure" or custom_llm_provider == "azure_text": if api_base is None or not isinstance(api_base, str): filtered_litellm_params = { @@ -190,7 +195,9 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): if azure_ad_token.startswith("oidc/"): azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) if api_version is None: - api_version = os.getenv("AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION) + api_version = os.getenv( + "AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION + ) if "gateway.ai.cloudflare.com" in api_base: if not api_base.endswith("/"): @@ -304,6 +311,11 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): "api_version": api_version, "azure_ad_token": azure_ad_token, } + + if azure_ad_token_provider is not None: + azure_client_params["azure_ad_token_provider"] = ( + azure_ad_token_provider + ) from litellm.llms.azure import select_azure_base_url_or_endpoint # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client @@ -493,3 +505,17 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): ttl=client_ttl, local_only=True, ) # cache for 1 hr + + +def get_azure_ad_token_from_entrata_id() -> Callable[[], str]: + from azure.identity import DefaultAzureCredential, get_bearer_token_provider + + verbose_router_logger.debug("Getting Azure AD Token from Entrata ID") + + token_provider = get_bearer_token_provider( + DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" + ) + + verbose_router_logger.debug("token_provider %s", token_provider) + + return token_provider From 08fa3f346aee872048d26fdc5a2674c648a57539 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 11:37:30 -0700 Subject: [PATCH 3/6] add new litellm params for client_id, tenant_id etc --- docs/my-website/docs/providers/azure.md | 127 +++++++++++++++++- litellm/main.py | 3 + litellm/proxy/proxy_config.yaml | 5 +- .../client_initalization_utils.py | 39 +++++- litellm/types/utils.py | 3 + 5 files changed, 168 insertions(+), 9 deletions(-) diff --git a/docs/my-website/docs/providers/azure.md b/docs/my-website/docs/providers/azure.md index be3401fd2e9..8a9603f5ce3 100644 --- a/docs/my-website/docs/providers/azure.md +++ b/docs/my-website/docs/providers/azure.md @@ -1,3 +1,8 @@ + +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + # Azure OpenAI ## API Keys, Params api_key, api_base, api_version etc can be passed directly to `litellm.completion` - see here or set as `litellm.api_key` params see here @@ -12,7 +17,7 @@ os.environ["AZURE_AD_TOKEN"] = "" os.environ["AZURE_API_TYPE"] = "" ``` -## Usage +## **Usage - LiteLLM Python SDK** Open In Colab @@ -64,6 +69,126 @@ response = litellm.completion( ) ``` + +## **Usage - LiteLLM Proxy Server** + +Here's how to call Azure OpenAI models with the LiteLLM Proxy Server + +### 1. Save key in your environment + +```bash +export AZURE_API_KEY="" +``` + +### 2. Start the proxy + + + + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: azure/chatgpt-v-2 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. +``` + + + +This uses the `ClientSecretCredential` + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: azure/chatgpt-v-2 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + tenant_id: os.environ/AZURE_TENANT_ID + client_id: os.environ/AZURE_CLIENT_ID + client_secret: os.environ/AZURE_CLIENT_SECRET +``` + + + + +### 3. Test it + + + + + +```shell +curl --location 'http://0.0.0.0:4000/chat/completions' \ +--header 'Content-Type: application/json' \ +--data ' { + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] + } +' +``` + + + +```python +import openai +client = openai.OpenAI( + api_key="anything", + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create(model="gpt-3.5-turbo", messages = [ + { + "role": "user", + "content": "this is a test request, write a short poem" + } +]) + +print(response) + +``` + + + +```python +from langchain.chat_models import ChatOpenAI +from langchain.prompts.chat import ( + ChatPromptTemplate, + HumanMessagePromptTemplate, + SystemMessagePromptTemplate, +) +from langchain.schema import HumanMessage, SystemMessage + +chat = ChatOpenAI( + openai_api_base="http://0.0.0.0:4000", # set openai_api_base to the LiteLLM Proxy + model = "gpt-3.5-turbo", + temperature=0.1 +) + +messages = [ + SystemMessage( + content="You are a helpful assistant that im using to make a test request to." + ), + HumanMessage( + content="test from litellm. tell me why it's amazing in 1 sentence" + ), +] +response = chat(messages) + +print(response) +``` + + + + + ## Azure OpenAI Chat Completion Models :::tip diff --git a/litellm/main.py b/litellm/main.py index cf37d46d159..45e164a8991 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3231,6 +3231,9 @@ def embedding( "cooldown_time", "tags", "azure_ad_token_provider", + "tenant_id", + "client_id", + "client_secret", ] default_params = openai_params + litellm_params non_default_params = { diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 7d9b97246cc..c8599d56ef8 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -4,8 +4,9 @@ model_list: model: azure/chatgpt-v-2 api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ api_version: "2023-05-15" - azure_ad_token_provider: True - + tenant_id: os.environ/AZURE_TENANT_ID + client_id: os.environ/AZURE_CLIENT_ID + client_secret: os.environ/AZURE_CLIENT_SECRET guardrails: - guardrail_name: "lakera-pre-guard" diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index b0b13fac7fa..1dda941b429 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -173,10 +173,13 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): organization = litellm.get_secret(organization_env_name) litellm_params["organization"] = organization - azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) - if azure_ad_token_provider is not None: + if litellm_params.get("tenant_id"): verbose_router_logger.debug("Using Azure AD Token Provider for Azure Auth") - azure_ad_token_provider = get_azure_ad_token_from_entrata_id() + azure_ad_token_provider = get_azure_ad_token_from_entrata_id( + tenant_id=litellm_params.get("tenant_id"), + client_id=litellm_params.get("client_id"), + client_secret=litellm_params.get("client_secret"), + ) if custom_llm_provider == "azure" or custom_llm_provider == "azure_text": if api_base is None or not isinstance(api_base, str): @@ -507,13 +510,37 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): ) # cache for 1 hr -def get_azure_ad_token_from_entrata_id() -> Callable[[], str]: - from azure.identity import DefaultAzureCredential, get_bearer_token_provider +def get_azure_ad_token_from_entrata_id( + tenant_id: str, client_id: str, client_secret: str +) -> Callable[[], str]: + from azure.identity import ( + ClientSecretCredential, + DefaultAzureCredential, + get_bearer_token_provider, + ) verbose_router_logger.debug("Getting Azure AD Token from Entrata ID") + if tenant_id.startswith("os.environ/"): + tenant_id = litellm.get_secret(tenant_id) + + if client_id.startswith("os.environ/"): + client_id = litellm.get_secret(client_id) + + if client_secret.startswith("os.environ/"): + client_secret = litellm.get_secret(client_secret) + verbose_router_logger.debug( + "tenant_id %s, client_id %s, client_secret %s", + tenant_id, + client_id, + client_secret, + ) + credential = ClientSecretCredential(tenant_id, client_id, client_secret) + + verbose_router_logger.debug("credential %s", credential) + token_provider = get_bearer_token_provider( - DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" + credential, "https://cognitiveservices.azure.com/.default" ) verbose_router_logger.debug("token_provider %s", token_provider) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 62a3a5bcd4a..21eae868a00 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1117,6 +1117,9 @@ all_litellm_params = [ "cache_key", "max_retries", "azure_ad_token_provider", + "tenant_id", + "client_id", + "client_secret", ] From 38ce14c66b318db74fae89ea2a48dc930f64806a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 13:20:05 -0700 Subject: [PATCH 4/6] docs use entrata id with litellm proxy --- docs/my-website/docs/providers/azure.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/my-website/docs/providers/azure.md b/docs/my-website/docs/providers/azure.md index 8a9603f5ce3..dc64bffc1cc 100644 --- a/docs/my-website/docs/providers/azure.md +++ b/docs/my-website/docs/providers/azure.md @@ -95,9 +95,8 @@ model_list: api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. ``` - + -This uses the `ClientSecretCredential` ```yaml model_list: From f13a65be2dd431585c87f2350f25f3338b83b863 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 16:15:53 -0700 Subject: [PATCH 5/6] fix azure_ad_token_provider --- litellm/router_utils/client_initalization_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index 1dda941b429..e98b8b4ddc9 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -172,7 +172,7 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): organization_env_name = organization.replace("os.environ/", "") organization = litellm.get_secret(organization_env_name) litellm_params["organization"] = organization - + azure_ad_token_provider = None if litellm_params.get("tenant_id"): verbose_router_logger.debug("Using Azure AD Token Provider for Azure Auth") azure_ad_token_provider = get_azure_ad_token_from_entrata_id( From 2bf1fb3275771ed429c90a94c9f1e472b4f14b37 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 18:09:00 -0700 Subject: [PATCH 6/6] add test for test_azure_tenant_id_auth --- .circleci/config.yml | 1 + litellm/tests/test_azure_openai.py | 98 ++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 litellm/tests/test_azure_openai.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 24d826f4f61..3f457b30b3e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -51,6 +51,7 @@ jobs: pip install prisma==0.11.0 pip install "detect_secrets==1.5.0" pip install "httpx==0.24.1" + pip install "respx==0.21.1" pip install fastapi pip install "gunicorn==21.2.0" pip install "anyio==3.7.1" diff --git a/litellm/tests/test_azure_openai.py b/litellm/tests/test_azure_openai.py new file mode 100644 index 00000000000..9972f2833d0 --- /dev/null +++ b/litellm/tests/test_azure_openai.py @@ -0,0 +1,98 @@ +import json +import os +import sys +import traceback + +from dotenv import load_dotenv + +load_dotenv() +import io +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import os +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from openai import OpenAI +from openai.types.chat import ChatCompletionMessage +from openai.types.chat.chat_completion import ChatCompletion, Choice +from respx import MockRouter + +import litellm +from litellm import RateLimitError, Timeout, completion, completion_cost, embedding +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.prompt_templates.factory import anthropic_messages_pt +from litellm.router import Router + + +@pytest.mark.asyncio() +@pytest.mark.respx() +async def test_azure_tenant_id_auth(respx_mock: MockRouter): + """ + + Tests when we set tenant_id, client_id, client_secret they don't get sent with the request + + PROD Test + """ + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_base": os.getenv("AZURE_API_BASE"), + "tenant_id": os.getenv("AZURE_TENANT_ID"), + "client_id": os.getenv("AZURE_CLIENT_ID"), + "client_secret": os.getenv("AZURE_CLIENT_SECRET"), + }, + }, + ], + ) + + mock_response = AsyncMock() + obj = ChatCompletion( + id="foo", + model="gpt-4", + object="chat.completion", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage( + content="Hello world!", + role="assistant", + ), + ) + ], + created=int(datetime.now().timestamp()), + ) + mock_request = respx_mock.post(url__regex=r".*/chat/completions.*").mock( + return_value=httpx.Response(200, json=obj.model_dump(mode="json")) + ) + + await router.acompletion( + model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello world!"}] + ) + + # Ensure all mocks were called + respx_mock.assert_all_called() + + for call in mock_request.calls: + print(call) + print(call.request.content) + + json_body = json.loads(call.request.content) + print(json_body) + + assert json_body == { + "messages": [{"role": "user", "content": "Hello world!"}], + "model": "chatgpt-v-2", + "stream": False, + }