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

v0 add sagemaker

Browse source
This commit is contained in:
ishaan-jaff 2023-09-04 11:02:19 -07:00
parent 33cc1e2cab
commit 138c26d98d
2 changed files with 130 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/main.py
View file

@ -22,6 +22,7 @@ from litellm.utils import (
from .llms import anthropic
from .llms import together_ai
from .llms import ai21
from .llms import sagemaker
from .llms.huggingface_restapi import HuggingfaceRestAPILLM
from .llms.baseten import BasetenLLM
from .llms.aleph_alpha import AlephAlphaLLM
@ -680,6 +681,32 @@ def completion(
## RESPONSE OBJECT
response = model_response
elif custom_llm_provider == "sagemaker":
# boto3 reads keys from .env
model_response = sagemaker.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=optional_params,
litellm_params=litellm_params,
logger_fn=logger_fn,
encoding=encoding,
api_key=ai21_key,
logging_obj=logging
)
# TODO: Add streaming for sagemaker
# if "stream" in optional_params and optional_params["stream"] == True:
# # don't try to access stream object,
# response = CustomStreamWrapper(
# model_response, model, custom_llm_provider="ai21", logging_obj=logging
# )
# return response
## RESPONSE OBJECT
response = model_response
elif custom_llm_provider == "ollama":
endpoint = (
litellm.api_base if litellm.api_base is not None else api_base