From b1318703150e3fb07d49632ab3e29d405341a5cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ianar=C3=A9=20S=C3=A9vi?= <ianare@mindee.co>
Date: Fri, 18 Jul 2025 17:18:59 +0200
Subject: [PATCH] :sparkles: add generic response loader

---
 mindee/client_v2.py            | 14 --------------
 mindee/input/local_response.py | 16 +++++++++++++++-
 tests/test_client_v2.py        |  8 +++++---
 3 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/mindee/client_v2.py b/mindee/client_v2.py
index 0b757e50..288ec0e3 100644
--- a/mindee/client_v2.py
+++ b/mindee/client_v2.py
@@ -5,7 +5,6 @@
 from mindee.error.mindee_error import MindeeError
 from mindee.error.mindee_http_error_v2 import handle_error_v2
 from mindee.input.inference_parameters import InferenceParameters
-from mindee.input.local_response import LocalResponse
 from mindee.input.polling_options import PollingOptions
 from mindee.input.sources.local_input_source import LocalInputSource
 from mindee.logger import logger
@@ -137,16 +136,3 @@ def enqueue_and_get_inference(
             sleep(params.polling_options.delay_sec)
 
         raise MindeeError(f"Couldn't retrieve document after {try_counter + 1} tries.")
-
-    @staticmethod
-    def load_inference(local_response: LocalResponse) -> InferenceResponse:
-        """
-        Load a prediction from the V2 API.
-
-        :param local_response: Local response to load.
-        :return: A valid prediction.
-        """
-        try:
-            return InferenceResponse(local_response.as_dict)
-        except KeyError as exc:
-            raise MindeeError("No prediction found in local response.") from exc
diff --git a/mindee/input/local_response.py b/mindee/input/local_response.py
index bb57420c..c8a27fbb 100644
--- a/mindee/input/local_response.py
+++ b/mindee/input/local_response.py
@@ -4,9 +4,10 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, BinaryIO, Dict, Union
+from typing import Any, BinaryIO, Dict, Type, TypeVar, Union
 
 from mindee.error.mindee_error import MindeeError
+from mindee.parsing.v2.common_response import CommonResponse
 
 
 class LocalResponse:
@@ -102,3 +103,16 @@ def is_valid_hmac_signature(
         :return: True if the HMAC signature is valid.
         """
         return signature == self.get_hmac_signature(secret_key)
+
+    ResponseT = TypeVar("ResponseT", bound=CommonResponse)
+
+    def deserialize_response(self, response_class: Type[ResponseT]) -> ResponseT:
+        """
+        Load a local inference.
+
+        Typically used when wanting to load a V2 webhook callback.
+        """
+        try:
+            return response_class(self.as_dict)
+        except KeyError as exc:
+            raise MindeeError("Invalid class specified for deserialization.") from exc
diff --git a/tests/test_client_v2.py b/tests/test_client_v2.py
index 22b36bbb..f70165bc 100644
--- a/tests/test_client_v2.py
+++ b/tests/test_client_v2.py
@@ -3,7 +3,7 @@
 import pytest
 
 from mindee import ClientV2, InferenceParameters, InferenceResponse, LocalResponse
-from mindee.error.mindee_error import MindeeApiV2Error
+from mindee.error.mindee_error import MindeeApiV2Error, MindeeError
 from mindee.error.mindee_http_error_v2 import MindeeHTTPErrorV2
 from mindee.input import LocalInputSource, PathInput
 from mindee.mindee_http.base_settings import USER_AGENT
@@ -140,12 +140,14 @@ def _assert_findoc_inference(response: InferenceResponse):
 
 
 @pytest.mark.v2
-def test_loads_from_prediction(env_client):
+def test_loads_from_prediction():
     input_inference = LocalResponse(
         V2_DATA_DIR / "products" / "financial_document" / "complete.json"
     )
-    response = env_client.load_inference(input_inference)
+    response = input_inference.deserialize_response(InferenceResponse)
     _assert_findoc_inference(response)
+    with pytest.raises(MindeeError):
+        input_inference.deserialize_response(JobResponse)
 
 
 @pytest.mark.v2