jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: lazy import sagemaker (#10342)

Browse source
This commit is contained in:
Bowen Liang 2024-11-06 12:45:22 +08:00 • committed by GitHub
commit d45d90e8ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 5 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

6
api/core/model_runtime/model_providers/sagemaker/llm/llm.py
View file

@ -4,10 +4,7 @@ import re
from collections.abc import Generator, Iterator
from typing import Any, Optional, Union, cast
# from openai.types.chat import ChatCompletion, ChatCompletionChunk
import boto3
from sagemaker import Predictor, serializers
from sagemaker.session import Session
from core.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
@ -212,6 +209,9 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
:param user: unique user id
:return: full response or stream response chunk generator result
"""
from sagemaker import Predictor, serializers
from sagemaker.session import Session
if not self.sagemaker_session:
access_key = credentials.get("aws_access_key_id")
secret_key = credentials.get("aws_secret_access_key")