Merge pull request #2247 from BerriAI/litellm_mistral_azure_ai_fixes

[FEAT]  mistral azure - use with env vars + added pricing

docs/my-website/docs/providers/azure_ai.md

@@ -2,16 +2,36 @@
 
 ## Using Mistral models deployed on Azure AI Studio
 
-**Ensure you have the `/v1` in your api_base**
+### Sample Usage - setting env vars 
+
+Set `MISTRAL_AZURE_API_KEY` and `MISTRAL_AZURE_API_BASE` in your env
+
+```shell
+MISTRAL_AZURE_API_KEY = "zE************""
+MISTRAL_AZURE_API_BASE = "https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com"
+```
 
-### Sample Usage
 ```python
 from litellm import completion
 import os
 
 response = completion(
     model="mistral/Mistral-large-dfgfj", 
-    api_base="https://Mistral-large-dfgfj-serverless.eastus2.inference.ai.azure.com/v1",
+    messages=[
+       {"role": "user", "content": "hello from litellm"}
+   ],
+)
+print(response)
+```
+
+### Sample Usage - passing `api_base` and `api_key` to `litellm.completion`
+```python
+from litellm import completion
+import os
+
+response = completion(
+    model="mistral/Mistral-large-dfgfj", 
+    api_base="https://Mistral-large-dfgfj-serverless.eastus2.inference.ai.azure.com",
     api_key = "JGbKodRcTp****"
     messages=[
        {"role": "user", "content": "hello from litellm"}
@@ -23,14 +43,12 @@ print(response)
 ### [LiteLLM Proxy] Using Mistral Models 
 
 Set this on your litellm proxy config.yaml
-
-**Ensure you have the `/v1` in your api_base**
 ```yaml
 model_list:
   - model_name: mistral
     litellm_params:
       model: mistral/Mistral-large-dfgfj
-      api_base: https://Mistral-large-dfgfj-serverless.eastus2.inference.ai.azure.com/v1
+      api_base: https://Mistral-large-dfgfj-serverless.eastus2.inference.ai.azure.com
       api_key: JGbKodRcTp****
 ```
 

litellm/model_prices_and_context_window_backup.json

@@ -424,6 +424,22 @@
         "mode": "chat",
         "supports_function_calling": true
     },
+    "azure/mistral-large-latest": {
+        "max_tokens": 32000,
+        "input_cost_per_token": 0.000008,
+        "output_cost_per_token": 0.000024,
+        "litellm_provider": "azure",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
+    "azure/mistral-large-2402": {
+        "max_tokens": 32000,
+        "input_cost_per_token": 0.000008,
+        "output_cost_per_token": 0.000024,
+        "litellm_provider": "azure",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
     "azure/ada": {
         "max_tokens": 8191,
         "input_cost_per_token": 0.0000001,

litellm/tests/test_completion.py

@@ -107,6 +107,31 @@ def test_completion_mistral_api():
         pytest.fail(f"Error occurred: {e}")
 
 
+@pytest.mark.skip(
+    reason="Since we already test mistral/mistral-tiny in test_completion_mistral_api. This is only for locally verifying azure mistral works"
+)
+def test_completion_mistral_azure():
+    try:
+        litellm.set_verbose = True
+        response = completion(
+            model="mistral/Mistral-large-nmefg",
+            api_key=os.environ["MISTRAL_AZURE_API_KEY"],
+            api_base=os.environ["MISTRAL_AZURE_API_BASE"],
+            max_tokens=5,
+            messages=[
+                {
+                    "role": "user",
+                    "content": "Hi from litellm",
+                }
+            ],
+        )
+        # Add any assertions here to check the response
+        print(response)
+
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 # test_completion_mistral_api()
 
 

litellm/utils.py

@@ -5003,8 +5003,21 @@ def get_llm_provider(
                 dynamic_api_key = get_secret("GROQ_API_KEY")
             elif custom_llm_provider == "mistral":
                 # mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai
-                api_base = api_base or "https://api.mistral.ai/v1"
-                dynamic_api_key = get_secret("MISTRAL_API_KEY")
+                api_base = (
+                    api_base
+                    or get_secret("MISTRAL_AZURE_API_BASE")  # for Azure AI Mistral
+                    or "https://api.mistral.ai/v1"
+                )
+                # if api_base does not end with /v1 we add it
+                if api_base is not None and not api_base.endswith(
+                    "/v1"
+                ):  # Mistral always needs a /v1 at the end
+                    api_base = api_base + "/v1"
+                dynamic_api_key = (
+                    api_key
+                    or get_secret("MISTRAL_AZURE_API_KEY")  # for Azure AI Mistral
+                    or get_secret("MISTRAL_API_KEY")
+                )
             elif custom_llm_provider == "voyage":
                 # voyage is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.voyageai.com/v1
                 api_base = "https://api.voyageai.com/v1"

model_prices_and_context_window.json

@@ -424,6 +424,22 @@
         "mode": "chat",
         "supports_function_calling": true
     },
+    "azure/mistral-large-latest": {
+        "max_tokens": 32000,
+        "input_cost_per_token": 0.000008,
+        "output_cost_per_token": 0.000024,
+        "litellm_provider": "azure",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
+    "azure/mistral-large-2402": {
+        "max_tokens": 32000,
+        "input_cost_per_token": 0.000008,
+        "output_cost_per_token": 0.000024,
+        "litellm_provider": "azure",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
     "azure/ada": {
         "max_tokens": 8191,
         "input_cost_per_token": 0.0000001,