From 3776dc3d44bc09eb68da99f45e36e058dca2607e Mon Sep 17 00:00:00 2001
From: box-sdk-build <94016436+box-sdk-build@users.noreply.github.com>
Date: Wed, 18 Sep 2024 14:00:39 +0200
Subject: [PATCH] feat: Introduce `raw_data` field for storing raw `json`
 response (box/box-codegen#566) (#319)

---
 .codegen.json                                |  2 +-
 box_sdk_gen/internal/base_object.py          | 12 +++++++++
 box_sdk_gen/internal/utils.py                |  4 +++
 box_sdk_gen/serialization/json/serializer.py |  5 ++--
 test/ai.py                                   | 26 +++++++++++---------
 5 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/.codegen.json b/.codegen.json
index 87ffcbc1..9e8c6154 100644
--- a/.codegen.json
+++ b/.codegen.json
@@ -1 +1 @@
-{ "engineHash": "0338480", "specHash": "d9f90b8", "version": "1.4.1" }
+{ "engineHash": "5ceb679", "specHash": "d9f90b8", "version": "1.4.1" }
diff --git a/box_sdk_gen/internal/base_object.py b/box_sdk_gen/internal/base_object.py
index be247e74..47c40732 100644
--- a/box_sdk_gen/internal/base_object.py
+++ b/box_sdk_gen/internal/base_object.py
@@ -10,6 +10,7 @@ class BaseObject:
     _fields_to_json_mapping = {}
 
     def __init__(self, **kwargs):
+        self._raw_data: dict = {}
         self.__dict__.update(kwargs)
 
     @classmethod
@@ -23,9 +24,20 @@ def from_dict(cls, data: dict):
             )
         return cls(**unpacked_attributes)
 
+    @property
+    def raw_data(self):
+        """
+        Returns the raw json representation returned by the API
+        :return: dict with the raw json data
+        """
+        return self._raw_data
+
     def to_dict(self) -> dict:
         result_dict = {}
         for k, v in vars(self).items():
+            # Skip private and protected attributes
+            if k.startswith("_"):
+                continue
             if v is None:
                 continue
             if isinstance(v, NullValue):
diff --git a/box_sdk_gen/internal/utils.py b/box_sdk_gen/internal/utils.py
index 404dc33b..f79c6fee 100644
--- a/box_sdk_gen/internal/utils.py
+++ b/box_sdk_gen/internal/utils.py
@@ -237,6 +237,10 @@ def get_epoch_time_in_seconds() -> int:
     return int(time.time())
 
 
+def get_value_from_object_raw_data(obj: BaseObject, key: str) -> Any:
+    return obj.raw_data.get(key, None)
+
+
 class JwtAlgorithm(str, Enum):
     HS256 = 'HS256'
     HS384 = 'HS384'
diff --git a/box_sdk_gen/serialization/json/serializer.py b/box_sdk_gen/serialization/json/serializer.py
index 69aafb6b..408d20f7 100644
--- a/box_sdk_gen/serialization/json/serializer.py
+++ b/box_sdk_gen/serialization/json/serializer.py
@@ -20,5 +20,6 @@ def serialize(obj: Union[BaseObject, dict, list]) -> SerializedData:
 def deserialize(value: SerializedData, type: Type[BaseObject]):
     if get_origin(type) == Union:
         type = BaseObject._deserialize_union('', value, type)
-
-    return type.from_dict(value)
+    obj = type.from_dict(value)
+    obj._raw_data = value
+    return obj
diff --git a/test/ai.py b/test/ai.py
index 8650c966..eb281feb 100644
--- a/test/ai.py
+++ b/test/ai.py
@@ -58,6 +58,8 @@
 
 from box_sdk_gen.internal.utils import date_time_to_string
 
+from box_sdk_gen.internal.utils import get_value_from_object_raw_data
+
 from test.commons import upload_new_file
 
 from box_sdk_gen.schemas.ai_agent_ask import AiAgentAsk
@@ -210,7 +212,7 @@ def testAIExtract():
         ),
     )
     file: FileFull = uploaded_files.entries[0]
-    delay_in_seconds(1)
+    delay_in_seconds(2)
     response: AiResponse = client.ai.create_ai_extract(
         'firstName, lastName, location, yearOfBirth, company',
         [AiItemBase(id=file.id)],
@@ -240,7 +242,7 @@ def testAIExtractStructuredWithFields():
         ),
     )
     file: FileFull = uploaded_files.entries[0]
-    delay_in_seconds(1)
+    delay_in_seconds(2)
     response: AiExtractResponse = client.ai.create_ai_extract_structured(
         [AiItemBase(id=file.id)],
         fields=[
@@ -282,11 +284,11 @@ def testAIExtractStructuredWithFields():
         ],
         ai_agent=ai_extract_structured_agent_config,
     )
-    assert response.firstName == 'John'
-    assert response.lastName == 'Doe'
-    assert response.dateOfBirth == '1990-07-04'
-    assert response.age == 34
-    assert response.hobby == ['guitar', 'books']
+    assert get_value_from_object_raw_data(response, 'firstName') == 'John'
+    assert get_value_from_object_raw_data(response, 'lastName') == 'Doe'
+    assert get_value_from_object_raw_data(response, 'dateOfBirth') == '1990-07-04'
+    assert get_value_from_object_raw_data(response, 'age') == 34
+    assert get_value_from_object_raw_data(response, 'hobby') == ['guitar', 'books']
     client.files.delete_file_by_id(file.id)
 
 
@@ -345,11 +347,11 @@ def testAIExtractStructuredWithMetadataTemplate():
             template_key=template_key, scope='enterprise'
         ),
     )
-    assert response.firstName == 'John'
-    assert response.lastName == 'Doe'
-    assert response.dateOfBirth == '1990-07-04'
-    assert response.age == 34
-    assert response.hobby == ['guitar', 'books']
+    assert get_value_from_object_raw_data(response, 'firstName') == 'John'
+    assert get_value_from_object_raw_data(response, 'lastName') == 'Doe'
+    assert get_value_from_object_raw_data(response, 'dateOfBirth') == '1990-07-04'
+    assert get_value_from_object_raw_data(response, 'age') == 34
+    assert get_value_from_object_raw_data(response, 'hobby') == ['guitar', 'books']
     client.metadata_templates.delete_metadata_template(
         DeleteMetadataTemplateScope.ENTERPRISE.value, template.template_key
     )