catching errors at prediction time

Author: isabelizimm

vetiver/server.py

@@ -191,15 +191,27 @@ def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw):
     else: 
         requester = requests
 
-    if isinstance(data, (pd.DataFrame,pd.Series)):
+    # TO DO: arrow format 
+
+    if isinstance(data, pd.DataFrame):
         data_json = data.to_json(orient="records")
         response = requester.post(endpoint, data=data_json, **kw)
+    elif isinstance(data, pd.Series):
+        data_dict = data.to_json()
+        response = requester.post(endpoint, data=data_dict, **kw)
     elif isinstance(data, dict):
         response = requester.post(endpoint, json=data, **kw)
     else:
-        raise TypeError(f"Given type is {type(data)}")
-        
+        try:
+            response = requester.post(endpoint, json=data, **kw)
+        except:
+            raise TypeError(f"Predict expects a DataFrame or dict. Given type is {type(data)}")
+
     response_df = pd.DataFrame.from_dict(response.json())
+    
+    if isinstance(response_df.iloc[0,0], dict):
+        if 'type_error.dict' in response_df.iloc[0,0].values():
+            raise TypeError(f"Predict expects a DataFrame or dict. Given type is {type(data)}")
 
     return response_df
 

vetiver/tests/test_predict.py

@@ -28,6 +28,7 @@ def test_predict_sklearn_dict_ptype():
 
     assert isinstance(response, pd.DataFrame), response
     assert response.iloc[0,0] == 44.47
+    assert len(response) == 1
 
 
 def test_predict_sklearn_no_ptype():
@@ -74,9 +75,10 @@ def test_predict_sklearn_df_check_ptype():
     assert len(response) == 100
 
 
-def test_predict_sklearn_df_no_ptype():
+def test_predict_sklearn_series_check_ptype():
     np.random.seed(500)
     X, y = mock.get_mock_data()
+    ser = pd.Series(data=[0,0,0])
     model = mock.get_mock_model().fit(X, y)
     v = VetiverModel(
         model=model,
@@ -86,13 +88,14 @@ def test_predict_sklearn_df_no_ptype():
         versioned=None,
         description="A regression model for testing purposes",
     )
-    app = VetiverAPI(v, check_ptype=False)
+    app = VetiverAPI(v, check_ptype=True)
     client = TestClient(app.app)
 
-    response = predict(endpoint=client, data=X)
+    response = predict(endpoint=client, data=ser)
 
     assert isinstance(response, pd.DataFrame), response
     assert response.iloc[0,0] == 44.47
+    assert len(response) == 1
 
 
 def test_predict_sklearn_type_error():
@@ -109,7 +112,7 @@ def test_predict_sklearn_type_error():
     )
     app = VetiverAPI(v, check_ptype=True)
     client = TestClient(app.app)
-    data = '0,0,0'
+    data = (0,0)
 
     with pytest.raises(TypeError):
         predict(endpoint=client, data=data)