mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-26 01:33:05 +00:00
feat: aws向量模型
This commit is contained in:
wxg0103
wxg0103
parent
cb3c064463
commit
6618c6baf3
|
|
@ -6,12 +6,10 @@ from common.util.file_util import get_file_content
|
|||
from setting.models_provider.base_model_provider import (
|
||||
IModelProvider, ModelProvideInfo, ModelInfo, ModelTypeConst, ModelInfoManage
|
||||
)
|
||||
from setting.models_provider.impl.aws_bedrock_model_provider.credential.embedding import BedrockEmbeddingCredential
|
||||
from setting.models_provider.impl.aws_bedrock_model_provider.credential.llm import BedrockLLMModelCredential
|
||||
from setting.models_provider.impl.aws_bedrock_model_provider.model.embedding import BedrockEmbeddingModel
|
||||
from setting.models_provider.impl.aws_bedrock_model_provider.model.llm import BedrockModel
|
||||
from setting.models_provider.impl.tencent_model_provider.credential.embedding import TencentEmbeddingCredential
|
||||
from setting.models_provider.impl.tencent_model_provider.credential.llm import TencentLLMModelCredential
|
||||
from setting.models_provider.impl.tencent_model_provider.model.embedding import TencentEmbeddingModel
|
||||
from setting.models_provider.impl.tencent_model_provider.model.llm import TencentModel
|
||||
from smartdoc.conf import PROJECT_DIR
|
||||
|
||||
|
||||
|
|
@ -118,10 +116,21 @@ def _initialize_model_info():
|
|||
BedrockLLMModelCredential,
|
||||
BedrockModel),
|
||||
]
|
||||
embedded_model_info_list = [
|
||||
_create_model_info(
|
||||
'amazon.titan-embed-text-v1',
|
||||
'Titan Embed Text 是 Amazon Titan Embed 系列中最大的嵌入模型，可以处理各种文本嵌入任务，如文本分类、文本相似度计算等。',
|
||||
ModelTypeConst.EMBEDDING,
|
||||
BedrockEmbeddingCredential,
|
||||
BedrockEmbeddingModel
|
||||
),
|
||||
]
|
||||
|
||||
model_info_manage = ModelInfoManage.builder() \
|
||||
.append_model_info_list(model_info_list) \
|
||||
.append_default_model_info(model_info_list[0]) \
|
||||
.append_model_info_list(embedded_model_info_list) \
|
||||
.append_default_model_info(embedded_model_info_list[0]) \
|
||||
.build()
|
||||
|
||||
return model_info_manage
|
||||
|
|
|
|||
|
|
@ -1,64 +1,64 @@
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Dict
|
||||
|
||||
from tencentcloud.common import credential
|
||||
from tencentcloud.common.profile.client_profile import ClientProfile
|
||||
from tencentcloud.common.profile.http_profile import HttpProfile
|
||||
from tencentcloud.hunyuan.v20230901 import hunyuan_client, models
|
||||
|
||||
from common import forms
|
||||
from common.exception.app_exception import AppApiException
|
||||
from common.forms import BaseForm
|
||||
from setting.models_provider.base_model_provider import BaseModelCredential, ValidCode
|
||||
from setting.models_provider.impl.aws_bedrock_model_provider.model.embedding import BedrockEmbeddingModel
|
||||
|
||||
|
||||
class TencentEmbeddingCredential(BaseForm, BaseModelCredential):
|
||||
@classmethod
|
||||
def _validate_model_type(cls, model_type: str, provider) -> bool:
|
||||
model_type_list = provider.get_model_type_list()
|
||||
if not any(mt.get('value') == model_type for mt in model_type_list):
|
||||
raise AppApiException(ValidCode.valid_error.value, f'{model_type} 模型类型不支持')
|
||||
return True
|
||||
class BedrockEmbeddingCredential(BaseForm, BaseModelCredential):
|
||||
|
||||
@classmethod
|
||||
def _validate_credential(cls, model_credential: Dict[str, object]) -> credential.Credential:
|
||||
for key in ['SecretId', 'SecretKey']:
|
||||
if key not in model_credential:
|
||||
raise AppApiException(ValidCode.valid_error.value, f'{key} 字段为必填字段')
|
||||
return credential.Credential(model_credential['SecretId'], model_credential['SecretKey'])
|
||||
@staticmethod
|
||||
def _update_aws_credentials(profile_name, access_key_id, secret_access_key):
|
||||
credentials_path = os.path.join(os.path.expanduser("~"), ".aws", "credentials")
|
||||
os.makedirs(os.path.dirname(credentials_path), exist_ok=True)
|
||||
|
||||
@classmethod
|
||||
def _test_credentials(cls, client, model_name: str):
|
||||
req = models.GetEmbeddingRequest()
|
||||
params = {
|
||||
"Model": model_name,
|
||||
"Input": "测试"
|
||||
}
|
||||
req.from_json_string(json.dumps(params))
|
||||
try:
|
||||
res = client.GetEmbedding(req)
|
||||
print(res.to_json_string())
|
||||
except Exception as e:
|
||||
raise AppApiException(ValidCode.valid_error.value, f'校验失败,请检查参数是否正确: {str(e)}')
|
||||
content = open(credentials_path, 'r').read() if os.path.exists(credentials_path) else ''
|
||||
pattern = rf'\n*\[{profile_name}\]\n*(aws_access_key_id = .*)\n*(aws_secret_access_key = .*)\n*'
|
||||
content = re.sub(pattern, '', content, flags=re.DOTALL)
|
||||
|
||||
if not re.search(rf'\[{profile_name}\]', content):
|
||||
content += f"\n[{profile_name}]\naws_access_key_id = {access_key_id}\naws_secret_access_key = {secret_access_key}\n"
|
||||
|
||||
with open(credentials_path, 'w') as file:
|
||||
file.write(content)
|
||||
|
||||
def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], provider,
|
||||
raise_exception=True) -> bool:
|
||||
try:
|
||||
self._validate_model_type(model_type, provider)
|
||||
cred = self._validate_credential(model_credential)
|
||||
httpProfile = HttpProfile(endpoint="hunyuan.tencentcloudapi.com")
|
||||
clientProfile = ClientProfile(httpProfile=httpProfile)
|
||||
client = hunyuan_client.HunyuanClient(cred, "", clientProfile)
|
||||
self._test_credentials(client, model_name)
|
||||
return True
|
||||
except AppApiException as e:
|
||||
raise_exception=False):
|
||||
model_type_list = provider.get_model_type_list()
|
||||
if not any(mt.get('value') == model_type for mt in model_type_list):
|
||||
if raise_exception:
|
||||
raise e
|
||||
raise AppApiException(ValidCode.valid_error.value, f'{model_type} 模型类型不支持')
|
||||
return False
|
||||
|
||||
def encryption_dict(self, model: Dict[str, object]) -> Dict[str, object]:
|
||||
encrypted_secret_key = super().encryption(model.get('SecretKey', ''))
|
||||
return {**model, 'SecretKey': encrypted_secret_key}
|
||||
required_keys = ['region_name', 'access_key_id', 'secret_access_key']
|
||||
if not all(key in model_credential for key in required_keys):
|
||||
if raise_exception:
|
||||
raise AppApiException(ValidCode.valid_error.value, f'以下字段为必填字段: {", ".join(required_keys)}')
|
||||
return False
|
||||
|
||||
SecretId = forms.PasswordInputField('SecretId', required=True)
|
||||
SecretKey = forms.PasswordInputField('SecretKey', required=True)
|
||||
try:
|
||||
self._update_aws_credentials('aws-profile', model_credential['access_key_id'],
|
||||
model_credential['secret_access_key'])
|
||||
model_credential['credentials_profile_name'] = 'aws-profile'
|
||||
model: BedrockEmbeddingModel = provider.get_model(model_type, model_name, model_credential)
|
||||
aa = model.embed_query('你好')
|
||||
print(aa)
|
||||
except AppApiException:
|
||||
raise
|
||||
except Exception as e:
|
||||
if raise_exception:
|
||||
raise AppApiException(ValidCode.valid_error.value, f'校验失败,请检查参数是否正确: {str(e)}')
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def encryption_dict(self, model: Dict[str, object]):
|
||||
return {**model, 'secret_access_key': super().encryption(model.get('secret_access_key', ''))}
|
||||
|
||||
region_name = forms.TextInputField('Region Name', required=True)
|
||||
access_key_id = forms.TextInputField('Access Key ID', required=True)
|
||||
secret_access_key = forms.PasswordInputField('Secret Access Key', required=True)
|
||||
|
|
|
|||
|
|
@ -1,25 +1,56 @@
|
|||
from langchain_community.embeddings import BedrockEmbeddings
|
||||
|
||||
from setting.models_provider.base_model_provider import MaxKBBaseModel
|
||||
from typing import Dict
|
||||
import requests
|
||||
from typing import Dict, List
|
||||
|
||||
|
||||
class TencentEmbeddingModel(MaxKBBaseModel):
|
||||
def __init__(self, secret_id: str, secret_key: str, api_base: str, model_name: str):
|
||||
self.secret_id = secret_id
|
||||
self.secret_key = secret_key
|
||||
self.api_base = api_base
|
||||
self.model_name = model_name
|
||||
class BedrockEmbeddingModel(MaxKBBaseModel, BedrockEmbeddings):
|
||||
def __init__(self, model_id: str, region_name: str, credentials_profile_name: str,
|
||||
**kwargs):
|
||||
super().__init__(model_id=model_id, region_name=region_name,
|
||||
credentials_profile_name=credentials_profile_name, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def new_instance(model_type: str, model_name: str, model_credential: Dict[str, str], **model_kwargs):
|
||||
return TencentEmbeddingModel(
|
||||
secret_id=model_credential.get('SecretId'),
|
||||
secret_key=model_credential.get('SecretKey'),
|
||||
api_base=model_credential.get('api_base'),
|
||||
model_name=model_name,
|
||||
@classmethod
|
||||
def new_instance(cls, model_type: str, model_name: str, model_credential: Dict[str, str],
|
||||
**model_kwargs) -> 'BedrockModel':
|
||||
return cls(
|
||||
model_id=model_name,
|
||||
region_name=model_credential['region_name'],
|
||||
credentials_profile_name=model_credential['credentials_profile_name'],
|
||||
)
|
||||
|
||||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||
"""Compute doc embeddings using a Bedrock model.
|
||||
|
||||
def _generate_auth_token(self):
|
||||
# Example method to generate an authentication token for the model API
|
||||
return f"{self.secret_id}:{self.secret_key}"
|
||||
Args:
|
||||
texts: The list of texts to embed
|
||||
|
||||
Returns:
|
||||
List of embeddings, one for each text.
|
||||
"""
|
||||
results = []
|
||||
for text in texts:
|
||||
response = self._embedding_func(text)
|
||||
|
||||
if self.normalize:
|
||||
response = self._normalize_vector(response)
|
||||
|
||||
results.append(response)
|
||||
|
||||
return results
|
||||
|
||||
def embed_query(self, text: str) -> List[float]:
|
||||
"""Compute query embeddings using a Bedrock model.
|
||||
|
||||
Args:
|
||||
text: The text to embed.
|
||||
|
||||
Returns:
|
||||
Embeddings for the text.
|
||||
"""
|
||||
embedding = self._embedding_func(text)
|
||||
|
||||
if self.normalize:
|
||||
return self._normalize_vector(embedding)
|
||||
|
||||
return embedding
|
||||
|
|
|
|||
Loading…
Reference in New Issue