feat: support openllm embedding (#1293)

api/core/model_providers/models/embedding/openllm_embedding.py

@@ -0,0 +1,22 @@
+from core.third_party.langchain.embeddings.openllm_embedding import OpenLLMEmbeddings
+
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import BaseModelProvider
+from core.model_providers.models.embedding.base import BaseEmbedding
+
+
+class OpenLLMEmbedding(BaseEmbedding):
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = OpenLLMEmbeddings(
+            server_url=credentials['server_url']
+        )
+
+        super().__init__(model_provider, client, name)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"OpenLLM embedding: {str(ex)}")

api/core/model_providers/models/embedding/xinference_embedding.py

@@ -1,5 +1,4 @@
 from core.third_party.langchain.embeddings.xinference_embedding import XinferenceEmbedding as XinferenceEmbeddings
-from replicate.exceptions import ModelError, ReplicateError
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.providers.base import BaseModelProvider
@@ -21,7 +20,4 @@ class XinferenceEmbedding(BaseEmbedding):
         super().__init__(model_provider, client, name)
 
     def handle_exceptions(self, ex: Exception) -> Exception:
-        if isinstance(ex, (ModelError, ReplicateError)):
-            return LLMBadRequestError(f"Xinference embedding: {str(ex)}")
-        else:
-            return ex
+        return LLMBadRequestError(f"Xinference embedding: {str(ex)}")

api/core/model_providers/providers/openllm_provider.py

@@ -2,11 +2,13 @@ import json
 from typing import Type
 
 from core.helper import encrypter
+from core.model_providers.models.embedding.openllm_embedding import OpenLLMEmbedding
 from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType
 from core.model_providers.models.llm.openllm_model import OpenLLMModel
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 
 from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.embeddings.openllm_embedding import OpenLLMEmbeddings
 from core.third_party.langchain.llms.openllm import OpenLLM
 from models.provider import ProviderType
 
@@ -31,6 +33,8 @@ class OpenLLMProvider(BaseModelProvider):
         """
         if model_type == ModelType.TEXT_GENERATION:
             model_class = OpenLLMModel
+        elif model_type== ModelType.EMBEDDINGS:
+            model_class = OpenLLMEmbedding
         else:
             raise NotImplementedError
 
@@ -69,14 +73,21 @@ class OpenLLMProvider(BaseModelProvider):
                 'server_url': credentials['server_url']
             }
 
-            llm = OpenLLM(
-                llm_kwargs={
-                    'max_new_tokens': 10
-                },
-                **credential_kwargs
-            )
+            if model_type == ModelType.TEXT_GENERATION:
+                llm = OpenLLM(
+                    llm_kwargs={
+                        'max_new_tokens': 10
+                    },
+                    **credential_kwargs
+                )
 
-            llm("ping")
+                llm("ping")
+            elif model_type == ModelType.EMBEDDINGS:
+                embedding = OpenLLMEmbeddings(
+                    **credential_kwargs
+                )
+
+                embedding.embed_query("ping")
         except Exception as ex:
             raise CredentialsValidateFailedError(str(ex))