unstructured-platform-plugins 0.0.0__py3-none-any.whl

test/test_utils.py

@@ -0,0 +1,140 @@
+import inspect
+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
+
+
+def test_get_func_simple():
+    instance = import_from_string("test.assets.typed_dict_response:sample_function")
+    func = utils.get_func(instance)
+    assert callable(func)
+    assert func.__name__ == "sample_function"
+    sig = inspect.signature(func)
+    response = sig.return_annotation
+    # Check for typed dict response
+    assert dict in inspect.getmro(response)
+    assert dict is not inspect.getmro(response)[0]
+
+
+def test_get_async_func():
+    instance = import_from_string("test.assets.async_typed_dict_response:async_sample_function")
+    func = utils.get_func(instance)
+    assert callable(func)
+    assert func.__name__ == "async_sample_function"
+    sig = inspect.signature(func)
+    response = sig.return_annotation
+    # Check for typed dict response
+    assert dict in inspect.getmro(response)
+    assert dict is not inspect.getmro(response)[0]
+
+
+def test_get_class_func():
+    instance = import_from_string("test.assets.pydantic_response_class_method:SampleClass")
+    func = utils.get_func(instance, method_name="sample_method")
+    assert callable(func)
+    assert func.__name__ == "sample_method"
+    sig = inspect.signature(func)
+    response = sig.return_annotation
+    # Check for pydantic model response
+    assert inspect.isclass(response)
+    assert issubclass(response, BaseModel)
+
+
+def test_get_func_dataclass_response():
+    instance = import_from_string("test.assets.dataclass_response:sample_function_with_path")
+    func = utils.get_func(instance)
+    assert callable(func)
+    assert func.__name__ == "sample_function_with_path"
+    sig = inspect.signature(func)
+    response = sig.return_annotation
+    # Check for dataclass response
+    assert is_dataclass(response)
+
+
+def test_get_plugin_id_value():
+    instance = import_from_string("test.assets.simple_hash_value:hash_value")
+    hash_value = utils.get_plugin_id(instance)
+    assert hash_value == "plugin_id_123"
+
+
+def test_get_plugin_id_lambda():
+    instance = import_from_string("test.assets.simple_hash_lambda:hash_lambda_fn")
+    hash_value = utils.get_plugin_id(instance)
+    assert hash_value == "plugin_id_hash_123"
+
+
+def test_get_plugin_id_function():
+    instance = import_from_string("test.assets.hash_function:get_hash")
+    hash_value = utils.get_plugin_id(instance)
+    assert hash_value == "plugin_id_fn_123"
+
+
+def test_get_plugin_id_class():
+    instance = import_from_string("test.assets.simple_hash_class:GetHash")
+    hash_value = utils.get_plugin_id(instance, method_name="my_hash")
+    assert hash_value == "plugin_id_class_123"
+
+
+def test_get_plugin_id_class_instance():
+    instance = import_from_string("test.assets.simple_hash_class:get_hash_class_instance")
+    hash_value = utils.get_plugin_id(instance, method_name="my_hash")
+    assert hash_value == "plugin_id_class_123"
+
+
+@dataclass
+class A:
+    b: int
+    c: float
+
+
+class B(BaseModel):
+    d: bool
+    e: dict
+
+
+class MyEnum(Enum):
+    VALUE = "value"
+
+
+def test_map_inputs():
+    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)
+    expected = {
+        "a": A(b=4, c=5.6),
+        "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/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.0.0"  # pragma: no cover

unstructured_platform_plugins/etl_uvicorn/api_generator.py

@@ -0,0 +1,173 @@
+import asyncio
+import hashlib
+import inspect
+import json
+import logging
+from typing import Any, Callable, Optional
+
+from fastapi import FastAPI, status
+from pydantic import BaseModel
+from starlette.responses import RedirectResponse
+from uvicorn.importer import import_from_string
+
+from unstructured_platform_plugins.etl_uvicorn.utils import (
+    get_func,
+    get_input_schema,
+    get_output_sig,
+    get_plugin_id,
+    get_schema_dict,
+    map_inputs,
+)
+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")
+
+
+async def invoke_func(func: Callable, kwargs: Optional[dict[str, Any]] = None) -> Any:
+    kwargs = kwargs or {}
+    if inspect.iscoroutinefunction(func):
+        return await func(**kwargs)
+    else:
+        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)
+    if id_str:
+        id_ref = import_from_string(id_str)
+        plugin_id = get_plugin_id(instance=id_ref, method_name=id_method)
+    else:
+        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()
+
+    response_type = get_output_sig(func)
+
+    class InvokeResponse(BaseModel):
+        usage: list[UsageData]
+        status_code: int
+        status_code_text: Optional[str] = None
+        output: Optional[response_type] = None
+
+    input_schema = get_input_schema(func, omit=["usage"])
+    input_schema_model = schema_to_base_model(input_schema)
+
+    logging.getLogger("etl_uvicorn.fastapi")
+
+    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
+        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)
+        except Exception as invoke_error:
+            logger.error(f"failed to invoke plugin: {invoke_error}", exc_info=True)
+            return InvokeResponse(
+                usage=usage,
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                status_code_text=f"failed to invoke plugin: "
+                f"[{invoke_error.__class__.__name__}] {invoke_error}",
+            )
+
+    if input_schema_model.model_fields:
+
+        @fastapi_app.post("/invoke", response_model=InvokeResponse)
+        async def run_job(request: input_schema_model) -> InvokeResponse:
+            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 wrap_fn(func=func, kwargs=request_dict)
+
+    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 wrap_fn(
+                func=func,
+            )
+
+    class SchemaOutputResponse(BaseModel):
+        inputs: dict[str, Any]
+        outputs: dict[str, Any]
+
+    @fastapi_app.get("/", include_in_schema=False)
+    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
+
+    # Run initial schema validation
+    try:
+        asyncio.run(get_schema())
+    except TypeError as e:
+        raise TypeError(f"failed to validate function schema: {e}") from e
+
+    return fastapi_app

unstructured_platform_plugins/etl_uvicorn/main.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass, field
+from typing import IO, Any, Optional
+
+import click
+from uvicorn.config import LOGGING_CONFIG, Config, RawConfigParser
+from uvicorn.main import main, run
+
+from unstructured_platform_plugins.etl_uvicorn.api_generator import generate_fast_api
+
+
+@dataclass
+class CustomConfig:
+    log_config: dict[str, Any] | str | RawConfigParser | IO[Any] | None = field(
+        default_factory=lambda: LOGGING_CONFIG
+    )
+    log_level: str | int | None = None
+    use_colors: bool | None = None
+    access_log: bool = True
+
+
+CustomConfig.configure_logging = Config.configure_logging
+
+
+def get_command() -> click.Command:
+    @click.command(context_settings={"auto_envvar_prefix": "UVICORN"})
+    def api_wrapper(
+        app: str,
+        log_config: str,
+        reload_dirs: list[str],
+        reload_includes: list[str],
+        reload_excludes: list[str],
+        headers: list[str],
+        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
+        config = CustomConfig(
+            log_config=LOGGING_CONFIG if log_config is None else log_config,
+            log_level=kwargs["log_level"],
+            use_colors=kwargs["use_colors"],
+            access_log=kwargs["access_log"],
+        )
+        config.configure_logging()
+        fastapi_app = generate_fast_api(
+            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
+        run(
+            fastapi_app,
+            log_config=LOGGING_CONFIG if log_config is None else log_config,
+            reload_dirs=reload_dirs or None,
+            reload_includes=reload_includes or None,
+            reload_excludes=reload_excludes or None,
+            headers=[header.split(":", 1) for header in headers],  # type: ignore[misc]
+            **kwargs,
+        )
+
+    cmd = api_wrapper
+    cmd.params = main.params
+    cmd.params.extend(
+        [
+            click.Option(
+                ["--method-name"],
+                required=False,
+                type=str,
+                default=None,
+                help="If passed in instance is a class, what method to wrap. "
+                "Will fall back to __call__ if none is provided.",
+            ),
+            click.Option(
+                ["--plugin-id"],
+                required=False,
+                type=str,
+                default=None,
+                help="Reference to either a value or function to get "
+                "the plugin id once instantiated",
+            ),
+            click.Option(
+                ["--plugin-id-method"],
+                required=False,
+                type=str,
+                default=None,
+                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
+
+
+cmd = get_command()
+
+if __name__ == "__main__":
+    cmd = get_command()
+    cmd()

unstructured_platform_plugins/etl_uvicorn/utils.py

@@ -0,0 +1,115 @@
+import inspect
+from dataclasses import is_dataclass
+from enum import EnumMeta
+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 (
+    parameters_to_json_schema,
+    response_to_json_schema,
+)
+from unstructured_platform_plugins.schema.utils import get_typed_parameters
+from unstructured_platform_plugins.type_hints import get_type_hints
+
+
+def get_func(instance: Any, method_name: Optional[str] = None) -> Callable:
+    method_name = method_name or "__call__"
+    if inspect.isfunction(instance):
+        return instance
+    elif inspect.isclass(instance):
+        i = instance()
+        return getattr(i, method_name)
+    elif isinstance(instance, object) and hasattr(instance, method_name):
+        func = getattr(instance, method_name)
+        if inspect.ismethod(func):
+            return func
+    raise ValueError(f"type of instance not recognized: {type(instance)}")
+
+
+def get_plugin_id(instance: Any, method_name: Optional[str] = None) -> str:
+    method_name = method_name or "__call__"
+    ref_id = None
+    if inspect.isfunction(instance):
+        ref_id = instance()
+    elif inspect.isclass(instance):
+        i = instance()
+        method_name = method_name or "__call__"
+        fn = getattr(i, method_name)
+        ref_id = fn()
+    elif isinstance(instance, object) and hasattr(instance, method_name):
+        func = getattr(instance, method_name)
+        if inspect.ismethod(func):
+            ref_id = func()
+    else:
+        ref_id = instance
+    if not ref_id:
+        raise ValueError(f"id could not be parsed from instance {instance}")
+    ref_id = str(ref_id)
+    if not ref_id.isidentifier():
+        raise ValueError(f"'{ref_id}' is not a valid identifier")
+    return ref_id
+
+
+def get_input_schema(func: Callable, omit: Optional[list[str]] = None) -> dict:
+
+    parameters = get_typed_parameters(func)
+    if omit:
+        parameters = [p for p in parameters if p.name not in omit]
+    return parameters_to_json_schema(parameters)
+
+
+def get_output_sig(func: Callable) -> Optional[Any]:
+    inspect.signature(func)
+    type_hints = get_type_hints(func)
+    return_typing = type_hints.get("return")
+    outputs = return_typing if return_typing is not NoneType else None
+    return outputs
+
+
+def get_output_schema(func: Callable) -> dict:
+    return response_to_json_schema(get_output_sig(func))
+
+
+def get_schema_dict(func, omit: list[str] = ["usage"]) -> dict:
+    return {
+        "inputs": get_input_schema(func, omit=omit),
+        "outputs": get_output_schema(func),
+    }
+
+
+def map_inputs(func: Callable, raw_inputs: dict[str, Any]) -> dict[str, Any]:
+    # deserializes the raw dictionary coming in from the api into the underlying data
+    # types expected by the function when being invoked
+    raw_inputs = raw_inputs.copy()
+    type_info = get_type_hints(func)
+    type_info.pop("return", None)
+    for field_name, type_data in type_info.items():
+        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