|
8 | 8 |
|
9 | 9 | from typing import Dict |
10 | 10 |
|
11 | | -from typing import Union |
12 | | - |
13 | 11 | from box_sdk_gen.serialization.json import serialize |
14 | 12 |
|
15 | 13 | from box_sdk_gen.internal.utils import to_string |
|
22 | 20 |
|
23 | 21 | from box_sdk_gen.schemas.ai_agent_reference import AiAgentReference |
24 | 22 |
|
| 23 | +from box_sdk_gen.schemas.ai_agent_ask import AiAgentAsk |
| 24 | + |
| 25 | +from box_sdk_gen.schemas.ai_ask_agent import AiAskAgent |
| 26 | + |
25 | 27 | from box_sdk_gen.networking.fetch_options import ResponseFormat |
26 | 28 |
|
| 29 | +from box_sdk_gen.schemas.ai_agent_text_gen import AiAgentTextGen |
| 30 | + |
| 31 | +from box_sdk_gen.schemas.ai_text_gen_agent import AiTextGenAgent |
| 32 | + |
| 33 | +from box_sdk_gen.schemas.ai_agent_extract import AiAgentExtract |
| 34 | + |
| 35 | +from box_sdk_gen.schemas.ai_agent_extract_structured import AiAgentExtractStructured |
| 36 | + |
27 | 37 | from box_sdk_gen.schemas.ai_item_base import AiItemBase |
28 | 38 |
|
| 39 | +from box_sdk_gen.schemas.ai_extract_agent import AiExtractAgent |
| 40 | + |
| 41 | +from box_sdk_gen.schemas.ai_extract_structured_agent import AiExtractStructuredAgent |
| 42 | + |
29 | 43 | from box_sdk_gen.schemas.ai_response_full import AiResponseFull |
30 | 44 |
|
31 | 45 | from box_sdk_gen.schemas.client_error import ClientError |
|
36 | 50 |
|
37 | 51 | from box_sdk_gen.schemas.ai_text_gen import AiTextGen |
38 | 52 |
|
39 | | -from box_sdk_gen.schemas.ai_agent_ask import AiAgentAsk |
40 | | - |
41 | | -from box_sdk_gen.schemas.ai_agent_text_gen import AiAgentTextGen |
42 | | - |
43 | | -from box_sdk_gen.schemas.ai_agent_extract import AiAgentExtract |
44 | | - |
45 | | -from box_sdk_gen.schemas.ai_agent_extract_structured import AiAgentExtractStructured |
| 53 | +from box_sdk_gen.schemas.ai_agent import AiAgent |
46 | 54 |
|
47 | 55 | from box_sdk_gen.schemas.ai_extract import AiExtract |
48 | 56 |
|
@@ -222,7 +230,7 @@ def create_ai_ask( |
222 | 230 | *, |
223 | 231 | dialogue_history: Optional[List[AiDialogueHistory]] = None, |
224 | 232 | include_citations: Optional[bool] = None, |
225 | | - ai_agent: Optional[Union[AiAgentAsk, AiAgentReference]] = None, |
| 233 | + ai_agent: Optional[AiAskAgent] = None, |
226 | 234 | extra_headers: Optional[Dict[str, Optional[str]]] = None |
227 | 235 | ) -> Optional[AiResponseFull]: |
228 | 236 | """ |
@@ -280,7 +288,7 @@ def create_ai_text_gen( |
280 | 288 | items: List[CreateAiTextGenItems], |
281 | 289 | *, |
282 | 290 | dialogue_history: Optional[List[AiDialogueHistory]] = None, |
283 | | - ai_agent: Optional[Union[AiAgentReference, AiAgentTextGen]] = None, |
| 291 | + ai_agent: Optional[AiTextGenAgent] = None, |
284 | 292 | extra_headers: Optional[Dict[str, Optional[str]]] = None |
285 | 293 | ) -> AiResponse: |
286 | 294 | """ |
@@ -330,7 +338,7 @@ def get_ai_agent_default_config( |
330 | 338 | language: Optional[str] = None, |
331 | 339 | model: Optional[str] = None, |
332 | 340 | extra_headers: Optional[Dict[str, Optional[str]]] = None |
333 | | - ) -> Union[AiAgentAsk, AiAgentTextGen, AiAgentExtract, AiAgentExtractStructured]: |
| 341 | + ) -> AiAgent: |
334 | 342 | """ |
335 | 343 | Get the AI agent default config. |
336 | 344 | :param mode: The mode to filter the agent config to return. |
@@ -366,17 +374,14 @@ def get_ai_agent_default_config( |
366 | 374 | network_session=self.network_session, |
367 | 375 | ) |
368 | 376 | ) |
369 | | - return deserialize( |
370 | | - response.data, |
371 | | - Union[AiAgentAsk, AiAgentTextGen, AiAgentExtract, AiAgentExtractStructured], |
372 | | - ) |
| 377 | + return deserialize(response.data, AiAgent) |
373 | 378 |
|
374 | 379 | def create_ai_extract( |
375 | 380 | self, |
376 | 381 | prompt: str, |
377 | 382 | items: List[AiItemBase], |
378 | 383 | *, |
379 | | - ai_agent: Optional[Union[AiAgentReference, AiAgentExtract]] = None, |
| 384 | + ai_agent: Optional[AiExtractAgent] = None, |
380 | 385 | extra_headers: Optional[Dict[str, Optional[str]]] = None |
381 | 386 | ) -> AiResponse: |
382 | 387 | """ |
@@ -420,7 +425,7 @@ def create_ai_extract_structured( |
420 | 425 | *, |
421 | 426 | metadata_template: Optional[CreateAiExtractStructuredMetadataTemplate] = None, |
422 | 427 | fields: Optional[List[CreateAiExtractStructuredFields]] = None, |
423 | | - ai_agent: Optional[Union[AiAgentReference, AiAgentExtractStructured]] = None, |
| 428 | + ai_agent: Optional[AiExtractStructuredAgent] = None, |
424 | 429 | extra_headers: Optional[Dict[str, Optional[str]]] = None |
425 | 430 | ) -> AiExtractStructuredResponse: |
426 | 431 | """ |
|
0 commit comments