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

(v0) sagemaker streaming

Browse source
This commit is contained in:
ishaan-jaff 2024-01-22 21:50:40 -08:00
parent 998094d38d
commit 3ee672c35f
1 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/main.py
View file

@ -1520,10 +1520,12 @@ def completion(
# fake streaming for sagemaker # fake streaming for sagemaker
print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER") print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER")
resp_string = model_response["choices"][0]["message"]["content"] from .llms.sagemaker import TokenIterator
tokenIterator = TokenIterator(model_response)
response = CustomStreamWrapper( response = CustomStreamWrapper(
resp_string, completion_stream=tokenIterator,
model, model=model,
custom_llm_provider="sagemaker", custom_llm_provider="sagemaker",
logging_obj=logging, logging_obj=logging,
) )