chore: lazy import sagemaker (#10342)

This commit is contained in:
Bowen Liang
2024-11-06 12:45:22 +08:00
committed by GitHub
parent 42a9374e71
commit d45d90e8ae
4 changed files with 5 additions and 12 deletions

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")