llama-cloud 0.1.18__py3-none-any.whl → 0.1.19__py3-none-any.whl

llama_cloud/resources/llama_extract/client.py

@@ -10,18 +10,19 @@ from ...core.jsonable_encoder import jsonable_encoder
 from ...core.remove_none_from_dict import remove_none_from_dict
 from ...errors.unprocessable_entity_error import UnprocessableEntityError
 from ...types.extract_agent import ExtractAgent
-from ...types.extract_agent_create import ExtractAgentCreate
-from ...types.extract_agent_update import ExtractAgentUpdate
+from ...types.extract_config import ExtractConfig
 from ...types.extract_job import ExtractJob
 from ...types.extract_job_create import ExtractJobCreate
-from ...types.extract_job_create_batch import ExtractJobCreateBatch
 from ...types.extract_resultset import ExtractResultset
 from ...types.extract_run import ExtractRun
-from ...types.extract_schema_validate_request import ExtractSchemaValidateRequest
 from ...types.extract_schema_validate_response import ExtractSchemaValidateResponse
 from ...types.http_validation_error import HttpValidationError
 from ...types.llama_extract_settings import LlamaExtractSettings
 from ...types.paginated_extract_runs_response import PaginatedExtractRunsResponse
+from .types.extract_agent_create_data_schema import ExtractAgentCreateDataSchema
+from .types.extract_agent_update_data_schema import ExtractAgentUpdateDataSchema
+from .types.extract_job_create_batch_data_schema_override import ExtractJobCreateBatchDataSchemaOverride
+from .types.extract_schema_validate_request_data_schema import ExtractSchemaValidateRequestDataSchema
 
 try:
     import pydantic
@@ -39,17 +40,19 @@ class LlamaExtractClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
         self._client_wrapper = client_wrapper
 
-    def list_extraction_agents(self, *, project_id: typing.Optional[str] = None) -> typing.List[ExtractAgent]:
+    def list_extraction_agents(self, *, project_id: str) -> typing.List[ExtractAgent]:
         """
         Parameters:
-            - project_id: typing.Optional[str].
+            - project_id: str.
         ---
         from llama_cloud.client import LlamaCloud
 
         client = LlamaCloud(
             token="YOUR_TOKEN",
         )
-        client.llama_extract.list_extraction_agents()
+        client.llama_extract.list_extraction_agents(
+            project_id="string",
+        )
         """
         _response = self._client_wrapper.httpx_client.request(
             "GET",
@@ -73,7 +76,9 @@ class LlamaExtractClient:
         *,
         project_id: typing.Optional[str] = None,
         organization_id: typing.Optional[str] = None,
-        request: ExtractAgentCreate,
+        name: str,
+        data_schema: ExtractAgentCreateDataSchema,
+        config: ExtractConfig,
     ) -> ExtractAgent:
         """
         Parameters:
@@ -81,26 +86,23 @@ class LlamaExtractClient:
 
             - organization_id: typing.Optional[str].
 
-            - request: ExtractAgentCreate.
+            - name: str. The name of the extraction schema
+
+            - data_schema: ExtractAgentCreateDataSchema. The schema of the data.
+
+            - config: ExtractConfig. The configuration parameters for the extraction agent.
         ---
-        from llama_cloud import (
-            ExtractAgentCreate,
-            ExtractConfig,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import LlamaCloud
 
         client = LlamaCloud(
             token="YOUR_TOKEN",
         )
         client.llama_extract.create_extraction_agent(
-            request=ExtractAgentCreate(
-                name="string",
-                config=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            name="string",
+            config=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
@@ -108,7 +110,7 @@ class LlamaExtractClient:
             "POST",
             urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/extraction-agents"),
             params=remove_none_from_dict({"project_id": project_id, "organization_id": organization_id}),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"name": name, "data_schema": data_schema, "config": config}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -122,30 +124,29 @@ class LlamaExtractClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
-    def validate_extraction_schema(self, *, request: ExtractSchemaValidateRequest) -> ExtractSchemaValidateResponse:
+    def validate_extraction_schema(
+        self, *, data_schema: ExtractSchemaValidateRequestDataSchema
+    ) -> ExtractSchemaValidateResponse:
         """
         Validates an extraction agent's schema definition.
         Returns the normalized and validated schema if valid, otherwise raises an HTTP 400.
 
         Parameters:
-            - request: ExtractSchemaValidateRequest.
+            - data_schema: ExtractSchemaValidateRequestDataSchema.
         ---
-        from llama_cloud import ExtractSchemaValidateRequest
         from llama_cloud.client import LlamaCloud
 
         client = LlamaCloud(
             token="YOUR_TOKEN",
         )
-        client.llama_extract.validate_extraction_schema(
-            request=ExtractSchemaValidateRequest(),
-        )
+        client.llama_extract.validate_extraction_schema()
         """
         _response = self._client_wrapper.httpx_client.request(
             "POST",
             urllib.parse.urljoin(
                 f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/extraction-agents/schema/validation"
             ),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"data_schema": data_schema}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -226,19 +227,18 @@ class LlamaExtractClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
-    def update_extraction_agent(self, extraction_agent_id: str, *, request: ExtractAgentUpdate) -> ExtractAgent:
+    def update_extraction_agent(
+        self, extraction_agent_id: str, *, data_schema: ExtractAgentUpdateDataSchema, config: ExtractConfig
+    ) -> ExtractAgent:
         """
         Parameters:
             - extraction_agent_id: str.
 
-            - request: ExtractAgentUpdate.
+            - data_schema: ExtractAgentUpdateDataSchema. The schema of the data
+
+            - config: ExtractConfig. The configuration parameters for the extraction agent.
         ---
-        from llama_cloud import (
-            ExtractAgentUpdate,
-            ExtractConfig,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import LlamaCloud
 
         client = LlamaCloud(
@@ -246,11 +246,9 @@ class LlamaExtractClient:
         )
         client.llama_extract.update_extraction_agent(
             extraction_agent_id="string",
-            request=ExtractAgentUpdate(
-                config=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            config=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
@@ -259,7 +257,7 @@ class LlamaExtractClient:
             urllib.parse.urljoin(
                 f"{self._client_wrapper.get_base_url()}/", f"api/v1/extraction/extraction-agents/{extraction_agent_id}"
             ),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"data_schema": data_schema, "config": config}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -357,6 +355,7 @@ class LlamaExtractClient:
         client.llama_extract.run_job(
             request=ExtractJobCreate(
                 extraction_agent_id="string",
+                file_id="string",
                 config_override=ExtractConfig(
                     extraction_target=ExtractTarget.PER_DOC,
                     extraction_mode=ExtractMode.FAST,
@@ -433,6 +432,7 @@ class LlamaExtractClient:
             ExtractJobCreate,
             ExtractMode,
             ExtractTarget,
+            FailPageMode,
             LlamaExtractSettings,
             LlamaParseParameters,
             ParsingMode,
@@ -445,6 +445,7 @@ class LlamaExtractClient:
         client.llama_extract.run_job_test_user(
             job_create=ExtractJobCreate(
                 extraction_agent_id="string",
+                file_id="string",
                 config_override=ExtractConfig(
                     extraction_target=ExtractTarget.PER_DOC,
                     extraction_mode=ExtractMode.FAST,
@@ -454,6 +455,7 @@ class LlamaExtractClient:
                 chunk_mode=ChunkMode.PAGE,
                 llama_parse_params=LlamaParseParameters(
                     parse_mode=ParsingMode.PARSE_PAGE_WITHOUT_LLM,
+                    replace_failed_page_mode=FailPageMode.RAW_TEXT,
                 ),
             ),
         )
@@ -479,42 +481,98 @@ class LlamaExtractClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
+    def run_job_on_file(
+        self,
+        *,
+        from_ui: typing.Optional[bool] = None,
+        extraction_agent_id: str,
+        file: typing.IO,
+        data_schema_override: typing.Optional[str] = None,
+        config_override: typing.Optional[str] = None,
+    ) -> ExtractJob:
+        """
+        Parameters:
+            - from_ui: typing.Optional[bool].
+
+            - extraction_agent_id: str.
+
+            - file: typing.IO.
+
+            - data_schema_override: typing.Optional[str].
+
+            - config_override: typing.Optional[str].
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            "POST",
+            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/jobs/file"),
+            params=remove_none_from_dict({"from_ui": from_ui}),
+            data=jsonable_encoder(
+                {
+                    "extraction_agent_id": extraction_agent_id,
+                    "data_schema_override": data_schema_override,
+                    "config_override": config_override,
+                }
+            ),
+            files={"file": file},
+            headers=self._client_wrapper.get_headers(),
+            timeout=60,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic.parse_obj_as(ExtractJob, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+        try:
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, body=_response.text)
+        raise ApiError(status_code=_response.status_code, body=_response_json)
+
     def run_batch_jobs(
-        self, *, from_ui: typing.Optional[bool] = None, request: ExtractJobCreateBatch
+        self,
+        *,
+        from_ui: typing.Optional[bool] = None,
+        extraction_agent_id: str,
+        file_ids: typing.List[str],
+        data_schema_override: typing.Optional[ExtractJobCreateBatchDataSchemaOverride] = OMIT,
+        config_override: typing.Optional[ExtractConfig] = OMIT,
     ) -> typing.List[ExtractJob]:
         """
         Parameters:
             - from_ui: typing.Optional[bool].
 
-            - request: ExtractJobCreateBatch.
+            - extraction_agent_id: str. The id of the extraction agent
+
+            - file_ids: typing.List[str]. The ids of the files
+
+            - data_schema_override: typing.Optional[ExtractJobCreateBatchDataSchemaOverride]. The data schema to override the extraction agent's data schema with
+
+            - config_override: typing.Optional[ExtractConfig].
         ---
-        from llama_cloud import (
-            ExtractConfig,
-            ExtractJobCreateBatch,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import LlamaCloud
 
         client = LlamaCloud(
             token="YOUR_TOKEN",
         )
         client.llama_extract.run_batch_jobs(
-            request=ExtractJobCreateBatch(
-                extraction_agent_id="string",
-                file_ids=[],
-                config_override=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            extraction_agent_id="string",
+            file_ids=[],
+            config_override=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
+        _request: typing.Dict[str, typing.Any] = {"extraction_agent_id": extraction_agent_id, "file_ids": file_ids}
+        if data_schema_override is not OMIT:
+            _request["data_schema_override"] = data_schema_override
+        if config_override is not OMIT:
+            _request["config_override"] = config_override
         _response = self._client_wrapper.httpx_client.request(
             "POST",
             urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/jobs/batch"),
             params=remove_none_from_dict({"from_ui": from_ui}),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder(_request),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -721,17 +779,19 @@ class AsyncLlamaExtractClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
         self._client_wrapper = client_wrapper
 
-    async def list_extraction_agents(self, *, project_id: typing.Optional[str] = None) -> typing.List[ExtractAgent]:
+    async def list_extraction_agents(self, *, project_id: str) -> typing.List[ExtractAgent]:
         """
         Parameters:
-            - project_id: typing.Optional[str].
+            - project_id: str.
         ---
         from llama_cloud.client import AsyncLlamaCloud
 
         client = AsyncLlamaCloud(
             token="YOUR_TOKEN",
         )
-        await client.llama_extract.list_extraction_agents()
+        await client.llama_extract.list_extraction_agents(
+            project_id="string",
+        )
         """
         _response = await self._client_wrapper.httpx_client.request(
             "GET",
@@ -755,7 +815,9 @@ class AsyncLlamaExtractClient:
         *,
         project_id: typing.Optional[str] = None,
         organization_id: typing.Optional[str] = None,
-        request: ExtractAgentCreate,
+        name: str,
+        data_schema: ExtractAgentCreateDataSchema,
+        config: ExtractConfig,
     ) -> ExtractAgent:
         """
         Parameters:
@@ -763,26 +825,23 @@ class AsyncLlamaExtractClient:
 
             - organization_id: typing.Optional[str].
 
-            - request: ExtractAgentCreate.
+            - name: str. The name of the extraction schema
+
+            - data_schema: ExtractAgentCreateDataSchema. The schema of the data.
+
+            - config: ExtractConfig. The configuration parameters for the extraction agent.
         ---
-        from llama_cloud import (
-            ExtractAgentCreate,
-            ExtractConfig,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import AsyncLlamaCloud
 
         client = AsyncLlamaCloud(
             token="YOUR_TOKEN",
         )
         await client.llama_extract.create_extraction_agent(
-            request=ExtractAgentCreate(
-                name="string",
-                config=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            name="string",
+            config=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
@@ -790,7 +849,7 @@ class AsyncLlamaExtractClient:
             "POST",
             urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/extraction-agents"),
             params=remove_none_from_dict({"project_id": project_id, "organization_id": organization_id}),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"name": name, "data_schema": data_schema, "config": config}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -805,31 +864,28 @@ class AsyncLlamaExtractClient:
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
     async def validate_extraction_schema(
-        self, *, request: ExtractSchemaValidateRequest
+        self, *, data_schema: ExtractSchemaValidateRequestDataSchema
     ) -> ExtractSchemaValidateResponse:
         """
         Validates an extraction agent's schema definition.
         Returns the normalized and validated schema if valid, otherwise raises an HTTP 400.
 
         Parameters:
-            - request: ExtractSchemaValidateRequest.
+            - data_schema: ExtractSchemaValidateRequestDataSchema.
         ---
-        from llama_cloud import ExtractSchemaValidateRequest
         from llama_cloud.client import AsyncLlamaCloud
 
         client = AsyncLlamaCloud(
             token="YOUR_TOKEN",
         )
-        await client.llama_extract.validate_extraction_schema(
-            request=ExtractSchemaValidateRequest(),
-        )
+        await client.llama_extract.validate_extraction_schema()
         """
         _response = await self._client_wrapper.httpx_client.request(
             "POST",
             urllib.parse.urljoin(
                 f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/extraction-agents/schema/validation"
             ),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"data_schema": data_schema}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -910,19 +966,18 @@ class AsyncLlamaExtractClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
-    async def update_extraction_agent(self, extraction_agent_id: str, *, request: ExtractAgentUpdate) -> ExtractAgent:
+    async def update_extraction_agent(
+        self, extraction_agent_id: str, *, data_schema: ExtractAgentUpdateDataSchema, config: ExtractConfig
+    ) -> ExtractAgent:
         """
         Parameters:
             - extraction_agent_id: str.
 
-            - request: ExtractAgentUpdate.
+            - data_schema: ExtractAgentUpdateDataSchema. The schema of the data
+
+            - config: ExtractConfig. The configuration parameters for the extraction agent.
         ---
-        from llama_cloud import (
-            ExtractAgentUpdate,
-            ExtractConfig,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import AsyncLlamaCloud
 
         client = AsyncLlamaCloud(
@@ -930,11 +985,9 @@ class AsyncLlamaExtractClient:
         )
         await client.llama_extract.update_extraction_agent(
             extraction_agent_id="string",
-            request=ExtractAgentUpdate(
-                config=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            config=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
@@ -943,7 +996,7 @@ class AsyncLlamaExtractClient:
             urllib.parse.urljoin(
                 f"{self._client_wrapper.get_base_url()}/", f"api/v1/extraction/extraction-agents/{extraction_agent_id}"
             ),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder({"data_schema": data_schema, "config": config}),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )
@@ -1041,6 +1094,7 @@ class AsyncLlamaExtractClient:
         await client.llama_extract.run_job(
             request=ExtractJobCreate(
                 extraction_agent_id="string",
+                file_id="string",
                 config_override=ExtractConfig(
                     extraction_target=ExtractTarget.PER_DOC,
                     extraction_mode=ExtractMode.FAST,
@@ -1117,6 +1171,7 @@ class AsyncLlamaExtractClient:
             ExtractJobCreate,
             ExtractMode,
             ExtractTarget,
+            FailPageMode,
             LlamaExtractSettings,
             LlamaParseParameters,
             ParsingMode,
@@ -1129,6 +1184,7 @@ class AsyncLlamaExtractClient:
         await client.llama_extract.run_job_test_user(
             job_create=ExtractJobCreate(
                 extraction_agent_id="string",
+                file_id="string",
                 config_override=ExtractConfig(
                     extraction_target=ExtractTarget.PER_DOC,
                     extraction_mode=ExtractMode.FAST,
@@ -1138,6 +1194,7 @@ class AsyncLlamaExtractClient:
                 chunk_mode=ChunkMode.PAGE,
                 llama_parse_params=LlamaParseParameters(
                     parse_mode=ParsingMode.PARSE_PAGE_WITHOUT_LLM,
+                    replace_failed_page_mode=FailPageMode.RAW_TEXT,
                 ),
             ),
         )
@@ -1163,42 +1220,98 @@ class AsyncLlamaExtractClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
+    async def run_job_on_file(
+        self,
+        *,
+        from_ui: typing.Optional[bool] = None,
+        extraction_agent_id: str,
+        file: typing.IO,
+        data_schema_override: typing.Optional[str] = None,
+        config_override: typing.Optional[str] = None,
+    ) -> ExtractJob:
+        """
+        Parameters:
+            - from_ui: typing.Optional[bool].
+
+            - extraction_agent_id: str.
+
+            - file: typing.IO.
+
+            - data_schema_override: typing.Optional[str].
+
+            - config_override: typing.Optional[str].
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            "POST",
+            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/jobs/file"),
+            params=remove_none_from_dict({"from_ui": from_ui}),
+            data=jsonable_encoder(
+                {
+                    "extraction_agent_id": extraction_agent_id,
+                    "data_schema_override": data_schema_override,
+                    "config_override": config_override,
+                }
+            ),
+            files={"file": file},
+            headers=self._client_wrapper.get_headers(),
+            timeout=60,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic.parse_obj_as(ExtractJob, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+        try:
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, body=_response.text)
+        raise ApiError(status_code=_response.status_code, body=_response_json)
+
     async def run_batch_jobs(
-        self, *, from_ui: typing.Optional[bool] = None, request: ExtractJobCreateBatch
+        self,
+        *,
+        from_ui: typing.Optional[bool] = None,
+        extraction_agent_id: str,
+        file_ids: typing.List[str],
+        data_schema_override: typing.Optional[ExtractJobCreateBatchDataSchemaOverride] = OMIT,
+        config_override: typing.Optional[ExtractConfig] = OMIT,
     ) -> typing.List[ExtractJob]:
         """
         Parameters:
             - from_ui: typing.Optional[bool].
 
-            - request: ExtractJobCreateBatch.
+            - extraction_agent_id: str. The id of the extraction agent
+
+            - file_ids: typing.List[str]. The ids of the files
+
+            - data_schema_override: typing.Optional[ExtractJobCreateBatchDataSchemaOverride]. The data schema to override the extraction agent's data schema with
+
+            - config_override: typing.Optional[ExtractConfig].
         ---
-        from llama_cloud import (
-            ExtractConfig,
-            ExtractJobCreateBatch,
-            ExtractMode,
-            ExtractTarget,
-        )
+        from llama_cloud import ExtractConfig, ExtractMode, ExtractTarget
         from llama_cloud.client import AsyncLlamaCloud
 
         client = AsyncLlamaCloud(
             token="YOUR_TOKEN",
         )
         await client.llama_extract.run_batch_jobs(
-            request=ExtractJobCreateBatch(
-                extraction_agent_id="string",
-                file_ids=[],
-                config_override=ExtractConfig(
-                    extraction_target=ExtractTarget.PER_DOC,
-                    extraction_mode=ExtractMode.FAST,
-                ),
+            extraction_agent_id="string",
+            file_ids=[],
+            config_override=ExtractConfig(
+                extraction_target=ExtractTarget.PER_DOC,
+                extraction_mode=ExtractMode.FAST,
             ),
         )
         """
+        _request: typing.Dict[str, typing.Any] = {"extraction_agent_id": extraction_agent_id, "file_ids": file_ids}
+        if data_schema_override is not OMIT:
+            _request["data_schema_override"] = data_schema_override
+        if config_override is not OMIT:
+            _request["config_override"] = config_override
         _response = await self._client_wrapper.httpx_client.request(
             "POST",
             urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v1/extraction/jobs/batch"),
             params=remove_none_from_dict({"from_ui": from_ui}),
-            json=jsonable_encoder(request),
+            json=jsonable_encoder(_request),
             headers=self._client_wrapper.get_headers(),
             timeout=60,
         )

llama_cloud/resources/llama_extract/types/__init__.py

@@ -0,0 +1,21 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from .extract_agent_create_data_schema import ExtractAgentCreateDataSchema
+from .extract_agent_create_data_schema_zero_value import ExtractAgentCreateDataSchemaZeroValue
+from .extract_agent_update_data_schema import ExtractAgentUpdateDataSchema
+from .extract_agent_update_data_schema_zero_value import ExtractAgentUpdateDataSchemaZeroValue
+from .extract_job_create_batch_data_schema_override import ExtractJobCreateBatchDataSchemaOverride
+from .extract_job_create_batch_data_schema_override_zero_value import ExtractJobCreateBatchDataSchemaOverrideZeroValue
+from .extract_schema_validate_request_data_schema import ExtractSchemaValidateRequestDataSchema
+from .extract_schema_validate_request_data_schema_zero_value import ExtractSchemaValidateRequestDataSchemaZeroValue
+
+__all__ = [
+    "ExtractAgentCreateDataSchema",
+    "ExtractAgentCreateDataSchemaZeroValue",
+    "ExtractAgentUpdateDataSchema",
+    "ExtractAgentUpdateDataSchemaZeroValue",
+    "ExtractJobCreateBatchDataSchemaOverride",
+    "ExtractJobCreateBatchDataSchemaOverrideZeroValue",
+    "ExtractSchemaValidateRequestDataSchema",
+    "ExtractSchemaValidateRequestDataSchemaZeroValue",
+]