diff --git a/litellm/health_check.py b/litellm/health_check.py index e814a8c60..dabe7215e 100644 --- a/litellm/health_check.py +++ b/litellm/health_check.py @@ -50,7 +50,6 @@ async def _perform_health_check(model_list: list): try: await litellm.aembedding(**model_params) except Exception as e: - print_verbose(f"\n\n Got Exception, {e}") print_verbose(f"Health check failed for model {model_params['model']}. Error: {e}") return False return True @@ -59,12 +58,7 @@ async def _perform_health_check(model_list: list): async def _check_model(model_params: dict): try: await litellm.acompletion(**model_params) - except Exception as e: - print_verbose(f"\n\n Got Exception, {e}") - error_str = (str(e)) - if "This is not a chat model" in error_str or "The chatCompletion operation does not work with the specified model" in error_str: - return await _check_embedding_model(model_params) - + except Exception as e: print_verbose(f"Health check failed for model {model_params['model']}. Error: {e}") return False @@ -78,7 +72,8 @@ async def _perform_health_check(model_list: list): litellm_params["messages"] = _get_random_llm_message() prepped_params.append(litellm_params) - if "embedding" in litellm_params["model"]: + if model.get("mode", None) == "embedding": + # this is an embedding model tasks.append(_check_embedding_model(litellm_params)) else: tasks.append(_check_model(litellm_params)) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index b5a9050ad..86191b095 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -11,6 +11,7 @@ model_list: api_base: os.environ/AZURE_API_BASE api_key: os.environ/AZURE_API_KEY api_version: "2023-07-01-preview" + mode: embedding - model_name: openai-gpt-3.5 litellm_params: model: gpt-3.5-turbo @@ -19,6 +20,7 @@ model_list: litellm_params: model: text-embedding-ada-002 api_key: os.environ/OPENAI_API_KEY + mode: embedding litellm_settings: diff --git a/litellm/proxy/types.py b/litellm/proxy/types.py index 3f7ae1db5..5d4669d7d 100644 --- a/litellm/proxy/types.py +++ b/litellm/proxy/types.py @@ -40,6 +40,7 @@ class ProxyChatCompletionRequest(BaseModel): class ModelInfo(BaseModel): id: Optional[str] + mode: Optional[str] class Config: extra = Extra.allow # Allow extra fields