From 12975cf128c07091b2031eaf4ecad01b1fa6731f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=BB=84=E8=85=BE?= <101850389+hangters@users.noreply.github.com> Date: Thu, 29 Aug 2024 16:21:32 +0800 Subject: [PATCH] Fix some security vulnerabilities. (#2160) ### What problem does this PR solve? Fix some security vulnerabilities ### Type of change - [x] Performance Improvement --------- Co-authored-by: Zhedong Cen --- api/apps/llm_app.py | 2 +- rag/llm/chat_model.py | 10 +++++----- rag/llm/embedding_model.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/api/apps/llm_app.py b/api/apps/llm_app.py index e48b7b3a9..26cbd4576 100644 --- a/api/apps/llm_app.py +++ b/api/apps/llm_app.py @@ -22,7 +22,7 @@ from api.db.db_models import TenantLLM from api.utils.api_utils import get_json_result from rag.llm import EmbeddingModel, ChatModel, RerankModel, CvModel, TTSModel import requests -import ast + @manager.route('/factories', methods=['GET']) @login_required diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 3af0e0257..9569daa24 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -457,8 +457,8 @@ class VolcEngineChat(Base): model_name is for display only """ base_url = base_url if base_url else 'https://ark.cn-beijing.volces.com/api/v3' - ark_api_key = eval(key).get('ark_api_key', '') - model_name = eval(key).get('ep_id', '') + ark_api_key = json.loads(key).get('ark_api_key', '') + model_name = json.loads(key).get('ep_id', '') super().__init__(ark_api_key, model_name, base_url) @@ -602,9 +602,9 @@ class BedrockChat(Base): def __init__(self, key, model_name, **kwargs): import boto3 - self.bedrock_ak = eval(key).get('bedrock_ak', '') - self.bedrock_sk = eval(key).get('bedrock_sk', '') - self.bedrock_region = eval(key).get('bedrock_region', '') + self.bedrock_ak = json.loads(key).get('bedrock_ak', '') + self.bedrock_sk = json.loads(key).get('bedrock_sk', '') + self.bedrock_region = json.loads(key).get('bedrock_region', '') self.model_name = model_name self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region, aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk) diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index 257a32891..0c9ce73fc 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -403,9 +403,9 @@ class BedrockEmbed(Base): def __init__(self, key, model_name, **kwargs): import boto3 - self.bedrock_ak = eval(key).get('bedrock_ak', '') - self.bedrock_sk = eval(key).get('bedrock_sk', '') - self.bedrock_region = eval(key).get('bedrock_region', '') + self.bedrock_ak = json.loads(key).get('bedrock_ak', '') + self.bedrock_sk = json.loads(key).get('bedrock_sk', '') + self.bedrock_region = json.loads(key).get('bedrock_region', '') self.model_name = model_name self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region, aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)