From 0bf35b4a91c7ee000b9979b8bf2cca69a08d37bc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 15 Feb 2024 20:55:21 -0800 Subject: [PATCH] fix(huggingface_restapi.py): catch streaming errors --- litellm/llms/huggingface_restapi.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index e945242bf2..e66627cccf 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -630,8 +630,12 @@ class Huggingface(BaseLLM): custom_llm_provider="huggingface", logging_obj=logging_obj, ) - async for transformed_chunk in streamwrapper: - yield transformed_chunk + + async def generator(): + async for transformed_chunk in streamwrapper: + yield transformed_chunk + + return generator() def embedding( self,