From 9ab96284a2bba274d039132c2bacf2bd7f6e656b Mon Sep 17 00:00:00 2001
From: John Harrison <john.harrison@autotrader.co.uk>
Date: Wed, 23 Aug 2023 17:38:48 +0100
Subject: [PATCH] Fix predictors for pydantic 2 API changes

Pydantic 2.x BaseModel has a slightly different API, and generates
slightly different JSON schemas.
https://docs.pydantic.dev/latest/migration/#changes-to-pydanticbasemodel
https://docs.pydantic.dev/latest/migration/#changes-to-json-schema-generation
---
 CHANGELOG.md                 |  2 ++
 fastapi_mlflow/predictors.py |  2 +-
 tests/test_predictors.py     | 12 ++++++------
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 262346b..2d7994b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Fixed
+- Adapt to pydantic 2.x API changes
 
 ## [0.6.1] - 2023-08-23
 ### Fixed
diff --git a/fastapi_mlflow/predictors.py b/fastapi_mlflow/predictors.py
index 6e21859..5c1ed77 100644
--- a/fastapi_mlflow/predictors.py
+++ b/fastapi_mlflow/predictors.py
@@ -55,7 +55,7 @@ class Response(BaseModel):
         data: List[output_model]
 
     def request_to_dataframe(request: Request) -> pd.DataFrame:
-        df = pd.DataFrame([row.dict() for row in request.data], dtype=object)
+        df = pd.DataFrame([row.model_dump() for row in request.data], dtype=object)
         for item in input_schema.to_dict():
             if item["type"] in ("integer", "int32"):
                 df[item["name"]] = df[item["name"]].astype(np.int32)
diff --git a/tests/test_predictors.py b/tests/test_predictors.py
index 1690102..f31495b 100644
--- a/tests/test_predictors.py
+++ b/tests/test_predictors.py
@@ -41,11 +41,11 @@ def test_predictor_has_correct_signature_for_input(
         "type in predictor function parameter `request` is not a"
         "subclass of pydantic.BaseModel"
     )
-    schema = request_type.schema()
+    schema = request_type.model_json_schema()
     assert "data" in schema["required"]
     assert schema["properties"]["data"]["type"] == "array"
-    assert "RequestRow" in schema["definitions"]
-    assert schema["definitions"]["RequestRow"]["required"] == list(model_input.columns)
+    assert "RequestRow" in schema["$defs"]
+    assert schema["$defs"]["RequestRow"]["required"] == list(model_input.columns)
 
 
 def test_predictor_signature_type_can_be_constructed(
@@ -90,10 +90,10 @@ def test_predictor_has_correct_return_type(
         "type in predictor function parameter `request` is not a"
         "subclass of pydantic.BaseModel"
     )
-    schema = return_type.schema()
+    schema = return_type.model_json_schema()
     assert "data" in schema["required"]
     assert schema["properties"]["data"]["type"] == "array"
-    assert "ResponseRow" in schema["definitions"]
+    assert "ResponseRow" in schema["$defs"]
 
 
 @pytest.mark.asyncio
@@ -119,7 +119,7 @@ async def test_predictor_correctly_applies_model(
     request_obj = request_type(data=model_input.to_dict(orient="records"))
     response = await predictor(request_obj)
     try:
-        assert response.data == model_output.to_dict(orient="records")  # type: ignore
+        assert [row.model_dump() for row in response.data] == model_output.to_dict(orient="records")  # type: ignore
     except (AttributeError, TypeError):
         predictions = [item.prediction for item in response.data]
         assert predictions == model_output.tolist()  # type: ignore