From bcad7c58a8ac61a0feb03860eb8a34ca2c6cd88e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 25 Jan 2024 20:07:31 -0800
Subject: [PATCH] feat(main.py): support auto-infering mode if not set

---
 litellm/main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index 84b1b79d1..05b352fa6 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -3175,6 +3175,9 @@ async def ahealth_check(
         if model is None:
             raise Exception("model not set")
 
+        if model in litellm.model_cost and mode is None:
+            mode = litellm.model_cost[model]["mode"]
+
         model, custom_llm_provider, _, _ = get_llm_provider(model=model)
         mode = mode or "chat"  # default to chat completion calls