From 9c2cce56f7d57d41796748a4c57a225a102df7b8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 16 Feb 2024 13:25:13 -0800 Subject: [PATCH] fix(huggingface_restapi.py): return streamed response correctly --- litellm/llms/huggingface_restapi.py | 7 ++----- litellm/tests/test_streaming.py | 2 +- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index d898ed8c7f..123b8ecbc4 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -631,11 +631,8 @@ class Huggingface(BaseLLM): logging_obj=logging_obj, ) - async def generator(): - async for transformed_chunk in streamwrapper: - yield transformed_chunk - - return generator() + async for transformed_chunk in streamwrapper: + yield transformed_chunk def embedding( self, diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 30d777d799..4df4f20e92 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -899,7 +899,7 @@ async def test_sagemaker_streaming_async(): pytest.fail(f"An exception occurred - {str(e)}") -asyncio.run(test_sagemaker_streaming_async()) +# asyncio.run(test_sagemaker_streaming_async()) def test_completion_sagemaker_stream():