Fix typing and minor issues

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -0,0 +1,51 @@
+import inspect
+from typing import Any, Awaitable, Callable, Type, Union, cast
+
+from fastapi import APIRouter, FastAPI, HTTPException, status
+from pydantic import BaseModel, create_model
+
+from ...helper import get_function_metadata_store
+from ...views import Trace
+
+
+def bootstrap_prediction_endpoint(
+    app: FastAPI, func: Callable[..., Union[Trace, Awaitable[Trace]]]
+) -> None:
+    router = APIRouter(
+        tags=["predictions"],
+    )
+
+    schema = _get_schema(func)
+
+    @router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace)
+    async def predict(input_value: schema) -> Trace:  # type: ignore
+        try:
+            if inspect.iscoroutinefunction(func):
+                return await cast(Callable[..., Awaitable[Trace]], func)(
+                    **cast(BaseModel, input_value).dict()
+                )
+            return cast(Callable[..., Trace], func)(
+                **cast(BaseModel, input_value).dict()
+            )
+        except Exception as e:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=f"The following exception has occurred: {type(e).__name__}: {e}",
+            )
+
+    app.include_router(router)
+
+
+def _get_schema(func: Callable) -> Type[BaseModel]:
+    signature = inspect.signature(func)
+    parameters = {
+        p.name: (
+            p.annotation if p.annotation != inspect._empty else Any,
+            p.default if p.default != inspect._empty else ...,
+        )
+        for p in signature.parameters.values()
+        if p.name in get_function_metadata_store(func).input_parameter_names
+    }
+
+    schema: Type[BaseModel] = create_model("InputModel", **parameters)  # type: ignore
+    return schema