From 7c3141a66c73452216cc091e801e18fa1c332f44 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 29 Feb 2024 08:15:47 -0800
Subject: [PATCH] (feat) mistral allow setting API base in env

---
 litellm/utils.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/litellm/utils.py b/litellm/utils.py
index 1106d2fec9a..83764961af6 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -5003,8 +5003,15 @@ 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_API_BASE")
+                    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"):
+                    api_base = api_base + "/v1"
+                dynamic_api_key = api_key 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"