From 51110bfb62ff4c7613c26bf32edb0ffa383f86ef Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 12 Jan 2024 11:24:31 +0530 Subject: [PATCH] fix(main.py): support text completion routing --- litellm/main.py | 6 ++++-- litellm/tests/test_health_check.py | 3 +-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index fa0e2cd4e7..99e5590327 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3066,7 +3066,9 @@ def image_generation( async def ahealth_check( model_params: dict, - mode: Optional[Literal["completion", "embedding", "image_generation"]] = None, + mode: Optional[ + Literal["completion", "embedding", "image_generation", "chat"] + ] = None, prompt: Optional[str] = None, input: Optional[List] = None, default_timeout: float = 6000, @@ -3084,7 +3086,7 @@ async def ahealth_check( raise Exception("model not set") model, custom_llm_provider, _, _ = get_llm_provider(model=model) - mode = mode or "completion" # default to completion calls + mode = mode or "chat" # default to chat completion calls if custom_llm_provider == "azure": api_key = ( diff --git a/litellm/tests/test_health_check.py b/litellm/tests/test_health_check.py index 05166b8888..21b72d2ac3 100644 --- a/litellm/tests/test_health_check.py +++ b/litellm/tests/test_health_check.py @@ -20,8 +20,7 @@ async def test_azure_health_check(): "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), - }, - mode="chat", + } ) print(f"response: {response}")