Fix some security vulnerabilities. (#2160)
Browse files### What problem does this PR solve?
Fix some security vulnerabilities
### Type of change
- [x] Performance Improvement
---------
Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
- api/apps/llm_app.py +1 -1
- rag/llm/chat_model.py +5 -5
- rag/llm/embedding_model.py +3 -3
api/apps/llm_app.py
CHANGED
|
@@ -22,7 +22,7 @@ from api.db.db_models import TenantLLM
|
|
| 22 |
from api.utils.api_utils import get_json_result
|
| 23 |
from rag.llm import EmbeddingModel, ChatModel, RerankModel, CvModel, TTSModel
|
| 24 |
import requests
|
| 25 |
-
|
| 26 |
|
| 27 |
@manager.route('/factories', methods=['GET'])
|
| 28 |
@login_required
|
|
|
|
| 22 |
from api.utils.api_utils import get_json_result
|
| 23 |
from rag.llm import EmbeddingModel, ChatModel, RerankModel, CvModel, TTSModel
|
| 24 |
import requests
|
| 25 |
+
|
| 26 |
|
| 27 |
@manager.route('/factories', methods=['GET'])
|
| 28 |
@login_required
|
rag/llm/chat_model.py
CHANGED
|
@@ -457,8 +457,8 @@ class VolcEngineChat(Base):
|
|
| 457 |
model_name is for display only
|
| 458 |
"""
|
| 459 |
base_url = base_url if base_url else 'https://ark.cn-beijing.volces.com/api/v3'
|
| 460 |
-
ark_api_key =
|
| 461 |
-
model_name =
|
| 462 |
super().__init__(ark_api_key, model_name, base_url)
|
| 463 |
|
| 464 |
|
|
@@ -602,9 +602,9 @@ class BedrockChat(Base):
|
|
| 602 |
|
| 603 |
def __init__(self, key, model_name, **kwargs):
|
| 604 |
import boto3
|
| 605 |
-
self.bedrock_ak =
|
| 606 |
-
self.bedrock_sk =
|
| 607 |
-
self.bedrock_region =
|
| 608 |
self.model_name = model_name
|
| 609 |
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
|
| 610 |
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
|
|
|
|
| 457 |
model_name is for display only
|
| 458 |
"""
|
| 459 |
base_url = base_url if base_url else 'https://ark.cn-beijing.volces.com/api/v3'
|
| 460 |
+
ark_api_key = json.loads(key).get('ark_api_key', '')
|
| 461 |
+
model_name = json.loads(key).get('ep_id', '')
|
| 462 |
super().__init__(ark_api_key, model_name, base_url)
|
| 463 |
|
| 464 |
|
|
|
|
| 602 |
|
| 603 |
def __init__(self, key, model_name, **kwargs):
|
| 604 |
import boto3
|
| 605 |
+
self.bedrock_ak = json.loads(key).get('bedrock_ak', '')
|
| 606 |
+
self.bedrock_sk = json.loads(key).get('bedrock_sk', '')
|
| 607 |
+
self.bedrock_region = json.loads(key).get('bedrock_region', '')
|
| 608 |
self.model_name = model_name
|
| 609 |
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
|
| 610 |
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
|
rag/llm/embedding_model.py
CHANGED
|
@@ -403,9 +403,9 @@ class BedrockEmbed(Base):
|
|
| 403 |
def __init__(self, key, model_name,
|
| 404 |
**kwargs):
|
| 405 |
import boto3
|
| 406 |
-
self.bedrock_ak =
|
| 407 |
-
self.bedrock_sk =
|
| 408 |
-
self.bedrock_region =
|
| 409 |
self.model_name = model_name
|
| 410 |
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
|
| 411 |
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
|
|
|
|
| 403 |
def __init__(self, key, model_name,
|
| 404 |
**kwargs):
|
| 405 |
import boto3
|
| 406 |
+
self.bedrock_ak = json.loads(key).get('bedrock_ak', '')
|
| 407 |
+
self.bedrock_sk = json.loads(key).get('bedrock_sk', '')
|
| 408 |
+
self.bedrock_region = json.loads(key).get('bedrock_region', '')
|
| 409 |
self.model_name = model_name
|
| 410 |
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
|
| 411 |
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
|