Merge pull request #1652 from BerriAI/litellm_auto_infer_health_check…

…_updated feat(main.py): support auto-infering mode if not set
BerriAI · Jan 28, 2024 · c86d97f · c86d97f
2 parents 0c9e56a + bcad7c5
commit c86d97f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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