add support for filedata_meta to expose more data other than typical input/output

Author: rbiseck3

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.9
+
+* **Support optionally exposing addition metadata around FileData**
+
 ## 0.0.8
 
 * **reduce usage data log level** We do not want to have so much verbosity for something that might happen a lot

test/assets/filedata_meta.py

@@ -0,0 +1,33 @@
+import math
+from typing import Optional
+
+from pydantic import BaseModel
+from unstructured_ingest.v2.interfaces import FileData
+
+from unstructured_platform_plugins.schema import FileDataMeta, NewRecord
+
+
+class Input(BaseModel):
+    m: int
+
+
+class Output(BaseModel):
+    n: float
+
+
+def process_input(i: Input, file_data: FileData, filedata_meta: FileDataMeta) -> Optional[Output]:
+    if i.m > 10:
+        filedata_meta.terminate_current = True
+        new_content = [
+            NewRecord(
+                file_data=FileData(
+                    identifier=str(i.m + x), connector_type=file_data.connector_type
+                ),
+                contents=Output(n=float(i.m + x)),
+            )
+            for x in range(5)
+        ]
+        filedata_meta.new_records.extend(new_content)
+        return None
+    else:
+        return Output(n=math.pi + i.m)

unstructured_platform_plugins/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.8"  # pragma: no cover
+__version__ = "0.0.9"  # pragma: no cover

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -23,11 +23,10 @@
     get_schema_dict,
     map_inputs,
 )
-from unstructured_platform_plugins.exceptions import UnrecoverableException
+from unstructured_platform_plugins.schema import FileDataMeta, UsageData
 from unstructured_platform_plugins.schema.json_schema import (
     schema_to_base_model,
 )
-from unstructured_platform_plugins.schema.usage import UsageData
 
 logger = logging.getLogger("uvicorn.error")
 
@@ -85,10 +84,11 @@ def wrap_in_fastapi(
     class InvokeResponse(BaseModel):
         usage: list[UsageData]
         status_code: int
+        filedata_meta: FileDataMeta
         status_code_text: Optional[str] = None
         output: Optional[response_type] = None
 
-    input_schema = get_input_schema(func, omit=["usage"])
+    input_schema = get_input_schema(func, omit=["usage", "filedata_meta"])
     input_schema_model = schema_to_base_model(input_schema)
 
     logging.getLogger("etl_uvicorn.fastapi")
@@ -97,36 +97,41 @@ class InvokeResponse(BaseModel):
 
     async def wrap_fn(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> ResponseType:
         usage: list[UsageData] = []
+        filedata_meta = FileDataMeta()
         request_dict = kwargs if kwargs else {}
         if "usage" in inspect.signature(func).parameters:
             request_dict["usage"] = usage
         else:
-            logger.debug("usage data not an expected parameter, omitting")
+            logger.warning("usage data not an expected parameter, omitting")
+        if "filedata_meta" in inspect.signature(func).parameters:
+            request_dict["filedata_meta"] = filedata_meta
         try:
             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
+                            usage=usage,
+                            filedata_meta=filedata_meta,
+                            status_code=status.HTTP_200_OK,
+                            output=output,
                         ).model_dump_json() + "\n"
 
                 return StreamingResponse(_stream_response(), media_type="application/x-ndjson")
             else:
-                try:
-                    output = await invoke_func(func=func, kwargs=request_dict)
-                    return InvokeResponse(
-                        usage=usage, status_code=status.HTTP_200_OK, output=output
-                    )
-                except UnrecoverableException as ex:
-                    # Thrower of this exception is responsible for logging necessary information
-                    logger.info("Unrecoverable error occurred during plugin invocation")
-                    return InvokeResponse(usage=usage, status_code=512, status_code_text=ex.message)
+                output = await invoke_func(func=func, kwargs=request_dict)
+                return InvokeResponse(
+                    usage=usage,
+                    filedata_meta=filedata_meta,
+                    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(
                 usage=usage,
+                filedata_meta=filedata_meta,
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 status_code_text=f"failed to invoke plugin: "
                 f"[{invoke_error.__class__.__name__}] {invoke_error}",

unstructured_platform_plugins/etl_uvicorn/otel.py

@@ -42,7 +42,6 @@ def get_settings() -> OtelSettings:
 
 def get_trace_provider() -> TracerProvider:
     settings = get_settings()
-    print(settings)
     provider = TracerProvider(resource=Resource({SERVICE_NAME: settings["service_name"]}))
 
     for trace_exporter_type in settings["trace_exporters"]:

unstructured_platform_plugins/etl_uvicorn/utils.py

@@ -111,10 +111,8 @@ def map_inputs(func: Callable, raw_inputs: dict[str, Any]) -> dict[str, Any]:
                 and not isinstance(type_data, GenericAlias)
                 and issubclass(type_data, BaseModel)
             ):
-                raw_inputs[field_name] = type_data.model_validate(raw_inputs[field_name])
+                field_value = raw_inputs[field_name]
+                raw_inputs[field_name] = type_data.model_validate(field_value.model_dump())
         except Exception as e:
-            exception_type = type(e)
-            raise exception_type(
-                f"failed to map input for field {field_name}: {field_value}"
-            ) from e
+            raise ValueError(f"failed to map input for field {field_name}: {field_value}") from e
     return raw_inputs

unstructured_platform_plugins/schema/__init__.py

@@ -0,0 +1,4 @@
+from .filedata_meta import FileDataMeta, NewRecord
+from .usage import UsageData
+
+__all__ = ["UsageData", "FileDataMeta", "NewRecord"]

unstructured_platform_plugins/schema/filedata_meta.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+from unstructured_ingest.v2.interfaces import FileData
+
+
+class NewRecord(BaseModel):
+    file_data: FileData
+    contents: Any
+
+
+class FileDataMeta(BaseModel):
+    terminate_current: bool = False
+    new_records: list[NewRecord] = Field(default_factory=list)

unstructured_platform_plugins/schema/json_schema.py

@@ -113,6 +113,12 @@ def dataclass_to_json_schema(class_or_instance: Any) -> dict:
         f_resp = to_json_schema(t)
         if f.default is not MISSING:
             f_resp["default"] = f.default.value if isinstance(f.default, Enum) else f.default
+        elif f.default_factory is not MISSING:
+            default = f.default_factory()
+            try:
+                f_resp["default"] = to_json_schema(default)
+            except Exception:
+                f_resp["default"] = default
         else:
             required.append(f.name)
         properties[f.name] = f_resp
@@ -339,6 +345,10 @@ def schema_to_base_model(schema: dict, name: str = "reconstructed_model") -> Typ
                 t = schema_to_base_model_type(
                     json_type_name=json_type_name, name=k, type_info=entry_info
                 )
+        elif "properties" in v:
+            t = schema_to_base_model(schema=v, name=k)
+        elif not v:
+            t = schema_to_base_model(schema={"type": "null"}, name=k)
         else:
             json_type_name = v["type"]
             t = schema_to_base_model_type(json_type_name=json_type_name, name=k, type_info=v)