typing

Author: vangheem

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -111,7 +111,9 @@ class InvokeResponse(BaseModel):
 
     logging.getLogger("etl_uvicorn.fastapi")
 
-    async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> InvokeResponse:
+    ResponseType = StreamingResponse if inspect.isasyncgenfunction(func) else InvokeResponse
+
+    async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> ResponseType:
         usage: list[UsageData] = []
         request_dict = kwargs if kwargs else {}
         if "usage" in inspect.signature(func).parameters:
@@ -144,7 +146,7 @@ async def _stream_response():
     if input_schema_model.model_fields:
 
         @fastapi_app.post("/invoke", response_model=InvokeResponse)
-        async def run_job(request: input_schema_model) -> InvokeResponse:
+        async def run_job(request: input_schema_model) -> ResponseType:
             log_func_and_body(func=func, body=request.json())
             # Create dictionary from pydantic model while preserving underlying types
             request_dict = {f: getattr(request, f) for f in request.model_fields}
@@ -156,7 +158,7 @@ async def run_job(request: input_schema_model) -> InvokeResponse:
     else:
 
         @fastapi_app.post("/invoke", response_model=InvokeResponse)
-        async def run_job() -> InvokeResponse:
+        async def run_job() -> ResponseType:
             log_func_and_body(func=func)
             return await wrap_fn(
                 func=func,