feat: Support ai/extract and ai/extract_structured endpoints (box…

…/box-codegen#564) (#317)

.codegen.json

@@ -1 +1 @@
-{ "engineHash": "5b7aecf", "specHash": "b21666d", "version": "1.4.1" }
+{ "engineHash": "0338480", "specHash": "d9f90b8", "version": "1.4.1" }

box_sdk_gen/internal/utils.py

@@ -3,7 +3,7 @@
 import hashlib
 import os
 import uuid
-from time import time
+import time
 from enum import Enum
 from io import SEEK_CUR, SEEK_END, SEEK_SET, BufferedIOBase, BytesIO
 from typing import Any, Callable, Dict, Iterable, Optional, TypeVar
@@ -131,6 +131,10 @@ def decode_base_64_byte_stream(value: str) -> ByteStream:
  return BytesIO(base64.b64decode(value))
 
 
+def string_to_byte_stream(value: str) -> ByteStream:
+ return BytesIO(bytes(value, 'utf-8'))
+
+
 def read_byte_stream(byte_stream: ByteStream) -> Buffer:
  return Buffer(byte_stream.read())
 
@@ -230,7 +234,7 @@ def is_browser() -> bool:
 
 
 def get_epoch_time_in_seconds() -> int:
- return int(time())
+ return int(time.time())
 
 
 class JwtAlgorithm(str, Enum):
@@ -339,5 +343,9 @@ def date_time_from_string(date_time: str) -> DateTime:
  return DateTime.fromisoformat(date_time.replace('Z', '+00:00'))
 
 
+def delay_in_seconds(seconds: int):
+ time.sleep(seconds)
+
+
 def create_null():
  return null

box_sdk_gen/managers/ai.py

@@ -16,6 +16,8 @@
 
 from typing import Union
 
+from box_sdk_gen.schemas.ai_item_base import AiItemBase
+
 from box_sdk_gen.schemas.ai_dialogue_history import AiDialogueHistory
 
 from box_sdk_gen.schemas.ai_response_full import AiResponseFull
@@ -32,6 +34,16 @@
 
 from box_sdk_gen.schemas.ai_agent_text_gen import AiAgentTextGen
 
+from box_sdk_gen.schemas.ai_agent_extract import AiAgentExtract
+
+from box_sdk_gen.schemas.ai_agent_extract_structured import AiAgentExtractStructured
+
+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 import AiExtractStructured
+
 from box_sdk_gen.networking.auth import Authentication
 
 from box_sdk_gen.networking.network import NetworkSession
@@ -58,27 +70,27 @@ class CreateAiAskMode(str, Enum):
  SINGLE_ITEM_QA = 'single_item_qa'
 
 
-class CreateAiAskItemsTypeField(str, Enum):
+class CreateAiTextGenItemsTypeField(str, Enum):
  FILE = 'file'
 
 
-class CreateAiAskItems(BaseObject):
+class CreateAiTextGenItems(BaseObject):
  _discriminator = 'type', {'file'}
 
  def __init__(
  self,
  id: str,
  *,
- type: CreateAiAskItemsTypeField = CreateAiAskItemsTypeField.FILE.value,
+ type: CreateAiTextGenItemsTypeField = CreateAiTextGenItemsTypeField.FILE.value,
  content: Optional[str] = None,
  **kwargs
  ):
  """
- :param id: The id of the item.
+ :param id: The ID of the item.
  :type id: str
- :param type: The type of the item., defaults to CreateAiAskItemsTypeField.FILE.value
- :type type: CreateAiAskItemsTypeField, optional
- :param content: The content of the item, often the text representation., defaults to None
+ :param type: The type of the item., defaults to CreateAiTextGenItemsTypeField.FILE.value
+ :type type: CreateAiTextGenItemsTypeField, optional
+ :param content: The content to use as context for generating new text or editing existing text., defaults to None
  :type content: Optional[str], optional
  """
  super().__init__(**kwargs)
@@ -87,38 +99,99 @@ def __init__(
  self.content = content
 
 
-class CreateAiTextGenItemsTypeField(str, Enum):
- FILE = 'file'
+class GetAiAgentDefaultConfigMode(str, Enum):
+ ASK = 'ask'
+ TEXT_GEN = 'text_gen'
+ EXTRACT = 'extract'
+ EXTRACT_STRUCTURED = 'extract_structured'
 
 
-class CreateAiTextGenItems(BaseObject):
- _discriminator = 'type', {'file'}
+class CreateAiExtractStructuredMetadataTemplateTypeField(str, Enum):
+ METADATA_TEMPLATE = 'metadata_template'
+
+
+class CreateAiExtractStructuredMetadataTemplate(BaseObject):
+ _discriminator = 'type', {'metadata_template'}
 
  def __init__(
  self,
- id: str,
  *,
- type: CreateAiTextGenItemsTypeField = CreateAiTextGenItemsTypeField.FILE.value,
- content: Optional[str] = None,
+ template_key: Optional[str] = None,
+ type: Optional[CreateAiExtractStructuredMetadataTemplateTypeField] = None,
+ scope: Optional[str] = None,
  **kwargs
  ):
  """
- :param id: The ID of the item.
- :type id: str
- :param type: The type of the item., defaults to CreateAiTextGenItemsTypeField.FILE.value
- :type type: CreateAiTextGenItemsTypeField, optional
- :param content: The content to use as context for generating new text or editing existing text., defaults to None
- :type content: Optional[str], optional
+ :param template_key: The name of the metadata template., defaults to None
+ :type template_key: Optional[str], optional
+ :param type: Value is always `metadata_template`., defaults to None
+ :type type: Optional[CreateAiExtractStructuredMetadataTemplateTypeField], optional
+ :param scope: The scope of the metadata template that can either be global or
+ enterprise.
+ * The **global** scope is used for templates that are
+ available to any Box enterprise.
+ * The **enterprise** scope represents templates created within a specific enterprise,
+ containing the ID of that enterprise., defaults to None
+ :type scope: Optional[str], optional
  """
  super().__init__(**kwargs)
- self.id = id
+ self.template_key = template_key
  self.type = type
- self.content = content
+ self.scope = scope
 
 
-class GetAiAgentDefaultConfigMode(str, Enum):
- ASK = 'ask'
- TEXT_GEN = 'text_gen'
+class CreateAiExtractStructuredFieldsOptionsField(BaseObject):
+ def __init__(self, key: str, **kwargs):
+ """
+ :param key: A unique identifier for the field.
+ :type key: str
+ """
+ super().__init__(**kwargs)
+ self.key = key
+
+
+class CreateAiExtractStructuredFields(BaseObject):
+ _fields_to_json_mapping: Dict[str, str] = {
+ 'display_name': 'displayName',
+ **BaseObject._fields_to_json_mapping,
+ }
+ _json_to_fields_mapping: Dict[str, str] = {
+ 'displayName': 'display_name',
+ **BaseObject._json_to_fields_mapping,
+ }
+
+ def __init__(
+ self,
+ key: str,
+ *,
+ description: Optional[str] = None,
+ display_name: Optional[str] = None,
+ prompt: Optional[str] = None,
+ type: Optional[str] = None,
+ options: Optional[List[CreateAiExtractStructuredFieldsOptionsField]] = None,
+ **kwargs
+ ):
+ """
+ :param key: A unique identifier for the field.
+ :type key: str
+ :param description: A description of the field., defaults to None
+ :type description: Optional[str], optional
+ :param display_name: The display name of the field., defaults to None
+ :type display_name: Optional[str], optional
+ :param prompt: The context about the key that may include how to find and format it., defaults to None
+ :type prompt: Optional[str], optional
+ :param type: The type of the field. It include but is not limited to string, float, date, enum, and multiSelect., defaults to None
+ :type type: Optional[str], optional
+ :param options: A list of options for this field. This is most often used in combination with the enum and multiSelect field types., defaults to None
+ :type options: Optional[List[CreateAiExtractStructuredFieldsOptionsField]], optional
+ """
+ super().__init__(**kwargs)
+ self.key = key
+ self.description = description
+ self.display_name = display_name
+ self.prompt = prompt
+ self.type = type
+ self.options = options
 
 
 class AiManager:
@@ -137,7 +210,7 @@ def create_ai_ask(
  self,
  mode: CreateAiAskMode,
  prompt: str,
- items: List[CreateAiAskItems],
+ items: List[AiItemBase],
  *,
  dialogue_history: Optional[List[AiDialogueHistory]] = None,
  include_citations: Optional[bool] = None,
@@ -155,7 +228,7 @@ def create_ai_ask(
  **Note**: Box AI handles documents with text representations up to 1MB in size, or a maximum of 25 files, whichever comes first.
  If the file size exceeds 1MB, the first 1MB of text representation will be processed.
  If you set `mode` parameter to `single_item_qa`, the `items` array can have one element only.
- :type items: List[CreateAiAskItems]
+ :type items: List[AiItemBase]
  :param dialogue_history: The history of prompts and answers previously passed to the LLM. This provides additional context to the LLM in generating the response., defaults to None
  :type dialogue_history: Optional[List[AiDialogueHistory]], optional
  :param include_citations: A flag to indicate whether citations should be returned., defaults to None
@@ -198,7 +271,7 @@ def create_ai_text_gen(
  extra_headers: Optional[Dict[str, Optional[str]]] = None
  ) -> AiResponse:
  """
- Sends an AI request to supported LLMs and returns an answer specifically focused on the creation of new text.
+ Sends an AI request to supported Large Language Models (LLMs) and returns generated text based on the provided prompt.
  :param prompt: The prompt provided by the client to be answered by the LLM. The prompt's length is limited to 10000 characters.
  :type prompt: str
  :param items: The items to be processed by the LLM, often files.
@@ -244,7 +317,7 @@ def get_ai_agent_default_config(
  language: Optional[str] = None,
  model: Optional[str] = None,
  extra_headers: Optional[Dict[str, Optional[str]]] = None
- ) -> Union[AiAgentAsk, AiAgentTextGen]:
+ ) -> Union[AiAgentAsk, AiAgentTextGen, AiAgentExtract, AiAgentExtractStructured]:
  """
  Get the AI agent default config
  :param mode: The mode to filter the agent config to return.
@@ -280,4 +353,106 @@ def get_ai_agent_default_config(
  network_session=self.network_session,
  )
  )
- return deserialize(response.data, Union[AiAgentAsk, AiAgentTextGen])
+ return deserialize(
+ response.data,
+ Union[AiAgentAsk, AiAgentTextGen, AiAgentExtract, AiAgentExtractStructured],
+ )
+
+ def create_ai_extract(
+ self,
+ prompt: str,
+ items: List[AiItemBase],
+ *,
+ ai_agent: Optional[AiAgentExtract] = None,
+ extra_headers: Optional[Dict[str, Optional[str]]] = None
+ ) -> AiResponse:
+ """
+ Sends an AI request to supported Large Language Models (LLMs) and extracts metadata in form of key-value pairs.
+
+ Freeform metadata extraction does not require any metadata template setup before sending the request.
+
+ :param prompt: The prompt provided to a Large Language Model (LLM) in the request. The prompt can be up to 10000 characters long and it can be an XML or a JSON schema.
+ :type prompt: str
+ :param items: The items that LLM will process. Currently, you can use files only.
+ :type items: List[AiItemBase]
+ :param extra_headers: Extra headers that will be included in the HTTP request., defaults to None
+ :type extra_headers: Optional[Dict[str, Optional[str]]], optional
+ """
+ if extra_headers is None:
+ extra_headers = {}
+ request_body: Dict = {'prompt': prompt, 'items': items, 'ai_agent': ai_agent}
+ headers_map: Dict[str, str] = prepare_params({**extra_headers})
+ response: FetchResponse = fetch(
+ FetchOptions(
+ url=''.join(
+ [self.network_session.base_urls.base_url, '/2.0/ai/extract']
+ ),
+ method='POST',
+ headers=headers_map,
+ data=serialize(request_body),
+ content_type='application/json',
+ response_format='json',
+ auth=self.auth,
+ network_session=self.network_session,
+ )
+ )
+ return deserialize(response.data, AiResponse)
+
+ def create_ai_extract_structured(
+ self,
+ items: List[AiItemBase],
+ *,
+ metadata_template: Optional[CreateAiExtractStructuredMetadataTemplate] = None,
+ fields: Optional[List[CreateAiExtractStructuredFields]] = None,
+ ai_agent: Optional[AiAgentExtractStructured] = None,
+ extra_headers: Optional[Dict[str, Optional[str]]] = None
+ ) -> AiExtractResponse:
+ """
+ Sends an AI request to supported Large Language Models (LLMs) and returns extracted metadata as a set of key-value pairs.
+
+ For this request, you need to use an already defined metadata template or a define a schema yourself.
+
+
+ To learn more about creating templates, see [Creating metadata templates in the Admin Console](https://support.box.com/hc/en-us/articles/360044194033-Customizing-Metadata-Templates)
+
+
+ or use the [metadata template API](g://metadata/templates/create).
+
+ :param items: The items to be processed by the LLM. Currently you can use files only.
+ :type items: List[AiItemBase]
+ :param metadata_template: The metadata template containing the fields to extract.
+ For your request to work, you must provide either `metadata_template` or `fields`, but not both., defaults to None
+ :type metadata_template: Optional[CreateAiExtractStructuredMetadataTemplate], optional
+ :param fields: The fields to be extracted from the provided items.
+ For your request to work, you must provide either `metadata_template` or `fields`, but not both., defaults to None
+ :type fields: Optional[List[CreateAiExtractStructuredFields]], optional
+ :param extra_headers: Extra headers that will be included in the HTTP request., defaults to None
+ :type extra_headers: Optional[Dict[str, Optional[str]]], optional
+ """
+ if extra_headers is None:
+ extra_headers = {}
+ request_body: Dict = {
+ 'items': items,
+ 'metadata_template': metadata_template,
+ 'fields': fields,
+ 'ai_agent': ai_agent,
+ }
+ headers_map: Dict[str, str] = prepare_params({**extra_headers})
+ response: FetchResponse = fetch(
+ FetchOptions(
+ url=''.join(
+ [
+ self.network_session.base_urls.base_url,
+ '/2.0/ai/extract_structured',
+ ]
+ ),
+ method='POST',
+ headers=headers_map,
+ data=serialize(request_body),
+ content_type='application/json',
+ response_format='json',
+ auth=self.auth,
+ network_session=self.network_session,
+ )
+ )
+ return deserialize(response.data, AiExtractResponse)

box_sdk_gen/schemas/__init__.py

@@ -18,6 +18,8 @@
 
 from box_sdk_gen.schemas.ai_response import *
 
+from box_sdk_gen.schemas.ai_extract_response import *
+
 from box_sdk_gen.schemas.app_item import *
 
 from box_sdk_gen.schemas.classification import *
@@ -332,6 +334,8 @@
 
 from box_sdk_gen.schemas.zip_download_status import *
 
+from box_sdk_gen.schemas.ai_item_base import *
+
 from box_sdk_gen.schemas.ai_llm_endpoint_params_google import *
 
 from box_sdk_gen.schemas.ai_llm_endpoint_params_open_ai import *
@@ -350,6 +354,14 @@
 
 from box_sdk_gen.schemas.ai_agent_long_text_tool import *
 
+from box_sdk_gen.schemas.ai_agent_extract_structured import *
+
+from box_sdk_gen.schemas.ai_extract_structured import *
+
+from box_sdk_gen.schemas.ai_agent_extract import *
+
+from box_sdk_gen.schemas.ai_extract import *
+
 from box_sdk_gen.schemas.ai_agent_ask import *
 
 from box_sdk_gen.schemas.ai_citation import *