Skip to content

Exposing docstrings of handler_predict method on API docs in when using custom handlers for predict #177

@krumeto

Description

@krumeto

Is your feature request related to a problem? Please describe.

Vetiver API (and FastAPI) nicely display the docstrings of the function used to produce a given endpoint as a documentation on the API docs. However, when overwriting the predict method with a custom one, I have troubles making the docstring of the predict_handler available as docs.

For example, I have the following handler that is used as a basis for my predict and I'd like the docstring of handler_predict to be available as docs in the API docs:

from vetiver.handlers.base import BaseHandler
from vetiver import VetiverModel
from pydantic import BaseModel, Field, constr
from vetiver.prototype import vetiver_create_prototype
from pydantic.main import ModelMetaclass
import pandas as pd
import numpy as np


# Define a Pydantic model to validate the request body data
class RecommendationsRequest(BaseModel):
    param_1: constr(min_length=36, max_length=36, strip_whitespace=True)
    param_2: int = Field(gt=0,  default=5)
    param_3: int = Field(gt=0,  default=60)
    param_4: int = Field(gt=0,  default=1)


@vetiver_create_prototype.register
def _(data: ModelMetaclass):
    return data


class CustomHandler(BaseHandler):
    """A handler to wrap the predict method required by vetiver."""

    def __init__(self, model, prototype_data = RecommendationsRequest):

        super().__init__(model, prototype_data)
        model_type = staticmethod(lambda: 'custom handler')

    def handler_predict(self, request: RecommendationsRequest, check_prototype):
        """
        A method to overwrite the default `predict` method of the Vetiver BaseHandler
        
        Here follows a documentation that I'd like to see in the API docs.
        """

        if isinstance(request, list):
            param_1 = request[0]
            param_2 = request[1]
            param_3 = request[2]
            param_4 = request[3]
        else:
            param_1 = request.loc[0, 'param_1']
            # The item() below because otherwise it returns numpy.int64 instead of python ints
            # All numerics need that
            param_2 = request.loc[0, 'param_2'].item()
            param_3 = request.loc[0, 'param_3'].item()
            param_4 = request.loc[0, 'param_4'].item()

        prediction = self.model.get_article_recommendations(param_1=param_1,
                                                                      param_2=param_2,
                                                                      param_3=param_3,
                                                                      param_4=param_4)

        try:
            results = pd.Series(prediction['results'][0]) # numpy array, because Vetiver requires the results to have a `tolist()` method
            
        #FIXME Improve response when things are wrong
        except AttributeError:
            results = np.array([None])
        return results

Describe the solution you'd like
I'd like the docstrings of the handler_predict method to be viewable in the API docs.

Thank you in advance!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions