utils.validate_environment to handle OLLAMA_API_BASE env

litellm/tests/test_utils.py

@@ -1,4 +1,6 @@
-import sys, os
+import sys
+from unittest import mock
+
 from dotenv import load_dotenv
 import copy
 
@@ -214,8 +216,19 @@ def test_validate_environment_empty_model():
         raise Exception()
 
 
-# test_validate_environment_empty_model()
+@mock.patch.dict(os.environ, {"OLLAMA_API_BASE": "foo"}, clear=True)
+def test_validate_environment_ollama():
+    for provider in ["ollama", "ollama_chat"]:
+        kv = validate_environment(provider+"/mistral")
+        assert kv["keys_in_environment"]
+        assert kv["missing_keys"] == []
 
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_validate_environment_ollama_failed():
+    for provider in ["ollama", "ollama_chat"]:
+        kv = validate_environment(provider+"/mistral")
+        assert not kv["keys_in_environment"]
+        assert kv["missing_keys"] == ["OLLAMA_API_BASE"]
 
 def test_function_to_dict():
     print("testing function to dict for get current weather")

litellm/utils.py

@@ -5602,6 +5602,11 @@ def validate_environment(model: Optional[str] = None) -> dict:
             else:
                 missing_keys.append("AWS_ACCESS_KEY_ID")
                 missing_keys.append("AWS_SECRET_ACCESS_KEY")
+        elif custom_llm_provider in ["ollama", "ollama_chat"]:
+            if "OLLAMA_API_BASE" in os.environ:
+                keys_in_environment = True
+            else:
+                missing_keys.append("OLLAMA_API_BASE")
     else:
         ## openai - chatcompletion + text completion
         if (