phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

utils.validate_environment to handle OLLAMA_API_BASE env

Browse source
This commit is contained in:
Mikhail Khludnev 2024-03-02 13:41:59 +03:00
parent d494941dcc
commit 381858139e
2 changed files with 20 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_utils.py
View file

@ -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")

5
litellm/utils.py
View file

@ -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 (