From 278f61f3ed7d02eb710f819feb61b50ead33b181 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 22 Dec 2023 11:08:44 +0530 Subject: [PATCH] fix(utils.py): handle 'os.environ/' being passed in as kwargs --- litellm/proxy/_health_check_test_config.yaml | 17 +++++++++++++ litellm/proxy/health_check.py | 14 +++++++++-- .../test_configs/test_config_no_auth.yaml | 25 +++++++++++++++++-- litellm/tests/test_proxy_server.py | 12 +++++++++ litellm/utils.py | 4 +++ 5 files changed, 68 insertions(+), 4 deletions(-) create mode 100644 litellm/proxy/_health_check_test_config.yaml diff --git a/litellm/proxy/_health_check_test_config.yaml b/litellm/proxy/_health_check_test_config.yaml new file mode 100644 index 000000000..56d532b1d --- /dev/null +++ b/litellm/proxy/_health_check_test_config.yaml @@ -0,0 +1,17 @@ +model_list: + - model_name: text-embedding-ada-002 + litellm_params: + model: azure/azure-embedding-model + api_base: "os.environ/AZURE_API_BASE" + api_key: "os.environ/AZURE_API_KEY" + api_version: "2023-07-01-preview" + model_info: + mode: embedding + base_model: text-embedding-ada-002 + +litellm_settings: + set_verbose: True + +general_settings: + background_health_checks: True # enable background health checks + health_check_interval: 300 # frequency of background health checks \ No newline at end of file diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 24ace2e94..3c7ff599e 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -44,6 +44,16 @@ async def _perform_health_check(model_list: list): """ Perform a health check for each model in the list. """ + async def _check_img_gen_model(model_params: dict): + model_params.pop("messages", None) + model_params["prompt"] = "test from litellm" + try: + await litellm.aimage_generation(**model_params) + except Exception as e: + print_verbose(f"Health check failed for model {model_params['model']}. Error: {e}") + return False + return True + async def _check_embedding_model(model_params: dict): model_params.pop("messages", None) model_params["input"] = ["test from litellm"] @@ -64,17 +74,17 @@ async def _perform_health_check(model_list: list): return True - prepped_params = [] tasks = [] for model in model_list: litellm_params = model["litellm_params"] model_info = model.get("model_info", {}) litellm_params["messages"] = _get_random_llm_message() - prepped_params.append(litellm_params) if model_info.get("mode", None) == "embedding": # this is an embedding model tasks.append(_check_embedding_model(litellm_params)) + elif model_info.get("mode", None) == "image_generation": + tasks.append(_check_img_gen_model(litellm_params)) else: tasks.append(_check_model(litellm_params)) diff --git a/litellm/tests/test_configs/test_config_no_auth.yaml b/litellm/tests/test_configs/test_config_no_auth.yaml index 76e7a294b..e3bf91456 100644 --- a/litellm/tests/test_configs/test_config_no_auth.yaml +++ b/litellm/tests/test_configs/test_config_no_auth.yaml @@ -49,6 +49,8 @@ model_list: api_version: 2023-07-01-preview model: azure/azure-embedding-model model_name: azure-embedding-model + model_info: + mode: "embedding" - litellm_params: model: gpt-3.5-turbo model_info: @@ -76,21 +78,40 @@ model_list: - model_name: amazon-embeddings litellm_params: model: "bedrock/amazon.titan-embed-text-v1" + model_info: + mode: embedding - model_name: "GPT-J 6B - Sagemaker Text Embedding (Internal)" litellm_params: model: "sagemaker/berri-benchmarking-gpt-j-6b-fp16" + model_info: + mode: embedding - model_name: dall-e-3 litellm_params: - model: dall-e-3 + model: dall-e-3 + model_info: + mode: image_generation - model_name: dall-e-3 litellm_params: model: "azure/dall-e-3-test" api_version: "2023-12-01-preview" api_base: "os.environ/AZURE_SWEDEN_API_BASE" api_key: "os.environ/AZURE_SWEDEN_API_KEY" + model_info: + mode: image_generation - model_name: dall-e-2 litellm_params: model: "azure/" api_version: "2023-06-01-preview" api_base: "os.environ/AZURE_API_BASE" - api_key: "os.environ/AZURE_API_KEY" \ No newline at end of file + api_key: "os.environ/AZURE_API_KEY" + model_info: + mode: image_generation +- model_name: text-embedding-ada-002 + litellm_params: + model: azure/azure-embedding-model + api_base: "os.environ/AZURE_API_BASE" + api_key: "os.environ/AZURE_API_KEY" + api_version: "2023-07-01-preview" + model_info: + mode: embedding + base_model: text-embedding-ada-002 \ No newline at end of file diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index b71f5b890..e38338314 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -211,6 +211,18 @@ def test_add_new_model(client_no_auth): except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}") + +def test_health(client_no_auth): + global headers + import time + try: + response = client_no_auth.get("/health") + assert response.status_code == 200 + result = response.json() + assert result["unhealthy_count"] == 0 + raise Exception(f"It worked!") + except Exception as e: + pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") # test_add_new_model() from litellm.integrations.custom_logger import CustomLogger diff --git a/litellm/utils.py b/litellm/utils.py index 3a48958fc..87c543fc8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1627,6 +1627,10 @@ def client(original_function): logging_obj = function_setup(start_time, *args, **kwargs) kwargs["litellm_logging_obj"] = logging_obj + # CHECK FOR 'os.environ/' in kwargs + for k,v in kwargs.items(): + if v is not None and isinstance(v, str) and v.startswith("os.environ/"): + kwargs[k] = litellm.get_secret(v) # [OPTIONAL] CHECK BUDGET if litellm.max_budget: if litellm._current_cost > litellm.max_budget: