From fd8687616412f42f08ec42029e99bfbf8215f3dc Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 6 Dec 2023 11:29:59 -0800
Subject: [PATCH] (feat) proxy: add mode in model info

---
 litellm/health_check.py         | 3 ++-
 litellm/proxy/proxy_config.yaml | 6 ++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/litellm/health_check.py b/litellm/health_check.py
index dabe7215e1..db07ae20a7 100644
--- a/litellm/health_check.py
+++ b/litellm/health_check.py
@@ -68,11 +68,12 @@ async def _perform_health_check(model_list: list):
     tasks = []
     for model in model_list:
         litellm_params = model["litellm_params"]
+        model_info = model.get("model_info", {})
         litellm_params["model"] = litellm.utils.remove_model_id(litellm_params["model"])
         litellm_params["messages"] = _get_random_llm_message()
 
         prepped_params.append(litellm_params)
-        if model.get("mode", None) == "embedding":
+        if model_info.get("mode", None) == "embedding":
             # this is an embedding model
             tasks.append(_check_embedding_model(litellm_params))
         else:
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 86191b0950..6d4e5b2e9d 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -11,7 +11,8 @@ 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_info:
+      mode: embedding
   - model_name: openai-gpt-3.5
     litellm_params:
       model: gpt-3.5-turbo
@@ -20,7 +21,8 @@ model_list:
     litellm_params:
       model: text-embedding-ada-002
       api_key: os.environ/OPENAI_API_KEY
-    mode: embedding
+    model_info:
+      mode: embedding
 
 
 litellm_settings: