feat: Update /ai/extract_structured response schema (box/box-codegen#641) (#459)

Author: box-sdk-build

.codegen.json

@@ -1 +1 @@
-{ "engineHash": "cf82faa", "specHash": "3dc3f1e", "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):

box_sdk_gen/managers/ai.py

@@ -42,7 +42,9 @@
 
 from box_sdk_gen.schemas.ai_extract import AiExtract
 
-from box_sdk_gen.schemas.ai_extract_response import AiExtractResponse
+from box_sdk_gen.schemas.ai_extract_structured_response import (
+    AiExtractStructuredResponse,
+)
 
 from box_sdk_gen.schemas.ai_extract_structured import AiExtractStructured
 
@@ -411,7 +413,7 @@ def create_ai_extract_structured(
         fields: Optional[List[CreateAiExtractStructuredFields]] = None,
         ai_agent: Optional[AiAgentExtractStructured] = None,
         extra_headers: Optional[Dict[str, Optional[str]]] = None
-    ) -> AiExtractResponse:
+    ) -> AiExtractStructuredResponse:
         """
                 Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as a set of key-value pairs.
 
@@ -463,4 +465,4 @@ def create_ai_extract_structured(
                 network_session=self.network_session,
             )
         )
-        return deserialize(response.data, AiExtractResponse)
+        return deserialize(response.data, AiExtractStructuredResponse)

box_sdk_gen/schemas/__init__.py

@@ -10,6 +10,8 @@
 
 from box_sdk_gen.schemas.ai_extract_response import *
 
+from box_sdk_gen.schemas.ai_extract_structured_response import *
+
 from box_sdk_gen.schemas.ai_item_base import *
 
 from box_sdk_gen.schemas.ai_llm_endpoint_params_aws import *

box_sdk_gen/schemas/ai_extract_structured_response.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from box_sdk_gen.internal.base_object import BaseObject
+
+from box_sdk_gen.schemas.ai_extract_response import AiExtractResponse
+
+from box_sdk_gen.schemas.ai_agent_info import AiAgentInfo
+
+from box_sdk_gen.box.errors import BoxSDKError
+
+from box_sdk_gen.internal.utils import DateTime
+
+
+class AiExtractStructuredResponse(BaseObject):
+    def __init__(
+        self,
+        answer: AiExtractResponse,
+        created_at: DateTime,
+        *,
+        completion_reason: Optional[str] = None,
+        ai_agent_info: Optional[AiAgentInfo] = None,
+        **kwargs
+    ):
+        """
+        :param created_at: The ISO date formatted timestamp of when the answer to the prompt was created.
+        :type created_at: DateTime
+        :param completion_reason: The reason the response finishes., defaults to None
+        :type completion_reason: Optional[str], optional
+        """
+        super().__init__(**kwargs)
+        self.answer = answer
+        self.created_at = created_at
+        self.completion_reason = completion_reason
+        self.ai_agent_info = ai_agent_info

docs/ai.md

@@ -233,6 +233,6 @@ client.ai.create_ai_extract_structured(
 
 ### Returns
 
-This function returns a value of type `AiExtractResponse`.
+This function returns a value of type `AiExtractStructuredResponse`.
 
 A successful response including the answer from the LLM.

test/ai.py

@@ -30,7 +30,9 @@
 
 from box_sdk_gen.managers.uploads import UploadFileAttributesParentField
 
-from box_sdk_gen.schemas.ai_extract_response import AiExtractResponse
+from box_sdk_gen.schemas.ai_extract_structured_response import (
+    AiExtractStructuredResponse,
+)
 
 from box_sdk_gen.managers.ai import CreateAiExtractStructuredFields
 
@@ -266,7 +268,7 @@ def testAIExtractStructuredWithFields():
     )
     file: FileFull = uploaded_files.entries[0]
     delay_in_seconds(5)
-    response: AiExtractResponse = client.ai.create_ai_extract_structured(
+    response: AiExtractStructuredResponse = client.ai.create_ai_extract_structured(
         [AiItemBase(id=file.id)],
         fields=[
             CreateAiExtractStructuredFields(
@@ -311,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, '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.lastName')) == 'Doe'
     )
+    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)
 
 
@@ -378,22 +386,28 @@ def testAIExtractStructuredWithMetadataTemplate():
             ),
         ],
     )
-    response: AiExtractResponse = client.ai.create_ai_extract_structured(
+    response: AiExtractStructuredResponse = client.ai.create_ai_extract_structured(
         [AiItemBase(id=file.id)],
         metadata_template=CreateAiExtractStructuredMetadataTemplate(
             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
     )