Skip to content

Commit 6e281f6

Browse files
committed
refactor: support volcanic engine embeddings
1 parent de85895 commit 6e281f6

File tree

3 files changed

+12
-6
lines changed

3 files changed

+12
-6
lines changed

apps/setting/models_provider/impl/volcanic_engine_model_provider/credential/embedding.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from setting.models_provider.base_model_provider import BaseModelCredential, ValidCode
1515

1616

17-
class OpenAIEmbeddingCredential(BaseForm, BaseModelCredential):
17+
class VolcanicEmbeddingCredential(BaseForm, BaseModelCredential):
1818
def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], model_params, provider,
1919
raise_exception=True):
2020
model_type_list = provider.get_model_type_list()
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
from typing import Dict
22

3-
from langchain_community.embeddings import VolcanoEmbeddings
3+
from langchain_openai import OpenAIEmbeddings
44

55
from setting.models_provider.base_model_provider import MaxKBBaseModel
66

77

8-
class VolcanicEngineEmbeddingModel(MaxKBBaseModel, VolcanoEmbeddings):
8+
class VolcanicEngineEmbeddingModel(MaxKBBaseModel, OpenAIEmbeddings):
99
@staticmethod
1010
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
1111
return VolcanicEngineEmbeddingModel(
12-
api_key=model_credential.get('api_key'),
12+
openai_api_key=model_credential.get('api_key'),
1313
model=model_name,
1414
openai_api_base=model_credential.get('api_base'),
15+
check_embedding_ctx_length=False,
1516
)

apps/setting/models_provider/impl/volcanic_engine_model_provider/volcanic_engine_model_provider.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,12 @@
1414
from setting.models_provider.impl.openai_model_provider.credential.embedding import OpenAIEmbeddingCredential
1515
from setting.models_provider.impl.openai_model_provider.credential.llm import OpenAILLMModelCredential
1616
from setting.models_provider.impl.openai_model_provider.model.embedding import OpenAIEmbeddingModel
17+
from setting.models_provider.impl.volcanic_engine_model_provider.credential.embedding import VolcanicEmbeddingCredential
1718
from setting.models_provider.impl.volcanic_engine_model_provider.credential.image import \
1819
VolcanicEngineImageModelCredential
1920
from setting.models_provider.impl.volcanic_engine_model_provider.credential.tti import VolcanicEngineTTIModelCredential
2021
from setting.models_provider.impl.volcanic_engine_model_provider.credential.tts import VolcanicEngineTTSModelCredential
22+
from setting.models_provider.impl.volcanic_engine_model_provider.model.embedding import VolcanicEngineEmbeddingModel
2123
from setting.models_provider.impl.volcanic_engine_model_provider.model.image import VolcanicEngineImage
2224
from setting.models_provider.impl.volcanic_engine_model_provider.model.llm import VolcanicEngineChatModel
2325
from setting.models_provider.impl.volcanic_engine_model_provider.credential.stt import VolcanicEngineSTTModelCredential
@@ -81,12 +83,13 @@
8183
),
8284
]
8385

84-
open_ai_embedding_credential = OpenAIEmbeddingCredential()
86+
open_ai_embedding_credential = VolcanicEmbeddingCredential()
8587
model_info_embedding_list = [
8688
ModelInfo('ep-xxxxxxxxxx-yyyy',
8789
'用户前往火山方舟的模型推理页面创建推理接入点,这里需要输入ep-xxxxxxxxxx-yyyy进行调用',
8890
ModelTypeConst.EMBEDDING, open_ai_embedding_credential,
89-
OpenAIEmbeddingModel)]
91+
VolcanicEngineEmbeddingModel)
92+
]
9093

9194
model_info_manage = (
9295
ModelInfoManage.builder()
@@ -96,6 +99,8 @@
9699
.append_default_model_info(model_info_list[2])
97100
.append_default_model_info(model_info_list[3])
98101
.append_default_model_info(model_info_list[4])
102+
.append_model_info_list(model_info_embedding_list)
103+
.append_default_model_info(model_info_embedding_list[0])
99104
.build()
100105
)
101106

0 commit comments

Comments
 (0)