Преглед изворни кода

chore: add input_type param desc for the _invoke method of text_embedding (#8778)

zhuhao пре 6 месеци
родитељ
комит
ac73763726
26 измењених фајлова са 64 додато и 1 уклоњено
  1. 10 0
      api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py
  2. 1 0
      api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py
  3. 1 0
      api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py
  4. 1 0
      api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py
  5. 10 0
      api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py
  6. 1 0
      api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py
  7. 1 0
      api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py
  8. 1 0
      api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py
  9. 2 1
      api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py
  10. 1 0
      api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py
  11. 1 0
      api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py
  12. 1 0
      api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py
  13. 1 0
      api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py
  14. 1 0
      api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py
  15. 1 0
      api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py
  16. 1 0
      api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py
  17. 1 0
      api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py
  18. 10 0
      api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py
  19. 1 0
      api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py
  20. 10 0
      api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py
  21. 1 0
      api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py
  22. 1 0
      api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py
  23. 2 0
      api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py
  24. 1 0
      api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py
  25. 1 0
      api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py
  26. 1 0
      api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py

+ 10 - 0
api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py

@@ -25,6 +25,16 @@ class AzureOpenAITextEmbeddingModel(_CommonAzureOpenAI, TextEmbeddingModel):
         user: Optional[str] = None,
         input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
     ) -> TextEmbeddingResult:
+        """
+        Invoke text embedding model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
+        :return: embeddings result
+        """
         base_model_name = credentials["base_model_name"]
         credentials_kwargs = self._to_credential_kwargs(credentials)
         client = AzureOpenAI(**credentials_kwargs)

+ 1 - 0
api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py

@@ -50,6 +50,7 @@ class BaichuanTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["api_key"]

+ 1 - 0
api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py

@@ -45,6 +45,7 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         client_config = Config(region_name=credentials["aws_region"])

+ 1 - 0
api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py

@@ -40,6 +40,7 @@ class CohereTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         # get model properties

+ 10 - 0
api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py

@@ -26,6 +26,16 @@ class HuggingfaceHubTextEmbeddingModel(_CommonHuggingfaceHub, TextEmbeddingModel
         user: Optional[str] = None,
         input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
     ) -> TextEmbeddingResult:
+        """
+        Invoke text embedding model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
+        :return: embeddings result
+        """
         client = InferenceClient(token=credentials["huggingfacehub_api_token"])
 
         execute_model = model

+ 1 - 0
api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py

@@ -44,6 +44,7 @@ class HuggingfaceTeiTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         server_url = credentials["server_url"]

+ 1 - 0
api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py

@@ -41,6 +41,7 @@ class HunyuanTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
 

+ 1 - 0
api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py

@@ -75,6 +75,7 @@ class JinaTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["api_key"]

+ 2 - 1
api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py

@@ -23,7 +23,7 @@ from core.model_runtime.model_providers.__base.text_embedding_model import TextE
 
 class LocalAITextEmbeddingModel(TextEmbeddingModel):
     """
-    Model class for Jina text embedding model.
+    Model class for LocalAI text embedding model.
     """
 
     def _invoke(
@@ -41,6 +41,7 @@ class LocalAITextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         if len(texts) != 1:

+ 1 - 0
api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py

@@ -49,6 +49,7 @@ class MinimaxTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["minimax_api_key"]

+ 1 - 0
api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py

@@ -42,6 +42,7 @@ class MixedBreadTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["api_key"]

+ 1 - 0
api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py

@@ -56,6 +56,7 @@ class NomicTextEmbeddingModel(_CommonNomic, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         embeddings, prompt_tokens, total_tokens = self.embed_text(

+ 1 - 0
api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py

@@ -42,6 +42,7 @@ class NvidiaTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         api_key = credentials["api_key"]

+ 1 - 0
api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py

@@ -53,6 +53,7 @@ class OllamaEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
 

+ 1 - 0
api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py

@@ -34,6 +34,7 @@ class OpenAITextEmbeddingModel(_CommonOpenAI, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         # transform credentials to kwargs for model instance

+ 1 - 0
api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py

@@ -43,6 +43,7 @@ class OAICompatEmbeddingModel(_CommonOaiApiCompat, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
 

+ 1 - 0
api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py

@@ -40,6 +40,7 @@ class OpenLLMTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         server_url = credentials["server_url"]

+ 10 - 0
api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py

@@ -22,6 +22,16 @@ class ReplicateEmbeddingModel(_CommonReplicate, TextEmbeddingModel):
         user: Optional[str] = None,
         input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
     ) -> TextEmbeddingResult:
+        """
+        Invoke text embedding model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
+        :return: embeddings result
+        """
         client = ReplicateClient(api_token=credentials["replicate_api_token"], timeout=30)
 
         if "model_version" in credentials:

+ 1 - 0
api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py

@@ -68,6 +68,7 @@ class SageMakerEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         # get model properties

+ 10 - 0
api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py

@@ -24,6 +24,16 @@ class SiliconflowTextEmbeddingModel(OAICompatEmbeddingModel):
         user: Optional[str] = None,
         input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
     ) -> TextEmbeddingResult:
+        """
+        Invoke text embedding model
+
+        :param model: model name
+        :param credentials: model credentials
+        :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
+        :return: embeddings result
+        """
         self._add_custom_parameters(credentials)
         return super()._invoke(model, credentials, texts, user)
 

+ 1 - 0
api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py

@@ -37,6 +37,7 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         credentials_kwargs = self._to_credential_kwargs(credentials)

+ 1 - 0
api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py

@@ -38,6 +38,7 @@ class UpstageTextEmbeddingModel(_CommonUpstage, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
 

+ 2 - 0
api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py

@@ -44,6 +44,8 @@ class VertexAiTextEmbeddingModel(_CommonVertexAi, TextEmbeddingModel):
         :param model: model name
         :param credentials: model credentials
         :param texts: texts to embed
+        :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"]))

+ 1 - 0
api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py

@@ -56,6 +56,7 @@ class VolcengineMaaSTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         if ArkClientV3.is_legacy(credentials):

+ 1 - 0
api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py

@@ -85,6 +85,7 @@ class WenxinTextEmbeddingModel(TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
 

+ 1 - 0
api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py

@@ -30,6 +30,7 @@ class ZhipuAITextEmbeddingModel(_CommonZhipuaiAI, TextEmbeddingModel):
         :param credentials: model credentials
         :param texts: texts to embed
         :param user: unique user id
+        :param input_type: input type
         :return: embeddings result
         """
         credentials_kwargs = self._to_credential_kwargs(credentials)