Support streaming invoke responses

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.6
+
+* **Support streaming response types for /invoke**
+
 ## 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
@@ -46,12 +47,29 @@ def log_func_and_body(func: Callable, body: Optional[str] = None) -> None:
         logger.log(level=logger.level, msg=msg)
 
 
+async def run_generator_in_executor(generator_func, **kwargs):
+    loop = asyncio.get_event_loop()
+    # Create a future that will yield the generator values one by one
+    gen = generator_func(**kwargs)
+    while True:
+        result = await loop.run_in_executor(None, next, gen, None)
+        if result is None:
+            break
+        yield result
+
+
 async def invoke_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> Any:
     kwargs = kwargs or {}
-    if inspect.iscoroutinefunction(func):
-        return await func(**kwargs)
+    if inspect.isasyncgenfunction(func):
+        async for val in func(**kwargs):
+            yield val
+    elif inspect.isgeneratorfunction(func):
+        async for val in run_generator_in_executor(func, **kwargs):
+            yield val
+    elif inspect.iscoroutinefunction(func):
+        yield await func(**kwargs)
     else:
-        return await asyncio.get_event_loop().run_in_executor(None, partial(func, **kwargs))
+        yield await asyncio.get_event_loop().run_in_executor(None, partial(func, **kwargs))
 
 
 def check_precheck_func(precheck_func: Callable):
@@ -118,8 +136,14 @@ 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)
+
+            async def _stream_response():
+                async for output in invoke_func(func=func, kwargs=request_dict):
+                    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")
         except Exception as invoke_error:
             logger.error(f"failed to invoke plugin: {invoke_error}", exc_info=True)
             return InvokeResponse(