add SklearnPredictionTypes

Author: isabelizimm

vetiver/handlers/base.py

@@ -121,7 +121,7 @@ def handler_startup():
         """
         ...
 
-    def handler_predict(self, input_data, check_prototype):
+    def handler_predict(self, input_data, check_prototype, **kw):
         """Generates method for /predict endpoint in VetiverAPI
 
         The `handler_predict` function executes at each API call. Use this

vetiver/handlers/spacy.py

@@ -53,7 +53,7 @@ def construct_prototype(self):
 
         return prototype
 
-    def handler_predict(self, input_data, check_prototype):
+    def handler_predict(self, input_data, check_prototype, **kw):
         """
         Generates method for /predict endpoint in VetiverAPI
 

vetiver/handlers/statsmodels.py

@@ -22,7 +22,7 @@ class StatsmodelsHandler(BaseHandler):
     if sm_exists:
         pip_name = "statsmodels"
 
-    def handler_predict(self, input_data, check_prototype):
+    def handler_predict(self, input_data, check_prototype, **kw):
         """
         Generates method for /predict endpoint in VetiverAPI
 

vetiver/handlers/torch.py

@@ -22,7 +22,7 @@ class TorchHandler(BaseHandler):
     if torch_exists:
         pip_name = "torch"
 
-    def handler_predict(self, input_data, check_prototype):
+    def handler_predict(self, input_data, check_prototype, **kw):
         """
         Generates method for /predict endpoint in VetiverAPI
 

vetiver/handlers/xgboost.py

@@ -22,7 +22,7 @@ class XGBoostHandler(BaseHandler):
     if xgb_exists:
         pip_name = "xgboost"
 
-    def handler_predict(self, input_data, check_prototype):
+    def handler_predict(self, input_data, check_prototype, **kw):
         """
         Generates method for /predict endpoint in VetiverAPI
 

vetiver/server.py

@@ -11,15 +11,16 @@
 import pandas as pd
 import requests
 import uvicorn
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from fastapi.exceptions import RequestValidationError
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import HTMLResponse, PlainTextResponse, RedirectResponse
-
 from .helpers import api_data_to_frame, response_to_frame
+from .handlers.sklearn import SKLearnHandler
 from .meta import VetiverMeta
 from .utils import _jupyter_nb, get_workbench_path
 from .vetiver_model import VetiverModel
+from .types import SklearnPredictionTypes
 
 
 class VetiverAPI:
@@ -111,7 +112,6 @@ async def startup_event():
 
         @app.get("/", include_in_schema=False)
         def docs_redirect():
-
             redirect = "__docs__"
 
             return RedirectResponse(redirect)
@@ -200,65 +200,75 @@ async def validation_exception_handler(request, exc):
 
         return app
 
-    def vetiver_post(self, endpoint_fx: Callable, endpoint_name: str = None, **kw):
-        """Create new POST endpoint that is aware of model input data
+    def vetiver_post(
+        self,
+        endpoint_fx: Union[Callable, SklearnPredictionTypes],
+        endpoint_name: str = None,
+        **kw,
+    ):
+        """Define a new POST endpoint that utilizes the model's input data.
 
         Parameters
         ----------
-        endpoint_fx : typing.Callable
-            Custom function to be run at endpoint
+        endpoint_fx : Union[typing.Callable, Literal["predict", "predict_proba", "predict_log_proba"]]
+            A callable function that specifies the custom logic to execute when the endpoint is called.
+            This function should take input data (e.g., a DataFrame or dictionary) and return the desired output
+            (e.g., predictions or transformed data). For scikit-learn models, endpoint_fx can also be one of
+            "predict", "predict_proba", or "predict_log_proba" if the model supports these methods.
+
         endpoint_name : str
-            Name of endpoint
+            The name of the endpoint to be created.
 
         Examples
         -------
-        ```{python}
+        ```python
         from vetiver import mock, VetiverModel, VetiverAPI
         X, y = mock.get_mock_data()
         model = mock.get_mock_model().fit(X, y)
 
-        v = VetiverModel(model = model, model_name = "model", prototype_data = X)
-        v_api = VetiverAPI(model = v, check_prototype = True)
+        v = VetiverModel(model=model, model_name="model", prototype_data=X)
+        v_api = VetiverAPI(model=v, check_prototype=True)
 
         def sum_values(x):
             return x.sum()
+
         v_api.vetiver_post(sum_values, "sums")
         ```
         """
-        if not endpoint_name:
-            endpoint_name = endpoint_fx.__name__
 
-        if endpoint_fx.__doc__ is not None:
-            api_desc = dedent(endpoint_fx.__doc__)
-        else:
-            api_desc = None
-
-        if self.check_prototype is True:
-
-            @self.app.post(
-                urljoin("/", endpoint_name),
-                name=endpoint_name,
-                description=api_desc,
+        if isinstance(endpoint_fx, SklearnPredictionTypes):
+            if not isinstance(self.model, SKLearnHandler):
+                raise ValueError(
+                    "The 'endpoint_fx' parameter can only be a string when using scikit-learn models."
+                )
+            self.vetiver_post(
+                self.model.handler_predict,
+                SklearnPredictionTypes,
+                check_prototype=self.check_prototype,
+                prediction_type=endpoint_fx,
             )
-            async def custom_endpoint(input_data: List[self.model.prototype]):
-                _to_frame = api_data_to_frame(input_data)
-                predictions = endpoint_fx(_to_frame, **kw)
-                if isinstance(predictions, List):
-                    return {endpoint_name: predictions}
-                else:
-                    return predictions
+            return
 
-        else:
+        endpoint_name = endpoint_name or endpoint_fx.__name__
+        endpoint_doc = dedent(endpoint_fx.__doc__) if endpoint_fx.__doc__ else None
 
-            @self.app.post(urljoin("/", endpoint_name))
-            async def custom_endpoint(input_data: Request):
+        @self.app.post(
+            urljoin("/", endpoint_name),
+            name=endpoint_name,
+            description=endpoint_doc,
+        )
+        async def custom_endpoint(input_data: List[self.model.prototype]):
+            if self.check_prototype:
+                served_data = api_data_to_frame(input_data)
+            else:
                 served_data = await input_data.json()
-                predictions = endpoint_fx(served_data, **kw)
 
-                if isinstance(predictions, List):
-                    return {endpoint_name: predictions}
-                else:
-                    return predictions
+            predictions = endpoint_fx(served_data, **kw)
+
+            if isinstance(predictions, List):
+                return {endpoint_name: predictions}
+            else:
+                return predictions
 
     def run(self, port: int = 8000, host: str = "127.0.0.1", quiet_open=False, **kw):
         """

vetiver/types.py

@@ -1,4 +1,5 @@
 from pydantic import BaseModel, create_model
+from typing import Literal
 
 all = ["Prototype", "create_prototype"]
 
@@ -7,5 +8,8 @@ class Prototype(BaseModel):
     pass
 
 
+SklearnPredictionTypes = Literal["predict", "predict_proba", "predict_log_proba"]
+
+
 def create_prototype(**dict_data):
     return create_model("prototype", __base__=Prototype, **dict_data)