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 <[email protected]>
- 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)
|