AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-13 00:29:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix llm integration problem: It doesn't work on docker env (#8701)

Browse Source 
Co-authored-by: Yuanbo Li <ybalbert@amazon.com>
This commit is contained in:
ybalbert001 2024-09-24 10:33:30 +08:00 committed by GitHub
parent 21e9608b23
commit 7c485f8bb8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
api/core/model_runtime/model_providers/sagemaker/llm/llm.py
View File

@ -85,7 +85,6 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
"""
sagemaker_client: Any = None
sagemaker_sess: Any = None
predictor: Any = None
def _handle_chat_generate_response(
@ -213,23 +212,22 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
:return: full response or stream response chunk generator result
"""
if not self.sagemaker_client:
access_key = credentials.get("access_key")
secret_key = credentials.get("secret_key")
access_key = credentials.get("aws_access_key_id")
secret_key = credentials.get("aws_secret_access_key")
aws_region = credentials.get("aws_region")
boto_session = None
if aws_region:
if access_key and secret_key:
self.sagemaker_client = boto3.client(
"sagemaker-runtime",
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
region_name=aws_region,
boto_session = boto3.Session(
aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name=aws_region
)
else:
self.sagemaker_client = boto3.client("sagemaker-runtime", region_name=aws_region)
boto_session = boto3.Session(region_name=aws_region)
else:
self.sagemaker_client = boto3.client("sagemaker-runtime")
boto_session = boto3.Session()
sagemaker_session = Session(sagemaker_runtime_client=self.sagemaker_client)
self.sagemaker_client = boto_session.client("sagemaker")
sagemaker_session = Session(boto_session=boto_session, sagemaker_client=self.sagemaker_client)
self.predictor = Predictor(
endpoint_name=credentials.get("sagemaker_endpoint"),
sagemaker_session=sagemaker_session,