diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py index 5da214dbc5..2e5ac889fb 100644 --- a/litellm/tests/test_utils.py +++ b/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") diff --git a/litellm/utils.py b/litellm/utils.py index 6d64368f87..3603495cd8 100644 --- a/litellm/utils.py +++ b/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 (