fix(huggingface_restapi.py): return streamed response correctly

BerriAI · Feb 16, 2024 · 5f9e141 · 5f9e141
1 parent 86886de
commit 5f9e141
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py
@@ -631,11 +631,8 @@ async def async_streaming(
                     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
@@ -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():