test: Change tests for new response in ai_extract_structured (box/b…

…ox-codegen#641)

.codegen.json

@@ -1 +1 @@
-{ "engineHash": "cf82faa", "specHash": "1fdcbef", "version": "1.10.0" }
+{ "engineHash": "a74691d", "specHash": "1fdcbef", "version": "1.10.0" }

box_sdk_gen/internal/utils.py

@@ -263,7 +263,11 @@ def get_epoch_time_in_seconds() -> int:
 
 
 def get_value_from_object_raw_data(obj: BaseObject, key: str) -> Any:
-    return obj.raw_data.get(key, None)
+    keys = key.split('.')
+    value: dict = obj.raw_data
+    for k in keys:
+        value = value.get(k, {})
+    return value
 
 
 class JwtAlgorithm(str, Enum):

test/ai.py

@@ -313,16 +313,22 @@ def testAIExtractStructuredWithFields():
         ],
         ai_agent=agent_ignoring_overriding_embeddings_model,
     )
-    assert to_string(get_value_from_object_raw_data(response, 'firstName')) == 'John'
-    assert to_string(get_value_from_object_raw_data(response, 'lastName')) == 'Doe'
+    assert to_string(
+        get_value_from_object_raw_data(response, 'answer.hobby')
+    ) == to_string(['guitar'])
     assert (
-        to_string(get_value_from_object_raw_data(response, 'dateOfBirth'))
-        == '1990-07-04'
+        to_string(get_value_from_object_raw_data(response, 'answer.firstName'))
+        == 'John'
+    )
+    assert (
+        to_string(get_value_from_object_raw_data(response, 'answer.lastName')) == 'Doe'
     )
-    assert to_string(get_value_from_object_raw_data(response, 'age')) == '34'
-    assert to_string(get_value_from_object_raw_data(response, 'hobby')) == to_string(
-        ['guitar']
+    assert (
+        to_string(get_value_from_object_raw_data(response, 'answer.dateOfBirth'))
+        == '1990-07-04'
     )
+    assert to_string(get_value_from_object_raw_data(response, 'answer.age')) == '34'
+    assert response.completion_reason == 'done'
     client.files.delete_file_by_id(file.id)
 
 
@@ -386,16 +392,22 @@ def testAIExtractStructuredWithMetadataTemplate():
             template_key=template_key, scope='enterprise'
         ),
     )
-    assert to_string(get_value_from_object_raw_data(response, 'firstName')) == 'John'
-    assert to_string(get_value_from_object_raw_data(response, 'lastName')) == 'Doe'
     assert (
-        to_string(get_value_from_object_raw_data(response, 'dateOfBirth'))
-        == '1990-07-04T00:00:00Z'
+        to_string(get_value_from_object_raw_data(response, 'answer.firstName'))
+        == 'John'
+    )
+    assert (
+        to_string(get_value_from_object_raw_data(response, 'answer.lastName')) == 'Doe'
     )
-    assert to_string(get_value_from_object_raw_data(response, 'age')) == '34'
-    assert to_string(get_value_from_object_raw_data(response, 'hobby')) == to_string(
-        ['guitar']
+    assert (
+        to_string(get_value_from_object_raw_data(response, 'answer.dateOfBirth'))
+        == '1990-07-04T00:00:00Z'
     )
+    assert to_string(get_value_from_object_raw_data(response, 'answer.age')) == '34'
+    assert to_string(
+        get_value_from_object_raw_data(response, 'answer.hobby')
+    ) == to_string(['guitar'])
+    assert response.completion_reason == 'done'
     client.metadata_templates.delete_metadata_template(
         DeleteMetadataTemplateScope.ENTERPRISE, template.template_key
     )