Support streaming invoke responses

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.6
+
+* **Support streaming response types for /invoke if callable is async generator**
+
 ## 0.0.5
 
 * **Improve logging to hide body in case of sensitive data unless TRACE level**

unstructured_platform_plugins/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.5"  # pragma: no cover
+__version__ = "0.0.6"  # pragma: no cover

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -7,6 +7,7 @@
 from typing import Any, Callable, Optional
 
 from fastapi import FastAPI, status
+from fastapi.responses import StreamingResponse
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
@@ -118,8 +119,19 @@ async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> In
         else:
             logger.warning("usage data not an expected parameter, omitting")
         try:
-            output = await invoke_func(func=func, kwargs=request_dict)
-            return InvokeResponse(usage=usage, status_code=status.HTTP_200_OK, output=output)
+            if inspect.isasyncgenfunction(func):
+                # Stream response if function is an async generator
+
+                async def _stream_response():
+                    async for output in func(**(request_dict or {})):
+                        yield InvokeResponse(
+                            usage=usage, status_code=status.HTTP_200_OK, output=output
+                        ).model_dump_json() + "\n"
+
+                return StreamingResponse(_stream_response(), media_type="application/x-ndjson")
+            else:
+                output = await invoke_func(func=func, kwargs=request_dict)
+                return InvokeResponse(usage=usage, status_code=status.HTTP_200_OK, output=output)
         except Exception as invoke_error:
             logger.error(f"failed to invoke plugin: {invoke_error}", exc_info=True)
             return InvokeResponse(