phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(v0) sagemaker streaming

Browse source
This commit is contained in:
ishaan-jaff 2024-01-22 21:50:40 -08:00
parent 802fd6cebf
commit f29de0024a
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

@ -1519,10 +1519,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,
) )