From 4669eb24befe840789de1dc9ae32f631a31a765d Mon Sep 17 00:00:00 2001 From: Jyong <76649700+JohnJyong@users.noreply.github.com> Date: Tue, 24 Sep 2024 21:53:50 +0800 Subject: [PATCH] add embedding input type parameter (#8724) --- api/core/embedding/cached_embedding.py | 9 ++++++-- api/core/embedding/embedding_constant.py | 10 ++++++++ api/core/model_manager.py | 7 +++++- .../__base/text_embedding_model.py | 23 +++++++++++++++---- .../text_embedding/text_embedding.py | 8 ++++++- .../baichuan/text_embedding/text_embedding.py | 8 ++++++- .../bedrock/text_embedding/text_embedding.py | 8 ++++++- .../cohere/text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../hunyuan/text_embedding/text_embedding.py | 8 ++++++- .../jina/text_embedding/text_embedding.py | 8 ++++++- .../localai/text_embedding/text_embedding.py | 8 ++++++- .../minimax/text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../nomic/text_embedding/text_embedding.py | 2 ++ .../nvidia/text_embedding/text_embedding.py | 8 ++++++- .../oci/text_embedding/text_embedding.py | 9 +++++++- .../ollama/text_embedding/text_embedding.py | 8 ++++++- .../openai/text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../openllm/text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 9 +++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../tongyi/text_embedding/text_embedding.py | 2 ++ .../upstage/text_embedding/text_embedding.py | 10 +++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 8 ++++++- .../wenxin/text_embedding/text_embedding.py | 8 ++++++- .../text_embedding/text_embedding.py | 9 +++++++- .../zhipuai/text_embedding/text_embedding.py | 8 ++++++- 33 files changed, 239 insertions(+), 35 deletions(-) create mode 100644 api/core/embedding/embedding_constant.py diff --git a/api/core/embedding/cached_embedding.py b/api/core/embedding/cached_embedding.py index 8ce12fd59f576..75219051cd30c 100644 --- a/api/core/embedding/cached_embedding.py +++ b/api/core/embedding/cached_embedding.py @@ -5,6 +5,7 @@ import numpy as np from sqlalchemy.exc import IntegrityError +from core.embedding.embedding_constant import EmbeddingInputType from core.model_manager import ModelInstance from core.model_runtime.entities.model_entities import ModelPropertyKey from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel @@ -56,7 +57,9 @@ def embed_documents(self, texts: list[str]) -> list[list[float]]: for i in range(0, len(embedding_queue_texts), max_chunks): batch_texts = embedding_queue_texts[i : i + max_chunks] - embedding_result = self._model_instance.invoke_text_embedding(texts=batch_texts, user=self._user) + embedding_result = self._model_instance.invoke_text_embedding( + texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT + ) for vector in embedding_result.embeddings: try: @@ -100,7 +103,9 @@ def embed_query(self, text: str) -> list[float]: redis_client.expire(embedding_cache_key, 600) return list(np.frombuffer(base64.b64decode(embedding), dtype="float")) try: - embedding_result = self._model_instance.invoke_text_embedding(texts=[text], user=self._user) + embedding_result = self._model_instance.invoke_text_embedding( + texts=[text], user=self._user, input_type=EmbeddingInputType.QUERY + ) embedding_results = embedding_result.embeddings[0] embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist() diff --git a/api/core/embedding/embedding_constant.py b/api/core/embedding/embedding_constant.py new file mode 100644 index 0000000000000..9b4934646bc0e --- /dev/null +++ b/api/core/embedding/embedding_constant.py @@ -0,0 +1,10 @@ +from enum import Enum + + +class EmbeddingInputType(Enum): + """ + Enum for embedding input type. + """ + + DOCUMENT = "document" + QUERY = "query" diff --git a/api/core/model_manager.py b/api/core/model_manager.py index 990efd36c609c..74b445236268d 100644 --- a/api/core/model_manager.py +++ b/api/core/model_manager.py @@ -3,6 +3,7 @@ from collections.abc import Callable, Generator, Sequence from typing import IO, Optional, Union, cast +from core.embedding.embedding_constant import EmbeddingInputType from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle from core.entities.provider_entities import ModelLoadBalancingConfiguration from core.errors.error import ProviderTokenNotInitError @@ -158,12 +159,15 @@ def get_llm_num_tokens( tools=tools, ) - def invoke_text_embedding(self, texts: list[str], user: Optional[str] = None) -> TextEmbeddingResult: + def invoke_text_embedding( + self, texts: list[str], user: Optional[str] = None, input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT + ) -> TextEmbeddingResult: """ Invoke large language model :param texts: texts to embed :param user: unique user id + :param input_type: input type :return: embeddings result """ if not isinstance(self.model_type_instance, TextEmbeddingModel): @@ -176,6 +180,7 @@ def invoke_text_embedding(self, texts: list[str], user: Optional[str] = None) -> credentials=self.credentials, texts=texts, user=user, + input_type=input_type, ) def get_text_embedding_num_tokens(self, texts: list[str]) -> int: diff --git a/api/core/model_runtime/model_providers/__base/text_embedding_model.py b/api/core/model_runtime/model_providers/__base/text_embedding_model.py index 54a4486023691..a948dca20d69a 100644 --- a/api/core/model_runtime/model_providers/__base/text_embedding_model.py +++ b/api/core/model_runtime/model_providers/__base/text_embedding_model.py @@ -4,6 +4,7 @@ from pydantic import ConfigDict +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import ModelPropertyKey, ModelType from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult from core.model_runtime.model_providers.__base.ai_model import AIModel @@ -20,35 +21,47 @@ class TextEmbeddingModel(AIModel): model_config = ConfigDict(protected_namespaces=()) def invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ - Invoke large language model + 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.started_at = time.perf_counter() try: - return self._invoke(model, credentials, texts, user) + return self._invoke(model, credentials, texts, user, input_type) except Exception as e: raise self._transform_invoke_error(e) @abstractmethod def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ - Invoke large language model + 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 """ raise NotImplementedError diff --git a/api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py index d9cff8ecbbadb..6b270b65ff00a 100644 --- a/api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/azure_openai/text_embedding/text_embedding.py @@ -7,6 +7,7 @@ import tiktoken from openai import AzureOpenAI +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import AIModelEntity, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.validate import CredentialsValidateFailedError @@ -17,7 +18,12 @@ class AzureOpenAITextEmbeddingModel(_CommonAzureOpenAI, TextEmbeddingModel): def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: base_model_name = credentials["base_model_name"] credentials_kwargs = self._to_credential_kwargs(credentials) diff --git a/api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py index 779dfbb608163..210c274bdf5e3 100644 --- a/api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/baichuan/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ from requests import post +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -35,7 +36,12 @@ class BaichuanTextEmbeddingModel(TextEmbeddingModel): api_base: str = "http://api.baichuan-ai.com/v1/embeddings" def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py index 251170d1aec49..8c4c50b26958c 100644 --- a/api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py @@ -13,6 +13,7 @@ UnknownServiceError, ) +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -30,7 +31,12 @@ class BedrockTextEmbeddingModel(TextEmbeddingModel): def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py index a1c5e98118e4c..1f93068a8c9e5 100644 --- a/api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/cohere/text_embedding/text_embedding.py @@ -5,6 +5,7 @@ import numpy as np from cohere.core import RequestOptions +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -25,7 +26,12 @@ class CohereTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py index 4ad96c4233fb1..cf18f84ac8cb2 100644 --- a/api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/huggingface_hub/text_embedding/text_embedding.py @@ -6,6 +6,7 @@ import requests from huggingface_hub import HfApi, InferenceClient +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -18,7 +19,12 @@ class HuggingfaceHubTextEmbeddingModel(_CommonHuggingfaceHub, TextEmbeddingModel): def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: client = InferenceClient(token=credentials["huggingfacehub_api_token"]) diff --git a/api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py index 55f3c25804ea0..58baf4933c241 100644 --- a/api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/huggingface_tei/text_embedding/text_embedding.py @@ -1,6 +1,7 @@ import time from typing import Optional +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -23,7 +24,12 @@ class HuggingfaceTeiTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py index 1396e59e188bc..3e14371f897e0 100644 --- a/api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/hunyuan/text_embedding/text_embedding.py @@ -9,6 +9,7 @@ from tencentcloud.common.profile.http_profile import HttpProfile from tencentcloud.hunyuan.v20230901 import hunyuan_client, models +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -26,7 +27,12 @@ class HunyuanTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py index 6c96699ea2e3f..9120f26b8d92f 100644 --- a/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/jina/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ from requests import post +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -60,7 +61,12 @@ def transform_jina_input_text(model, text): return data def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py index 7d258be81e058..d8878c7be85bc 100644 --- a/api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/localai/text_embedding/text_embedding.py @@ -5,6 +5,7 @@ from requests import post from yarl import URL +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -26,7 +27,12 @@ class LocalAITextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py index 76fd1342bdb92..d0d1d2aea1ccf 100644 --- a/api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/minimax/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ from requests import post +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -34,7 +35,12 @@ class MinimaxTextEmbeddingModel(TextEmbeddingModel): api_base: str = "https://api.minimax.chat/v1/embeddings" def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py index 05d9a9a0c69f4..cdc2d58d0c7e3 100644 --- a/api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/mixedbread/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ import requests +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -27,7 +28,12 @@ class MixedBreadTextEmbeddingModel(TextEmbeddingModel): api_base: str = "https://api.mixedbread.ai/v1" def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py index ccbfd196a9522..a797521576b23 100644 --- a/api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/nomic/text_embedding/text_embedding.py @@ -5,6 +5,7 @@ from nomic import embed from nomic import login as nomic_login +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import ( EmbeddingUsage, @@ -46,6 +47,7 @@ def _invoke( credentials: dict, texts: list[str], user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py index 00cec265d5ed9..a4ea28bd1073b 100644 --- a/api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/nvidia/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ from requests import post +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -27,7 +28,12 @@ class NvidiaTextEmbeddingModel(TextEmbeddingModel): models: list[str] = ["NV-Embed-QA"] def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/oci/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/oci/text_embedding/text_embedding.py index 80ad2be9f5687..4de9296ccaa92 100644 --- a/api/core/model_runtime/model_providers/oci/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/oci/text_embedding/text_embedding.py @@ -6,6 +6,7 @@ import numpy as np import oci +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -41,7 +42,12 @@ class OCITextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model @@ -50,6 +56,7 @@ def _invoke( :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 diff --git a/api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py index b4c61d8a6dc79..0501c8b841666 100644 --- a/api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/ollama/text_embedding/text_embedding.py @@ -8,6 +8,7 @@ import numpy as np import requests +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import ( AIModelEntity, @@ -38,7 +39,12 @@ class OllamaEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py index 535d8388bc7d7..7945723636ced 100644 --- a/api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/openai/text_embedding/text_embedding.py @@ -6,6 +6,7 @@ import tiktoken from openai import OpenAI +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.validate import CredentialsValidateFailedError @@ -19,7 +20,12 @@ class OpenAITextEmbeddingModel(_CommonOpenAI, TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py index e83cfdf873cb4..68b5773e16afe 100644 --- a/api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py @@ -7,6 +7,7 @@ import numpy as np import requests +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import ( AIModelEntity, @@ -28,7 +29,12 @@ class OAICompatEmbeddingModel(_CommonOaiApiCompat, TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py index 00e583cc797dc..c0a3efbb003dc 100644 --- a/api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/openllm/text_embedding/text_embedding.py @@ -5,6 +5,7 @@ from requests import post from requests.exceptions import ConnectionError, InvalidSchema, MissingSchema +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import ( @@ -25,7 +26,12 @@ class OpenLLMTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/perfxcloud/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/perfxcloud/text_embedding/text_embedding.py index b62a2d2aaf06f..1e86f351c8ae5 100644 --- a/api/core/model_runtime/model_providers/perfxcloud/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/perfxcloud/text_embedding/text_embedding.py @@ -7,6 +7,7 @@ import numpy as np import requests +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import ( AIModelEntity, @@ -28,7 +29,12 @@ class OAICompatEmbeddingModel(_CommonOaiApiCompat, TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model @@ -37,6 +43,7 @@ def _invoke( :param credentials: model credentials :param texts: texts to embed :param user: unique user id + :param input_type: input type :return: embeddings result """ diff --git a/api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py index 71b6fb99c4a1b..b6cf89bcd9868 100644 --- a/api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/replicate/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ from replicate import Client as ReplicateClient +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -14,7 +15,12 @@ class ReplicateEmbeddingModel(_CommonReplicate, TextEmbeddingModel): def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: client = ReplicateClient(api_token=credentials["replicate_api_token"], timeout=30) diff --git a/api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py index d55144f8a79be..957f2e5d0f78f 100644 --- a/api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/sagemaker/text_embedding/text_embedding.py @@ -6,6 +6,7 @@ import boto3 +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -53,7 +54,12 @@ def _sagemaker_embedding(self, sm_client, endpoint_name, content_list: list[str] return embeddings def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py index 6cdf4933b47c7..c6c681c15d9d4 100644 --- a/api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/siliconflow/text_embedding/text_embedding.py @@ -1,5 +1,6 @@ from typing import Optional +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult from core.model_runtime.model_providers.openai_api_compatible.text_embedding.text_embedding import ( OAICompatEmbeddingModel, @@ -16,7 +17,12 @@ def validate_credentials(self, model: str, credentials: dict) -> None: super().validate_credentials(model, credentials) def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: self._add_custom_parameters(credentials) return super()._invoke(model, credentials, texts, user) diff --git a/api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py index 5783d2e383e2d..0eef0db3e7df0 100644 --- a/api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/tongyi/text_embedding/text_embedding.py @@ -4,6 +4,7 @@ import dashscope import numpy as np +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import ( EmbeddingUsage, @@ -27,6 +28,7 @@ def _invoke( credentials: dict, texts: list[str], user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py index edd4a36d98f58..812bf92eeace7 100644 --- a/api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/upstage/text_embedding/text_embedding.py @@ -7,6 +7,7 @@ from openai import OpenAI from tokenizers import Tokenizer +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.validate import CredentialsValidateFailedError @@ -22,7 +23,14 @@ class UpstageTextEmbeddingModel(_CommonUpstage, TextEmbeddingModel): def _get_tokenizer(self) -> Tokenizer: return Tokenizer.from_pretrained("upstage/solar-1-mini-tokenizer") - def _invoke(self, model: str, credentials: dict, texts: list[str], user: str | None = None) -> TextEmbeddingResult: + def _invoke( + self, + model: str, + credentials: dict, + texts: list[str], + user: str | None = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, + ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py index 519373a7f31a3..509b41d9517bf 100644 --- a/api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py @@ -9,6 +9,7 @@ from google.oauth2 import service_account from vertexai.language_models import TextEmbeddingModel as VertexTextEmbeddingModel +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import ( AIModelEntity, @@ -30,7 +31,12 @@ class VertexAiTextEmbeddingModel(_CommonVertexAi, TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py index 9cba2cb879418..9d800af5f4d38 100644 --- a/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/volcengine_maas/text_embedding/text_embedding.py @@ -2,6 +2,7 @@ from decimal import Decimal from typing import Optional +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import ( AIModelEntity, @@ -41,7 +42,12 @@ class VolcengineMaaSTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py index 4d6f6dccd0cf7..1b5a0904db963 100644 --- a/api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/wenxin/text_embedding/text_embedding.py @@ -7,6 +7,7 @@ import numpy as np from requests import Response, post +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.invoke import InvokeError @@ -70,7 +71,12 @@ def _create_text_embedding(self, api_key: str, secret_key: str) -> TextEmbedding return WenxinTextEmbedding(api_key, secret_key) def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model diff --git a/api/core/model_runtime/model_providers/xinference/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/xinference/text_embedding/text_embedding.py index 8043af1d6cf11..16272391320d5 100644 --- a/api/core/model_runtime/model_providers/xinference/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/xinference/text_embedding/text_embedding.py @@ -3,6 +3,7 @@ from xinference_client.client.restful.restful_client import Client, RESTfulEmbeddingModelHandle +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.common_entities import I18nObject from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType, PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult @@ -25,7 +26,12 @@ class XinferenceTextEmbeddingModel(TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model @@ -40,6 +46,7 @@ def _invoke( :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"] diff --git a/api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py b/api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py index ee20954381053..707c08ef1b5f6 100644 --- a/api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py +++ b/api/core/model_runtime/model_providers/zhipuai/text_embedding/text_embedding.py @@ -1,6 +1,7 @@ import time from typing import Optional +from core.embedding.embedding_constant import EmbeddingInputType from core.model_runtime.entities.model_entities import PriceType from core.model_runtime.entities.text_embedding_entities import EmbeddingUsage, TextEmbeddingResult from core.model_runtime.errors.validate import CredentialsValidateFailedError @@ -15,7 +16,12 @@ class ZhipuAITextEmbeddingModel(_CommonZhipuaiAI, TextEmbeddingModel): """ def _invoke( - self, model: str, credentials: dict, texts: list[str], user: Optional[str] = None + self, + model: str, + credentials: dict, + texts: list[str], + user: Optional[str] = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, ) -> TextEmbeddingResult: """ Invoke text embedding model