From 069f0f3c838622bbdc693ee49cd5f57bdabd400e Mon Sep 17 00:00:00 2001
From: "Tuan Anh Nguyen Dang (Tadashi_Cin)" <tadashi@cinnamon.is>
Date: Fri, 6 Sep 2024 18:18:19 +0700
Subject: [PATCH] feat: expose Cohere and HF embedding support on UI (#236)

---
 flowsettings.py                      | 19 ++++++++++++++++++-
 libs/ktem/ktem/embeddings/manager.py | 10 +++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/flowsettings.py b/flowsettings.py
index e4b95b95..a0ffd0e3 100644
--- a/flowsettings.py
+++ b/flowsettings.py
@@ -167,7 +167,7 @@
         "default": False,
     }
 
-    KH_EMBEDDINGS["local-bge-en"] = {
+    KH_EMBEDDINGS["fast_embed"] = {
         "spec": {
             "__type__": "kotaemon.embeddings.FastEmbedEmbeddings",
             "model_name": "BAAI/bge-base-en-v1.5",
@@ -202,6 +202,23 @@
     "default": False,
 }
 
+# additional embeddings configurations
+KH_EMBEDDINGS["cohere"] = {
+    "spec": {
+        "__type__": "kotaemon.embeddings.LCCohereEmbeddings",
+        "model": "embed-multilingual-v2.0",
+        "cohere_api_key": "your-key",
+    },
+    "default": False,
+}
+# KH_EMBEDDINGS["huggingface"] = {
+#     "spec": {
+#         "__type__": "kotaemon.embeddings.LCHuggingFaceEmbeddings",
+#         "model_name": "sentence-transformers/all-mpnet-base-v2",
+#     },
+#     "default": False,
+# }
+
 KH_REASONINGS = [
     "ktem.reasoning.simple.FullQAPipeline",
     "ktem.reasoning.simple.FullDecomposeQAPipeline",
diff --git a/libs/ktem/ktem/embeddings/manager.py b/libs/ktem/ktem/embeddings/manager.py
index 4948aeb5..f1ff6a00 100644
--- a/libs/ktem/ktem/embeddings/manager.py
+++ b/libs/ktem/ktem/embeddings/manager.py
@@ -55,10 +55,18 @@ def load_vendors(self):
         from kotaemon.embeddings import (
             AzureOpenAIEmbeddings,
             FastEmbedEmbeddings,
+            LCCohereEmbeddings,
+            LCHuggingFaceEmbeddings,
             OpenAIEmbeddings,
         )
 
-        self._vendors = [AzureOpenAIEmbeddings, OpenAIEmbeddings, FastEmbedEmbeddings]
+        self._vendors = [
+            AzureOpenAIEmbeddings,
+            OpenAIEmbeddings,
+            FastEmbedEmbeddings,
+            LCCohereEmbeddings,
+            LCHuggingFaceEmbeddings,
+        ]
 
     def __getitem__(self, key: str) -> BaseEmbeddings:
         """Get model by name"""