-
Notifications
You must be signed in to change notification settings - Fork 18
Closed
Description
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!
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
No labels