test(boxsdkgen): Expand tests for ai extract structured with new fields (box/box-codegen#952)

.codegen.json

@@ -1 +1 @@
-{ "engineHash": "a87c352", "specHash": "dd7f7a9", "version": "4.11.0" }
+{ "engineHash": "e392e8c", "specHash": "dd7f7a9", "version": "4.11.0" }

docs/box_sdk_gen/ai.md

@@ -167,7 +167,7 @@ See the endpoint docs at
 <!-- sample post_ai_extract_structured -->
 
 ```python
-client.ai.create_ai_extract_structured([AiItemBase(id=file.id)], fields=[CreateAiExtractStructuredFields(key='firstName', display_name='First name', description='Person first name', prompt='What is the your first name?', type='string'), CreateAiExtractStructuredFields(key='lastName', display_name='Last name', description='Person last name', prompt='What is the your last name?', type='string'), CreateAiExtractStructuredFields(key='dateOfBirth', display_name='Birth date', description='Person date of birth', prompt='What is the date of your birth?', type='date'), CreateAiExtractStructuredFields(key='age', display_name='Age', description='Person age', prompt='How old are you?', type='float'), CreateAiExtractStructuredFields(key='hobby', display_name='Hobby', description='Person hobby', prompt='What is your hobby?', type='multiSelect', options=[CreateAiExtractStructuredFieldsOptionsField(key='guitar'), CreateAiExtractStructuredFieldsOptionsField(key='books')])], ai_agent=ai_extract_structured_agent_basic_text_config, include_confidence_score=True)
+client.ai.create_ai_extract_structured([AiItemBase(id=file.id)], fields=[CreateAiExtractStructuredFields(key='firstName', display_name='First name', description='Person first name', prompt='What is the your first name?', type='string'), CreateAiExtractStructuredFields(key='lastName', display_name='Last name', description='Person last name', prompt='What is the your last name?', type='string'), CreateAiExtractStructuredFields(key='dateOfBirth', display_name='Birth date', description='Person date of birth', prompt='What is the date of your birth?', type='date'), CreateAiExtractStructuredFields(key='age', display_name='Age', description='Person age', prompt='How old are you?', type='float'), CreateAiExtractStructuredFields(key='hobby', display_name='Hobby', description='Person hobby', prompt='What is your hobby?', type='multiSelect', options=[CreateAiExtractStructuredFieldsOptionsField(key='guitar'), CreateAiExtractStructuredFieldsOptionsField(key='books')]), CreateAiExtractStructuredFields(key='address', display_name='Address', description='Person address', type='struct', prompt='Extract the full mailing address.', fields=[AiExtractSubField(key='street', display_name='Street', type='string'), AiExtractSubField(key='city', display_name='City', type='string'), AiExtractSubField(key='state', display_name='State', type='string'), AiExtractSubField(key='zip', display_name='Zip', type='string'), AiExtractSubField(key='country', display_name='Country', type='string')]), CreateAiExtractStructuredFields(key='work_history', display_name='Work history', description='Person work history', type='table', prompt='Extract each job as a row.', fields=[AiExtractSubField(key='job_title', display_name='Job title', type='string'), AiExtractSubField(key='company', display_name='Company', type='string'), AiExtractSubField(key='start_year', display_name='Start year', type='string'), AiExtractSubField(key='end_year', display_name='End year', type='string')])], ai_agent=ai_extract_structured_agent_basic_text_config, include_confidence_score=True, include_reference=True)
 ```
 
 ### Arguments

test/box_sdk_gen/test/ai.py

@@ -78,6 +78,8 @@
 
 from box_sdk_gen.schemas.ai_agent_extract_structured import AiAgentExtractStructured
 
+from box_sdk_gen.schemas.ai_extract_sub_field import AiExtractSubField
+
 client: BoxClient = get_default_client()
 
 
@@ -254,7 +256,7 @@ def testAIExtractStructuredWithFields():
         string_to_byte_stream(
             ''.join(
                 [
-                    'My name is John Doe. I was born in 4th July 1990. I am 34 years old. My hobby is guitar. My UUID is ',
+                    'My name is John Doe. I was born in 4th July 1990. I am 34 years old. My hobby is guitar. I live at 900 Jefferson Ave, Redwood City, CA 94063, US. My work history: Software Engineer at Box from 2020 to 2024. My UUID is ',
                     get_uuid(),
                 ]
             )
@@ -304,16 +306,62 @@ def testAIExtractStructuredWithFields():
                     CreateAiExtractStructuredFieldsOptionsField(key='books'),
                 ],
             ),
+            CreateAiExtractStructuredFields(
+                key='address',
+                display_name='Address',
+                description='Person address',
+                type='struct',
+                prompt='Extract the full mailing address.',
+                fields=[
+                    AiExtractSubField(
+                        key='street', display_name='Street', type='string'
+                    ),
+                    AiExtractSubField(key='city', display_name='City', type='string'),
+                    AiExtractSubField(key='state', display_name='State', type='string'),
+                    AiExtractSubField(key='zip', display_name='Zip', type='string'),
+                    AiExtractSubField(
+                        key='country', display_name='Country', type='string'
+                    ),
+                ],
+            ),
+            CreateAiExtractStructuredFields(
+                key='work_history',
+                display_name='Work history',
+                description='Person work history',
+                type='table',
+                prompt='Extract each job as a row.',
+                fields=[
+                    AiExtractSubField(
+                        key='job_title', display_name='Job title', type='string'
+                    ),
+                    AiExtractSubField(
+                        key='company', display_name='Company', type='string'
+                    ),
+                    AiExtractSubField(
+                        key='start_year', display_name='Start year', type='string'
+                    ),
+                    AiExtractSubField(
+                        key='end_year', display_name='End year', type='string'
+                    ),
+                ],
+            ),
         ],
         ai_agent=ai_extract_structured_agent_basic_text_config,
         include_confidence_score=True,
+        include_reference=True,
     )
     assert not response.confidence_score == None
+    assert not response.reference == None
+    assert not response.ai_agent_info == None
     assert to_string(response.answer.get('hobby')) == to_string(['guitar'])
     assert to_string(response.answer.get('firstName')) == 'John'
     assert to_string(response.answer.get('lastName')) == 'Doe'
     assert to_string(response.answer.get('dateOfBirth')) == '1990-07-04'
     assert to_string(response.answer.get('age')) == '34'
+    assert 'Redwood City' in to_string(response.answer.get('address'))
+    assert 'CA' in to_string(response.answer.get('address'))
+    assert '94063' in to_string(response.answer.get('address'))
+    assert 'Box' in to_string(response.answer.get('work_history'))
     assert response.completion_reason == 'done'
     client.files.delete_file_by_id(file.id)