From 05a4555fe73bb727d5443b3f648e585f8de1d6f1 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 9 Aug 2024 12:11:02 -0700
Subject: [PATCH] ui add cohere embedding models

---
 .../src/components/model_dashboard.tsx        | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx
index d0d4b66a09e..de92805b0bf 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx
@@ -930,7 +930,26 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
             _providerModels.push(key);
           }
         });
+
+        // Special case for cohere_chat
+        // we need both cohere_chat and cohere models to show on dropdown
+        if (providerKey == Providers.Cohere) {
+          console.log("adding cohere chat model")
+          Object.entries(modelMap).forEach(([key, value]) => {
+            if (
+              value !== null &&
+              typeof value === "object" &&
+              "litellm_provider" in (value as object) &&
+              ((value as any)["litellm_provider"] === "cohere")
+            ) {
+              _providerModels.push(key);
+            }
+          });
+        }
       }
+
+      
+
       setProviderModels(_providerModels);
       console.log(`providerModels: ${providerModels}`);
     }
@@ -1787,7 +1806,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
                   </Row>
                   <Form.Item
                   label="LiteLLM Model Name(s)"
-                  tooltip="Actual model name used for making litellm.completion() call."
+                  tooltip="Actual model name used for making litellm.completion() / litellm.embedding() call."
                   className="mb-0"
                 >
                   <Form.Item