From 4c9fa6cf371f8768e8316817c2a60f77d2798e6d Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Mon, 24 Jun 2024 12:56:41 -0700 Subject: [PATCH] enh: pipe handling --- backend/apps/webui/main.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py index 260d305f0..f50d1dc36 100644 --- a/backend/apps/webui/main.py +++ b/backend/apps/webui/main.py @@ -233,6 +233,16 @@ async def generate_function_chat_completion(form_data, user): res = await pipe(**params) else: res = pipe(**params) + + # Directly return if the response is a StreamingResponse + if isinstance(res, StreamingResponse): + async for data in res.body_iterator: + yield data + return + if isinstance(res, dict): + yield f"data: {json.dumps(res)}\n\n" + return + except Exception as e: print(f"Error: {e}") yield f"data: {json.dumps({'error': {'detail':str(e)}})}\n\n" @@ -285,15 +295,13 @@ async def generate_function_chat_completion(form_data, user): res = await pipe(**params) else: res = pipe(**params) + + if isinstance(res, StreamingResponse): + return res except Exception as e: print(f"Error: {e}") return {"error": {"detail": str(e)}} - if inspect.iscoroutinefunction(pipe): - res = await pipe(**params) - else: - res = pipe(**params) - if isinstance(res, dict): return res elif isinstance(res, BaseModel):