51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
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
|