phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test: mock sagemaker tests

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-21 16:18:02 -07:00
parent 58f46d847c
commit 48e6a7036b
3 changed files with 56 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/llms/sagemaker/chat/handler.py
View file

@ -5,6 +5,7 @@ from typing import Callable, Optional, Union
import httpx
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.utils import ModelResponse, get_secret
from ..common_utils import AWSEventStreamDecoder
@ -125,6 +126,7 @@ class SagemakerChatHandler(BaseAWSLLM):
logger_fn=None,
acompletion: bool = False,
headers: dict = {},
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
):
# pop streaming if it's in the optional params as 'stream' raises an error with sagemaker
@ -173,4 +175,5 @@ class SagemakerChatHandler(BaseAWSLLM):
custom_endpoint=True,
custom_llm_provider="sagemaker_chat",
streaming_decoder=custom_stream_decoder, # type: ignore
client=client,
)