瀏覽代碼

enh: pipe handling

Timothy J. Baek 10 月之前
父節點
當前提交
4c9fa6cf37
共有 1 個文件被更改,包括 13 次插入5 次删除
  1. 13 5
      backend/apps/webui/main.py

+ 13 - 5
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):