scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge pull request #1569 from BerriAI/litellm_proxy_sagemaker_streaming

Browse source 
[Fix] LiteLLM Proxy - Streaming Sagemaker Support on /chat/completions
This commit is contained in:
Ishaan Jaff 2024-01-23 12:35:16 -08:00 • committed by GitHub
commit 861587e49d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 36 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/proxy/proxy_server.py
View file

@ -1422,6 +1422,28 @@ async def async_data_generator(response, user_api_key_dict):
yield f"data: {str(e)}\n\n"
def select_data_generator(response, user_api_key_dict):
try:
# since boto3 - sagemaker does not support async calls, we should use a sync data_generator
if (
hasattr(response, "custom_llm_provider")
and response.custom_llm_provider == "sagemaker"
):
return data_generator(
response=response,
)
else:
# default to async_data_generator
return async_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
except:
# worst case - use async_data_generator
return async_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
def get_litellm_model_info(model: dict = {}):
model_info = model.get("model_info", {})
model_to_lookup = model.get("litellm_params", {}).get("model", None)
@ -1658,11 +1680,12 @@ async def completion(
"stream" in data and data["stream"] == True
): # use generate_responses to stream responses
custom_headers = {"x-litellm-model-id": model_id}
selected_data_generator = select_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
return StreamingResponse(
async_data_generator(
user_api_key_dict=user_api_key_dict,
response=response,
),
selected_data_generator,
media_type="text/event-stream",
headers=custom_headers,
)
@ -1820,11 +1843,12 @@ async def chat_completion(
"stream" in data and data["stream"] == True
): # use generate_responses to stream responses
custom_headers = {"x-litellm-model-id": model_id}
selected_data_generator = select_data_generator(
response=response, user_api_key_dict=user_api_key_dict
)
return StreamingResponse(
async_data_generator(
user_api_key_dict=user_api_key_dict,
response=response,
),
selected_data_generator,
media_type="text/event-stream",
headers=custom_headers,
)

2
litellm/proxy/tests/test_openai_js.js
View file

@ -4,7 +4,7 @@ const openai = require('openai');
process.env.DEBUG=false;
async function runOpenAI() {
const client = new openai.OpenAI({
apiKey: 'sk-yPX56TDqBpr23W7ruFG3Yg',
apiKey: 'sk-JkKeNi6WpWDngBsghJ6B9g',
baseURL: 'http://0.0.0.0:8000'
});

8
litellm/utils.py
View file

@ -7732,10 +7732,8 @@ class CustomStreamWrapper:
]
self.sent_last_chunk = True
elif self.custom_llm_provider == "sagemaker":
print_verbose(f"ENTERS SAGEMAKER STREAMING")
new_chunk = next(self.completion_stream)
completion_obj["content"] = new_chunk
print_verbose(f"ENTERS SAGEMAKER STREAMING for chunk {chunk}")
completion_obj["content"] = chunk
elif self.custom_llm_provider == "petals":
if len(self.completion_stream) == 0:
if self.sent_last_chunk:
@ -7854,7 +7852,7 @@ class CustomStreamWrapper:
completion_obj["role"] = "assistant"
self.sent_first_chunk = True
model_response.choices[0].delta = Delta(**completion_obj)
print_verbose(f"model_response: {model_response}")
print_verbose(f"returning model_response: {model_response}")
return model_response
else:
return