1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161from abc import abstractmethod
from pydantic import ConfigDict
from dify_plugin.entities.model import EmbeddingInputType, ModelPropertyKey, ModelType
from dify_plugin.entities.model.text_embedding import (
MultiModalContent,
MultiModalEmbeddingResult,
TextEmbeddingResult,
)
from dify_plugin.interfaces.model.ai_model import AIModel
class TextEmbeddingModel(AIModel):
"""
Model class for text embedding model.
"""
model_type: ModelType = ModelType.TEXT_EMBEDDING
# pydantic configs
model_config = ConfigDict(protected_namespaces=())
############################################################
# Methods that can be implemented by plugin #
############################################################
@abstractmethod
def _invoke(
self,
model: str,
credentials: dict,
texts: list[str],
user: str | None = None,
input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
) -> TextEmbeddingResult:
"""
Invoke large language model
:param model: model name
:param credentials: model credentials
:param texts: texts to embed
:param user: unique user id
:param input_type: embedding input type
:return: embeddings result
"""
raise NotImplementedError
def _invoke_multimodal(
self,
model: str,
credentials: dict,
documents: list[MultiModalContent],
user: str | None = None,
input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
) -> MultiModalEmbeddingResult:
"""Invoke a multimodal embedding model."""
raise NotImplementedError(
f"{self.__class__.__name__} does not implement `_invoke_multimodal`. "
"Implement this method to support multimodal embeddings."
)
@abstractmethod
def get_num_tokens(self, model: str, credentials: dict, texts: list[str]) -> list[int]:
"""
Get number of tokens for given prompt messages
:param model: model name
:param credentials: model credentials
:param texts: texts to embed
:return:
"""
raise NotImplementedError
############################################################
# For plugin implementation use only #
############################################################
def _get_context_size(self, model: str, credentials: dict) -> int:
"""
Get context size for given embedding model
:param model: model name
:param credentials: model credentials
:return: context size
"""
model_schema = self.get_model_schema(model, credentials)
if model_schema and ModelPropertyKey.CONTEXT_SIZE in model_schema.model_properties:
return model_schema.model_properties[ModelPropertyKey.CONTEXT_SIZE]
return 1000
def _get_max_chunks(self, model: str, credentials: dict) -> int:
"""
Get max chunks for given embedding model
:param model: model name
:param credentials: model credentials
:return: max chunks
"""
model_schema = self.get_model_schema(model, credentials)
if model_schema and ModelPropertyKey.MAX_CHUNKS in model_schema.model_properties:
return model_schema.model_properties[ModelPropertyKey.MAX_CHUNKS]
return 1
############################################################
# For executor use only #
############################################################
def invoke(
self,
model: str,
credentials: dict,
texts: list[str],
user: str | None = None,
input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
) -> TextEmbeddingResult:
"""
Invoke large language model
:param model: model name
:param credentials: model credentials
:param texts: texts to embed
:param user: unique user id
:param input_type: embedding input type
:return: embeddings result
"""
with self.timing_context():
try:
return self._invoke(model, credentials, texts, user, input_type)
except Exception as e:
raise self._transform_invoke_error(e) from e
def invoke_multimodal(
self,
model: str,
credentials: dict,
documents: list[MultiModalContent],
user: str | None = None,
input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
) -> MultiModalEmbeddingResult:
"""Invoke a multimodal embedding model."""
with self.timing_context():
try:
return self._invoke_multimodal(
model,
credentials,
documents,
user,
input_type,
)
except NotImplementedError:
raise
except Exception as e:
raise self._transform_invoke_error(e) from e