84
84
85
85
from box_sdk_gen .schemas .ai_agent_extract_structured import AiAgentExtractStructured
86
86
87
+ from box_sdk_gen .schemas .ai_agent_long_text_tool import AiAgentLongTextTool
88
+
87
89
client : BoxClient = get_default_client ()
88
90
89
91
@@ -216,6 +218,17 @@ def testAIExtract():
216
218
ai_extract_agent_config : AiAgentExtract = client .ai .get_ai_agent_default_config (
217
219
GetAiAgentDefaultConfigMode .EXTRACT , language = 'en-US'
218
220
)
221
+ long_text_config_with_no_embeddings : AiAgentLongTextTool = AiAgentLongTextTool (
222
+ system_message = ai_extract_agent_config .long_text .system_message ,
223
+ prompt_template = ai_extract_agent_config .long_text .prompt_template ,
224
+ model = ai_extract_agent_config .long_text .model ,
225
+ num_tokens_for_completion = ai_extract_agent_config .long_text .num_tokens_for_completion ,
226
+ llm_endpoint_params = ai_extract_agent_config .long_text .llm_endpoint_params ,
227
+ )
228
+ agent_ignoring_overriding_embeddings_model : AiAgentExtract = AiAgentExtract (
229
+ basic_text = ai_extract_agent_config .basic_text ,
230
+ long_text = long_text_config_with_no_embeddings ,
231
+ )
219
232
uploaded_files : Files = client .uploads .upload_file (
220
233
UploadFileAttributes (
221
234
name = '' .join ([get_uuid (), '.txt' ]),
@@ -230,7 +243,7 @@ def testAIExtract():
230
243
response : AiResponse = client .ai .create_ai_extract (
231
244
'firstName, lastName, location, yearOfBirth, company' ,
232
245
[AiItemBase (id = file .id )],
233
- ai_agent = ai_extract_agent_config ,
246
+ ai_agent = agent_ignoring_overriding_embeddings_model ,
234
247
)
235
248
expected_response : str = (
236
249
'{"firstName": "John", "lastName": "Doe", "location": "San Francisco", "yearOfBirth": "1990", "company": "Box"}'
@@ -246,14 +259,12 @@ def testAIExtractStructuredWithFields():
246
259
GetAiAgentDefaultConfigMode .EXTRACT_STRUCTURED , language = 'en-US'
247
260
)
248
261
)
249
- long_text_config_with_no_embeddings : AiAgentExtractStructured = (
250
- AiAgentExtractStructured (
251
- system_message = ai_extract_structured_agent_config .long_text .system_message ,
252
- prompt_template = ai_extract_structured_agent_config .long_text .prompt_template ,
253
- model = ai_extract_structured_agent_config .long_text .model ,
254
- num_tokens_for_completion = ai_extract_structured_agent_config .long_text .num_tokens_for_completion ,
255
- llm_endpoint_params = ai_extract_structured_agent_config .long_text .llm_endpoint_params ,
256
- )
262
+ long_text_config_with_no_embeddings : AiAgentLongTextTool = AiAgentLongTextTool (
263
+ system_message = ai_extract_structured_agent_config .long_text .system_message ,
264
+ prompt_template = ai_extract_structured_agent_config .long_text .prompt_template ,
265
+ model = ai_extract_structured_agent_config .long_text .model ,
266
+ num_tokens_for_completion = ai_extract_structured_agent_config .long_text .num_tokens_for_completion ,
267
+ llm_endpoint_params = ai_extract_structured_agent_config .long_text .llm_endpoint_params ,
257
268
)
258
269
agent_ignoring_overriding_embeddings_model : AiAgentExtractStructured = (
259
270
AiAgentExtractStructured (
0 commit comments