feat: Add precheck endpoint (#15)

Author: rbiseck3

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -34,11 +34,25 @@ async def invoke_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -
         return func(**kwargs)
 
 
+def check_precheck_func(precheck_func: Callable):
+    sig = inspect.signature(precheck_func)
+    inputs = sig.parameters.values()
+    outputs = sig.return_annotation
+    if len(inputs) == 1:
+        i = inputs[0]
+        if i.name != "usage" or i.annotation is list:
+            raise ValueError("the only input available for precheck is usage which must be a list")
+    if outputs not in [None, sig.empty]:
+        raise ValueError(f"no output should exist for precheck function, found: {outputs}")
+
+
 def generate_fast_api(
     app: str,
     method_name: Optional[str] = None,
     id_str: Optional[str] = None,
     id_method: Optional[str] = None,
+    precheck_str: Optional[str] = None,
+    precheck_method: Optional[str] = None,
 ) -> FastAPI:
     instance = import_from_string(app)
     func = get_func(instance, method_name)
@@ -49,6 +63,16 @@ def generate_fast_api(
         plugin_id = hashlib.sha256(
             json.dumps(get_schema_dict(func), sort_keys=True).encode()
         ).hexdigest()[:32]
+
+    precheck_func = None
+    if precheck_str:
+        precheck_instance = import_from_string(precheck_str)
+        precheck_func = get_func(precheck_instance, precheck_method)
+    elif precheck_method:
+        precheck_func = get_func(instance, precheck_method)
+    if precheck_func is not None:
+        check_precheck_func(precheck_func=precheck_func)
+
     logger.debug(f"set static id response to: {plugin_id}")
 
     fastapi_app = FastAPI()
@@ -66,9 +90,8 @@ class InvokeResponse(BaseModel):
 
     logging.getLogger("etl_uvicorn.fastapi")
 
-    usage: list[UsageData] = []
-
     async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> InvokeResponse:
+        usage: list[UsageData] = []
         request_dict = kwargs if kwargs else {}
         if "usage" in inspect.signature(func).parameters:
             request_dict["usage"] = usage
@@ -114,12 +137,29 @@ class SchemaOutputResponse(BaseModel):
     async def docs_redirect():
         return RedirectResponse("/docs")
 
+    class InvokePrecheckResponse(BaseModel):
+        usage: list[UsageData]
+        status_code: int
+        status_code_text: Optional[str] = None
+
     @fastapi_app.get("/schema")
     async def get_schema() -> SchemaOutputResponse:
         schema = get_schema_dict(func)
         resp = SchemaOutputResponse(inputs=schema["inputs"], outputs=schema["outputs"])
         return resp
 
+    @fastapi_app.get("/precheck")
+    async def run_precheck() -> InvokePrecheckResponse:
+        if precheck_func:
+            fn_response = await wrap_fn(func=precheck_func)
+            return InvokePrecheckResponse(
+                status_code=fn_response.status_code,
+                status_code_text=fn_response.status_code_text,
+                usage=fn_response.usage,
+            )
+        else:
+            return InvokePrecheckResponse(status_code=status.HTTP_200_OK, usage=[])
+
     @fastapi_app.get("/id")
     async def get_id() -> str:
         return plugin_id

unstructured_platform_plugins/etl_uvicorn/main.py

@@ -33,6 +33,8 @@ def api_wrapper(
         method_name: Optional[str] = None,
         plugin_id: Optional[str] = None,
         plugin_id_method: Optional[str] = None,
+        precheck_app: Optional[str] = None,
+        precheck_app_method: Optional[str] = None,
         **kwargs,
     ):
         # Make sure logging is configured before the call to run() so any setup has the same format
@@ -44,7 +46,12 @@ def api_wrapper(
         )
         config.configure_logging()
         fastapi_app = generate_fast_api(
-            app=app, method_name=method_name, id_str=plugin_id, id_method=plugin_id_method
+            app=app,
+            method_name=method_name,
+            id_str=plugin_id,
+            id_method=plugin_id_method,
+            precheck_str=precheck_app,
+            precheck_method=precheck_app_method,
         )
         # Explicitly map values that are manipulated in the original
         # call to run(), preventing **kwargs reference
@@ -86,6 +93,22 @@ def api_wrapper(
                 help="If plugin id reference is a class, what method to wrap. "
                 "Will fall back to __call__ if none is provided.",
             ),
+            click.Option(
+                ["--precheck-app"],
+                required=False,
+                type=str,
+                default=None,
+                help="If provided, must point to code to run for precheck",
+            ),
+            click.Option(
+                ["--precheck-app-method"],
+                required=False,
+                type=str,
+                default=None,
+                help="If provided, points to a method to call on a class. "
+                "If precheck-app not provided, assumes method "
+                "lives on main class passes in.",
+            ),
         ]
     )
     return cmd

unstructured_platform_plugins/etl_uvicorn/utils.py

@@ -64,7 +64,7 @@ def get_input_schema(func: Callable, omit: Optional[list[str]] = None) -> dict:
 def get_output_sig(func: Callable) -> Optional[Any]:
     inspect.signature(func)
     type_hints = get_type_hints(func)
-    return_typing = type_hints["return"]
+    return_typing = type_hints.get("return")
     outputs = return_typing if return_typing is not NoneType else None
     return outputs
 
@@ -85,7 +85,7 @@ def map_inputs(func: Callable, raw_inputs: dict[str, Any]) -> dict[str, Any]:
     # types expected by the function when being invoked
     raw_inputs = raw_inputs.copy()
     type_info = get_type_hints(func)
-    type_info.pop("return")
+    type_info.pop("return", None)
     for field_name, type_data in type_info.items():
         if field_name not in raw_inputs:
             continue