Merge pull request #174 from rstudio/prototype-endpoint

Author: isabelizimm

.github/workflows/tests.yml

@@ -115,6 +115,23 @@ jobs:
         run: |
           pytest vetiver/tests/test_sklearn.py
 
+  test-pydantic-old:
+    name: "Test pydantic v1"
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e .[dev]
+          python -m pip install 'pydantic<2.0.0'
+
+      - name: Run tests
+        run: |
+          make test
   typecheck:
     runs-on: ubuntu-latest
 

.gitignore

@@ -144,3 +144,5 @@ NOTES.md
 .Rproj.user
 
 /.quarto/
+
+/.luarc.json

vetiver/__init__.py

@@ -2,7 +2,11 @@
 # Change to import.metadata when minimum python>=3.8
 from importlib_metadata import version as _version
 
-from .prototype import vetiver_create_prototype, vetiver_create_ptype  # noqa
+from .prototype import (  # noqa
+    vetiver_create_prototype,  # noqa
+    vetiver_create_ptype,  # noqa
+    InvalidPTypeError,  # noqa
+)  # noqa
 from .vetiver_model import VetiverModel  # noqa
 from .server import VetiverAPI, vetiver_endpoint, predict  # noqa
 from .mock import get_mock_data, get_mock_model  # noqa

vetiver/prototype.py

@@ -175,6 +175,22 @@ def _(data: dict):
     data : dict
         Dictionary
     """
+
+    # if it comes from vetiver's /prototype endpoint
+    dict_data = {}
+    if data.keys() >= {"properties", "title", "type"}:
+        # automatically create for simple prototypes
+        try:
+            for key, value in data["properties"].items():
+                dict_data.update({key: (type(value["default"]), value["default"])})
+        # error for complex objects
+        except KeyError:
+            raise InvalidPTypeError(
+                "Failed to create dict prototype for this data. "
+                "Please use pandas DataFrame or pydantic BaseModel."
+            )
+        return create_prototype(**dict_data)
+
     return create_prototype(**_to_field(data))
 
 

vetiver/server.py

@@ -3,6 +3,7 @@
 
 import re
 import httpx
+import json
 import pandas as pd
 import requests
 import uvicorn
@@ -27,6 +28,8 @@ class VetiverAPI:
     ----------
     model :  VetiverModel
         Model to be deployed in API
+    show_prototype: bool = True
+
     check_prototype : bool
         Determine if data prototype should be enforced
     app_factory :
@@ -60,6 +63,7 @@ class VetiverAPI:
     def __init__(
         self,
         model: VetiverModel,
+        show_prototype: bool = True,
         check_prototype: bool = True,
         app_factory=FastAPI,
         **kwargs,
@@ -80,6 +84,7 @@ def __init__(
             self.model.prototype = self.model.ptype
             delattr(self.model, "ptype")
 
+        self.show_prototype = show_prototype
         self.check_prototype = check_prototype
 
         self._init_app()
@@ -114,6 +119,23 @@ async def get_metadata():
             """Get metadata from model"""
             return self.model.metadata.to_dict()
 
+        if self.show_prototype is True:
+
+            @app.get("/prototype")
+            async def get_prototype():
+                # to handle pydantic<2 and >=2
+                prototype_schema = getattr(
+                    self.model.prototype,
+                    "model_json_schema",
+                    self.model.prototype.schema_json,
+                )()
+                # pydantic<2 returns a string, need to handle to json format
+                if isinstance(prototype_schema, str):
+                    prototype_schema = json.loads(prototype_schema)
+                for key, value in prototype_schema["properties"].items():
+                    value.pop("title", None)
+                return prototype_schema
+
         self.vetiver_post(
             self.model.handler_predict, "predict", check_prototype=self.check_prototype
         )

vetiver/tests/test_server.py

@@ -1,11 +1,20 @@
-from vetiver import mock, VetiverModel, VetiverAPI
+from vetiver import (
+    mock,
+    VetiverModel,
+    VetiverAPI,
+    vetiver_create_prototype,
+    InvalidPTypeError,
+)
+from pydantic import BaseModel, conint
 from fastapi.testclient import TestClient
+import numpy as np
 import pytest
 import sys
 
 
 @pytest.fixture
 def vetiver_model():
+    np.random.seed(500)
     X, y = mock.get_mock_data()
     model = mock.get_mock_model().fit(X, y)
     v = VetiverModel(
@@ -25,6 +34,31 @@ def client(vetiver_model):
     return TestClient(app.app)
 
 
+@pytest.fixture
+def complex_prototype_model():
+    np.random.seed(500)
+
+    class CustomPrototype(BaseModel):
+        B: conint(gt=42)
+        C: conint(gt=42)
+        D: conint(gt=42)
+
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        # move to model_construct for pydantic 3
+        prototype_data=CustomPrototype.construct(),
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    # dont actually want to make predictions, just for looking at schema
+    app = VetiverAPI(v, check_prototype=False)
+
+    return TestClient(app.app)
+
+
 def test_get_ping(client):
     response = client.get("/ping")
     assert response.status_code == 200, response.text
@@ -46,3 +80,40 @@ def test_get_metadata(client):
         "required_pkgs": ["scikit-learn"],
         "python_version": list(sys.version_info),  # JSON will return a list
     }
+
+
+def test_get_prototype(client, vetiver_model):
+    response = client.get("/prototype")
+    assert response.status_code == 200, response.text
+    assert response.json() == {
+        "properties": {
+            "B": {"default": 55, "type": "integer"},
+            "C": {"default": 65, "type": "integer"},
+            "D": {"default": 17, "type": "integer"},
+        },
+        "title": "prototype",
+        "type": "object",
+    }
+
+    assert (
+        vetiver_model.prototype.construct().dict()
+        == vetiver_create_prototype(response.json()).construct().dict()
+    )
+
+
+def test_complex_prototype(complex_prototype_model):
+    response = complex_prototype_model.get("/prototype")
+    assert response.status_code == 200, response.text
+    assert response.json() == {
+        "properties": {
+            "B": {"exclusiveMinimum": 42, "type": "integer"},
+            "C": {"exclusiveMinimum": 42, "type": "integer"},
+            "D": {"exclusiveMinimum": 42, "type": "integer"},
+        },
+        "required": ["B", "C", "D"],
+        "title": "CustomPrototype",
+        "type": "object",
+    }
+
+    with pytest.raises(InvalidPTypeError):
+        vetiver_create_prototype(response.json())