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

(fix) proxy - streaming sagemaker

Browse source
This commit is contained in:
ishaan-jaff 2024-01-23 11:12:16 -08:00
parent 01a2514b98
commit bd37a9cb5e
2 changed files with 19 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

26
litellm/proxy/proxy_server.py
View file

@ -1658,11 +1658,16 @@ async def completion(
"stream" in data and data["stream"] == True "stream" in data and data["stream"] == True
): # use generate_responses to stream responses ): # use generate_responses to stream responses
custom_headers = {"x-litellm-model-id": model_id} custom_headers = {"x-litellm-model-id": model_id}
return StreamingResponse( stream_content = async_data_generator(
async_data_generator( user_api_key_dict=user_api_key_dict,
user_api_key_dict=user_api_key_dict, response=response,
)
if response.custom_llm_provider == "sagemaker":
stream_content = data_generator(
response=response, response=response,
), )
return StreamingResponse(
stream_content,
media_type="text/event-stream", media_type="text/event-stream",
headers=custom_headers, headers=custom_headers,
) )
@ -1820,11 +1825,16 @@ async def chat_completion(
"stream" in data and data["stream"] == True "stream" in data and data["stream"] == True
): # use generate_responses to stream responses ): # use generate_responses to stream responses
custom_headers = {"x-litellm-model-id": model_id} custom_headers = {"x-litellm-model-id": model_id}
return StreamingResponse( stream_content = async_data_generator(
async_data_generator( user_api_key_dict=user_api_key_dict,
user_api_key_dict=user_api_key_dict, response=response,
)
if response.custom_llm_provider == "sagemaker":
stream_content = data_generator(
response=response, response=response,
), )
return StreamingResponse(
stream_content,
media_type="text/event-stream", media_type="text/event-stream",
headers=custom_headers, 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; process.env.DEBUG=false;
async function runOpenAI() { async function runOpenAI() {
const client = new openai.OpenAI({ const client = new openai.OpenAI({
apiKey: 'sk-yPX56TDqBpr23W7ruFG3Yg', apiKey: 'sk-JkKeNi6WpWDngBsghJ6B9g',
baseURL: 'http://0.0.0.0:8000' baseURL: 'http://0.0.0.0:8000'
}); });