diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e04122dc47..b874d7d80d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1672,11 +1672,16 @@ async def completion( "stream" in data and data["stream"] == True ): # use generate_responses to stream responses custom_headers = {"x-litellm-model-id": model_id} - return StreamingResponse( - async_data_generator( - user_api_key_dict=user_api_key_dict, + stream_content = async_data_generator( + user_api_key_dict=user_api_key_dict, + response=response, + ) + if response.custom_llm_provider == "sagemaker": + stream_content = data_generator( response=response, - ), + ) + return StreamingResponse( + stream_content, media_type="text/event-stream", headers=custom_headers, ) @@ -1834,11 +1839,16 @@ 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} - return StreamingResponse( - async_data_generator( - user_api_key_dict=user_api_key_dict, + stream_content = async_data_generator( + user_api_key_dict=user_api_key_dict, + response=response, + ) + if response.custom_llm_provider == "sagemaker": + stream_content = data_generator( response=response, - ), + ) + return StreamingResponse( + stream_content, media_type="text/event-stream", headers=custom_headers, ) diff --git a/litellm/proxy/tests/test_openai_js.js b/litellm/proxy/tests/test_openai_js.js index 7e74eeca3f..c0f25cf058 100644 --- a/litellm/proxy/tests/test_openai_js.js +++ b/litellm/proxy/tests/test_openai_js.js @@ -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' });