bugfix: fix mapping input content when associated with json mixin (#12)

Author: rbiseck3

test/test_schema.py

@@ -1,6 +1,7 @@
 import inspect
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from typing import Any, Optional, Union
 
 import pytest
@@ -408,6 +409,8 @@ def fn(
         e: Optional[dict[str, Any]] = None,
         f: list[float] = None,
         g: Optional[g_enum] = None,
+        h: FileData | None = None,
+        i: Path | None = None,
     ) -> None:
         pass
 
@@ -419,6 +422,8 @@ class ExpectedInputModel(BaseModel):
         e: Optional[dict[str, Any]] = None
         f: list[float] = None
         g: Optional[g_enum] = None
+        h: FileData | None = None
+        i: Path | None = None
 
     input_schema = get_input_schema(fn)
     input_model = js.schema_to_base_model(schema=input_schema)

test/test_utils.py

@@ -2,7 +2,9 @@
 from dataclasses import dataclass, is_dataclass
 from enum import Enum
 
+import pytest
 from pydantic import BaseModel
+from unstructured.ingest.v2.interfaces import FileData
 from uvicorn.importer import import_from_string
 
 from unstructured_platform_plugins.etl_uvicorn import utils
@@ -101,14 +103,20 @@ class MyEnum(Enum):
 
 
 def test_map_inputs():
-    def fn(a: A, b: B, c: MyEnum, d: list) -> None:
+    def fn(a: A, b: B, c: MyEnum, d: list, e: FileData) -> None:
         pass
 
+    file_data = FileData(
+        identifier="custom_file_data",
+        connector_type="mock_connector",
+        additional_metadata={"additional": "metadata"},
+    )
     inputs = {
         "a": {"b": 4, "c": 5.6},
         "b": {"d": True, "e": {"key": "value"}},
         "c": MyEnum.VALUE.value,
         "d": [1, 2, 3],
+        "e": file_data.to_dict(),
     }
 
     mapped_inputs = utils.map_inputs(func=fn, raw_inputs=inputs)
@@ -117,5 +125,16 @@ def fn(a: A, b: B, c: MyEnum, d: list) -> None:
         "b": B(d=True, e={"key": "value"}),
         "c": MyEnum.VALUE.value,
         "d": [1, 2, 3],
+        "e": file_data,
     }
     assert mapped_inputs == expected
+
+
+def test_map_inputs_error():
+    def fn(a: FileData) -> None:
+        pass
+
+    inputs = {"a": {"not": "the", "right": "values"}}
+
+    with pytest.raises(KeyError):
+        utils.map_inputs(func=fn, raw_inputs=inputs)

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -3,10 +3,9 @@
 import inspect
 import json
 import logging
-from pathlib import Path
 from typing import Any, Callable, Optional
 
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
 from uvicorn.importer import import_from_string
@@ -64,27 +63,35 @@ def generate_fast_api(
 
         @fastapi_app.post("/invoke", response_model=response_type)
         async def run_job(request: input_schema_model) -> response_type:
-            logger.debug(f"invoking function: {func}")
-            input_schema = get_input_schema(func)
-            request_dict = request.model_dump()
-            for k, v in request_dict.items():
-                if schema := input_schema.get(k):  # noqa: SIM102
-                    if (
-                        schema.get("type") == "string"
-                        and schema.get("is_path", False)
-                        and isinstance(v, str)
-                    ):
-                        request_dict[k] = Path(v)
+            logger.debug(f"invoking function {func} with input: {request.model_dump()}")
+            # Create dictionary from pydantic model while preserving underlying types
+            request_dict = {f: getattr(request, f) for f in request.model_fields}
             map_inputs(func=func, raw_inputs=request_dict)
             logger.debug(f"passing inputs to function: {request_dict}")
-            return await invoke_func(func=func, kwargs=request_dict)
+            try:
+                return await invoke_func(func=func, kwargs=request_dict)
+            except Exception as e:
+                logger.error(
+                    f"failed to invoke plugin with inputs {request_dict}: {e}", exc_info=True
+                )
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=f"failed to invoke plugin: {e}",
+                )
 
     else:
 
         @fastapi_app.post("/invoke", response_model=response_type)
         async def run_job() -> response_type:
             logger.debug(f"invoking function without inputs: {func}")
-            return await invoke_func(func=func)
+            try:
+                return await invoke_func(func=func)
+            except Exception as e:
+                logger.error(f"failed to invoke plugin: {e}", exc_info=True)
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=f"failed to invoke plugin: {e}",
+                )
 
     class SchemaOutputResponse(BaseModel):
         inputs: dict[str, Any]

unstructured_platform_plugins/etl_uvicorn/utils.py

@@ -4,6 +4,7 @@
 from types import GenericAlias, NoneType
 from typing import Any, Callable, Optional
 
+from dataclasses_json import DataClassJsonMixin
 from pydantic import BaseModel
 
 from unstructured_platform_plugins.schema.json_schema import (
@@ -83,18 +84,29 @@ def map_inputs(func: Callable, raw_inputs: dict[str, Any]) -> dict[str, Any]:
     type_info = get_type_hints(func)
     type_info.pop("return")
     for field_name, type_data in type_info.items():
-        if (
-            is_dataclass(type_data)
-            and field_name in raw_inputs
-            and isinstance(raw_inputs[field_name], dict)
-        ):
-            raw_inputs[field_name] = type_data(**raw_inputs[field_name])
-        elif isinstance(type_data, EnumMeta):
-            raw_inputs[field_name] = raw_inputs[field_name]
-        elif (
-            inspect.isclass(type_data)
-            and not isinstance(type_data, GenericAlias)
-            and issubclass(type_data, BaseModel)
-        ):
-            raw_inputs[field_name] = type_data.model_validate(raw_inputs[field_name])
+        if field_name not in raw_inputs:
+            continue
+        field_value = raw_inputs[field_name]
+        try:
+            if (
+                inspect.isclass(type_data)
+                and issubclass(type_data, DataClassJsonMixin)
+                and isinstance(field_value, dict)
+            ):
+                raw_inputs[field_name] = type_data.from_dict(raw_inputs[field_name])
+            elif is_dataclass(type_data) and isinstance(field_value, dict):
+                raw_inputs[field_name] = type_data(**raw_inputs[field_name])
+            elif isinstance(type_data, EnumMeta):
+                raw_inputs[field_name] = raw_inputs[field_name]
+            elif (
+                inspect.isclass(type_data)
+                and not isinstance(type_data, GenericAlias)
+                and issubclass(type_data, BaseModel)
+            ):
+                raw_inputs[field_name] = type_data.model_validate(raw_inputs[field_name])
+        except Exception as e:
+            exception_type = type(e)
+            raise exception_type(
+                f"failed to map input for field {field_name}: {field_value}"
+            ) from e
     return raw_inputs

unstructured_platform_plugins/schema/json_schema.py

@@ -262,8 +262,12 @@ def response_to_json_schema(return_annotation: Any) -> dict:
     return to_json_schema(val=return_annotation)
 
 
-def schema_to_base_model_type(json_type_name, name: str, type_info: dict) -> Type[BaseModel]:
+def schema_to_base_model_type(json_type_name, name: str, type_info: dict) -> Type:
     t = typed_map_reverse[json_type_name]
+    if t is dict and type_info.get("is_file_data", False):
+        return FileData
+    if t is str and type_info.get("is_path", False):
+        return Path
     if t is dict and "properties" in type_info:
         t = schema_to_base_model(
             schema=type_info["properties"],