|
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