From 0692d24354f2b0ad3f6fd677802861ca5723c04e Mon Sep 17 00:00:00 2001 From: ashish uppala Date: Mon, 6 Apr 2026 11:09:21 -0400 Subject: [PATCH 1/3] add new pipelines / processor --- datalab_sdk/__init__.py | 18 + datalab_sdk/client.py | 858 ++++++++++++++++++++++++++++++++++- datalab_sdk/models.py | 174 ++++++- tests/test_client_methods.py | 410 +++++++++++++++++ 4 files changed, 1438 insertions(+), 22 deletions(-) diff --git a/datalab_sdk/__init__.py b/datalab_sdk/__init__.py index 8097661..eb148c0 100644 --- a/datalab_sdk/__init__.py +++ b/datalab_sdk/__init__.py @@ -15,6 +15,7 @@ ConvertOptions, ExtractOptions, SegmentOptions, + CustomProcessorOptions, CustomPipelineOptions, TrackChangesOptions, OCROptions, @@ -25,6 +26,14 @@ WorkflowExecution, InputConfig, UploadedFileMetadata, + ExtractionSchema, + PipelineStep, + PipelineConfig, + PipelineVersion, + PipelineExecution, + PipelineExecutionStepResult, + CustomProcessor, + CustomProcessorVersion, ) from .settings import settings @@ -42,6 +51,7 @@ "ConvertOptions", "ExtractOptions", "SegmentOptions", + "CustomProcessorOptions", "CustomPipelineOptions", "TrackChangesOptions", "OCROptions", @@ -52,4 +62,12 @@ "WorkflowExecution", "InputConfig", "UploadedFileMetadata", + "ExtractionSchema", + "PipelineStep", + "PipelineConfig", + "PipelineVersion", + "PipelineExecution", + "PipelineExecutionStepResult", + "CustomProcessor", + "CustomProcessorVersion", ] diff --git a/datalab_sdk/client.py b/datalab_sdk/client.py index f1136a4..72df01a 100644 --- a/datalab_sdk/client.py +++ b/datalab_sdk/client.py @@ -36,6 +36,7 @@ ConvertOptions, ExtractOptions, SegmentOptions, + CustomProcessorOptions, CustomPipelineOptions, TrackChangesOptions, OCROptions, @@ -46,6 +47,14 @@ WorkflowExecution, InputConfig, UploadedFileMetadata, + ExtractionSchema, + PipelineStep, + PipelineConfig, + PipelineVersion, + PipelineExecution, + PipelineExecutionStepResult, + CustomProcessor, + CustomProcessorVersion, ) from datalab_sdk.settings import settings @@ -512,15 +521,15 @@ async def extract( poll_interval: int = 1, ) -> Union[ConversionResult, FileResult]: """ - Extract structured data from a document using a JSON schema + Extract structured data from a document using a JSON schema or saved extraction schema - Provide a file for end-to-end processing, or set checkpoint_id in options - (from a previous convert() call with save_checkpoint=True) to skip re-parsing. + Provide a page_schema for inline extraction, or a schema_id to use a saved + extraction schema. These are mutually exclusive. Args: file_path: Path to the file to extract from file_url: URL of the file to extract from - options: Extraction options (must include page_schema) + options: Extraction options (must include page_schema or schema_id) save_output: Optional path to save output files stream_response_to: Optional path to stream raw JSON response to disk max_polls: Maximum number of polling attempts @@ -536,7 +545,17 @@ async def extract( raise ValueError(f"Directory does not exist: {resolved_stream_response_to.parent}") if options is None: - raise ValueError("options must be provided with page_schema") + raise ValueError("options must be provided with page_schema or schema_id") + + has_page_schema = bool(options.page_schema) + has_schema_id = bool(options.schema_id) + + if has_page_schema and has_schema_id: + raise ValueError("page_schema and schema_id are mutually exclusive. Provide one or the other.") + if not has_page_schema and not has_schema_id: + raise ValueError("Either page_schema or schema_id must be provided in options.") + if options.schema_version is not None and not has_schema_id: + raise ValueError("schema_version can only be used with schema_id.") has_file = file_path is not None or file_url is not None has_checkpoint = options.checkpoint_id is not None @@ -635,23 +654,23 @@ async def segment( return result - async def run_custom_pipeline( + async def run_custom_processor( self, file_path: Optional[Union[str, Path]] = None, file_url: Optional[str] = None, - options: Optional[CustomPipelineOptions] = None, + options: Optional[CustomProcessorOptions] = None, save_output: Optional[Union[str, Path]] = None, stream_response_to: Optional[Union[str, Path]] = None, max_polls: int = 300, poll_interval: int = 1, ) -> Union[ConversionResult, FileResult]: """ - Execute a custom pipeline configuration + Execute a custom processor on a document Args: file_path: Path to the file to process file_url: URL of the file to process - options: Custom pipeline options (must include pipeline_id) + options: Custom processor options (must include pipeline_id) save_output: Optional path to save output files stream_response_to: Optional path to stream raw JSON response to disk max_polls: Maximum number of polling attempts @@ -670,7 +689,7 @@ async def run_custom_pipeline( raise ValueError("options must be provided with pipeline_id") result_data = await self._submit_and_poll( - "/api/v1/custom-pipeline", + "/api/v1/custom-processor", data=self.get_form_params( file_path=file_path, file_url=file_url, options=options ), @@ -691,6 +710,36 @@ async def run_custom_pipeline( return result + async def run_custom_pipeline( + self, + file_path: Optional[Union[str, Path]] = None, + file_url: Optional[str] = None, + options: Optional[CustomProcessorOptions] = None, + save_output: Optional[Union[str, Path]] = None, + stream_response_to: Optional[Union[str, Path]] = None, + max_polls: int = 300, + poll_interval: int = 1, + ) -> Union[ConversionResult, FileResult]: + """Execute a custom processor on a document + + .. deprecated:: + Use run_custom_processor() instead. + """ + warnings.warn( + "run_custom_pipeline() is deprecated. Use run_custom_processor() instead.", + DeprecationWarning, + stacklevel=2, + ) + return await self.run_custom_processor( + file_path=file_path, + file_url=file_url, + options=options, + save_output=save_output, + stream_response_to=stream_response_to, + max_polls=max_polls, + poll_interval=poll_interval, + ) + async def track_changes( self, file_path: Optional[Union[str, Path]] = None, @@ -1567,6 +1616,546 @@ async def delete_file( "message": response.get("message", f"File {file_id} deleted successfully"), } + # --- Extraction Schema methods --- + + async def create_extraction_schema( + self, + name: str, + schema_json: Dict[str, Any], + description: Optional[str] = None, + ) -> ExtractionSchema: + """ + Create a new extraction schema + + Args: + name: Name for the schema (max 200 characters) + schema_json: JSON schema for extraction (must contain 'properties' key) + description: Optional description + """ + payload: Dict[str, Any] = {"name": name, "schema_json": schema_json} + if description is not None: + payload["description"] = description + + response = await self._make_request("POST", "/api/v1/extraction_schemas", json=payload) + return self._build_extraction_schema(response) + + async def list_extraction_schemas( + self, + limit: int = 50, + offset: int = 0, + include_archived: bool = False, + ) -> Dict[str, Any]: + """ + List extraction schemas for the authenticated user's team + + Args: + limit: Maximum number of schemas to return (default: 50, max: 200) + offset: Offset for pagination (default: 0) + include_archived: Include archived schemas (default: False) + """ + params = f"limit={limit}&offset={offset}&include_archived={str(include_archived).lower()}" + response = await self._make_request("GET", f"/api/v1/extraction_schemas?{params}") + return { + "schemas": [self._build_extraction_schema(s) for s in response.get("schemas", [])], + "total": response.get("total", 0), + } + + async def get_extraction_schema(self, schema_id: str) -> ExtractionSchema: + """ + Get an extraction schema by its schema_id + + Args: + schema_id: Schema ID string (e.g. sch_k8Hx9mP2nQ4v) + """ + response = await self._make_request("GET", f"/api/v1/extraction_schemas/{schema_id}") + return self._build_extraction_schema(response) + + async def update_extraction_schema( + self, + schema_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + schema_json: Optional[Dict[str, Any]] = None, + archived: Optional[bool] = None, + create_new_version: bool = False, + ) -> ExtractionSchema: + """ + Update an extraction schema + + Args: + schema_id: Schema ID string (e.g. sch_k8Hx9mP2nQ4v) + name: New name (max 200 characters) + description: New description + schema_json: New JSON schema (must contain 'properties' key) + archived: Set archived status + create_new_version: If True, bump version and save current state to history + """ + payload: Dict[str, Any] = {} + if name is not None: + payload["name"] = name + if description is not None: + payload["description"] = description + if schema_json is not None: + payload["schema_json"] = schema_json + if archived is not None: + payload["archived"] = archived + if create_new_version: + payload["create_new_version"] = True + + response = await self._make_request("PUT", f"/api/v1/extraction_schemas/{schema_id}", json=payload) + return self._build_extraction_schema(response) + + async def delete_extraction_schema(self, schema_id: str) -> ExtractionSchema: + """ + Delete (archive) an extraction schema + + Args: + schema_id: Schema ID string (e.g. sch_k8Hx9mP2nQ4v) + """ + response = await self._make_request("DELETE", f"/api/v1/extraction_schemas/{schema_id}") + return self._build_extraction_schema(response) + + @staticmethod + def _build_extraction_schema(data: Dict[str, Any]) -> ExtractionSchema: + return ExtractionSchema( + id=data.get("id"), + schema_id=data["schema_id"], + name=data["name"], + description=data.get("description"), + schema_json=data["schema_json"], + version=data.get("version", 1), + version_history=data.get("version_history"), + archived=data.get("archived", False), + created=data.get("created"), + updated=data.get("updated"), + ) + + # --- Pipeline CRUD methods --- + + async def create_pipeline( + self, + steps: list[PipelineStep], + ) -> PipelineConfig: + """ + Create a new pipeline + + Args: + steps: Ordered list of PipelineStep objects + """ + payload = {"steps": [s.to_dict() for s in steps]} + response = await self._make_request("POST", "/api/v1/pipelines", json=payload) + return self._build_pipeline_config(response) + + async def list_pipelines( + self, + saved_only: bool = True, + include_archived: bool = False, + limit: int = 50, + offset: int = 0, + ) -> Dict[str, Any]: + """ + List pipelines for the authenticated user's team + + Args: + saved_only: Only return saved pipelines (default: True) + include_archived: Include archived pipelines (default: False) + limit: Maximum number to return (default: 50, max: 200) + offset: Offset for pagination (default: 0) + """ + params = ( + f"saved_only={str(saved_only).lower()}" + f"&include_archived={str(include_archived).lower()}" + f"&limit={limit}&offset={offset}" + ) + response = await self._make_request("GET", f"/api/v1/pipelines?{params}") + return { + "pipelines": [self._build_pipeline_config(p) for p in response.get("pipelines", [])], + "total": response.get("total", 0), + } + + async def get_pipeline(self, pipeline_id: str) -> PipelineConfig: + """ + Get a pipeline by its pipeline_id + + Args: + pipeline_id: Pipeline ID string (e.g. pl_k8Hx9mP2nQ4v) + """ + response = await self._make_request("GET", f"/api/v1/pipelines/{pipeline_id}") + return self._build_pipeline_config(response) + + async def update_pipeline( + self, + pipeline_id: str, + steps: list[PipelineStep], + ) -> PipelineConfig: + """ + Update pipeline steps (auto-save path for draft edits) + + Args: + pipeline_id: Pipeline ID string + steps: New ordered list of PipelineStep objects + """ + payload = {"steps": [s.to_dict() for s in steps]} + response = await self._make_request("PUT", f"/api/v1/pipelines/{pipeline_id}", json=payload) + return self._build_pipeline_config(response) + + async def save_pipeline( + self, + pipeline_id: str, + name: str = "", + ) -> PipelineConfig: + """ + Name and promote a pipeline to saved status + + Args: + pipeline_id: Pipeline ID string + name: Display name (auto-generated if empty) + """ + response = await self._make_request( + "PUT", f"/api/v1/pipelines/{pipeline_id}/save", json={"name": name} + ) + return self._build_pipeline_config(response) + + async def archive_pipeline(self, pipeline_id: str) -> Dict[str, Any]: + """Archive a pipeline, hiding it from the default list""" + return await self._make_request("POST", f"/api/v1/pipelines/{pipeline_id}/archive") + + async def unarchive_pipeline(self, pipeline_id: str) -> Dict[str, Any]: + """Unarchive a pipeline, restoring it to the default list""" + return await self._make_request("POST", f"/api/v1/pipelines/{pipeline_id}/unarchive") + + # --- Pipeline Versioning methods --- + + async def create_pipeline_version( + self, + pipeline_id: str, + description: Optional[str] = None, + ) -> PipelineVersion: + """ + Create a new version snapshot of the pipeline's current steps + + Args: + pipeline_id: Pipeline ID string + description: Optional description for this version + """ + payload: Dict[str, Any] = {} + if description is not None: + payload["description"] = description + response = await self._make_request( + "POST", f"/api/v1/pipelines/{pipeline_id}/versions", json=payload + ) + return PipelineVersion( + id=response.get("id"), + version=response["version"], + steps=response.get("steps", []), + description=response.get("description"), + created=response.get("created"), + ) + + async def list_pipeline_versions(self, pipeline_id: str) -> Dict[str, Any]: + """List all versions of a pipeline, newest first""" + response = await self._make_request("GET", f"/api/v1/pipelines/{pipeline_id}/versions") + return { + "versions": [ + PipelineVersion( + id=v.get("id"), + version=v["version"], + steps=v.get("steps", []), + description=v.get("description"), + created=v.get("created"), + ) + for v in response.get("versions", []) + ], + "total": response.get("total", 0), + } + + async def discard_pipeline_draft( + self, + pipeline_id: str, + version: Optional[int] = None, + ) -> PipelineConfig: + """ + Discard draft changes and revert to a published version + + Args: + pipeline_id: Pipeline ID string + version: Version to revert to (default: active published version) + """ + payload: Dict[str, Any] = {} + if version is not None: + payload["version"] = version + response = await self._make_request( + "POST", f"/api/v1/pipelines/{pipeline_id}/discard", json=payload + ) + return self._build_pipeline_config(response) + + async def get_pipeline_rate(self, pipeline_id: str) -> Dict[str, Any]: + """ + Get the per-page rate for a pipeline + + Returns dict with rate_per_1000_pages_cents and rate_breakdown. + """ + return await self._make_request("GET", f"/api/v1/pipelines/{pipeline_id}/rate") + + # --- Pipeline Execution methods --- + + async def run_pipeline( + self, + pipeline_id: str, + file_path: Optional[Union[str, Path]] = None, + file_url: Optional[str] = None, + page_range: Optional[str] = None, + output_format: Optional[str] = None, + run_evals: bool = False, + skip_cache: bool = False, + webhook_url: Optional[str] = None, + version: Optional[int] = None, + max_polls: int = 1, + poll_interval: int = 1, + ) -> PipelineExecution: + """ + Execute a pipeline on a file + + Args: + pipeline_id: Pipeline ID (pl_XXXXX) + file_path: Path to the file to process + file_url: URL of the file to process + page_range: Page range to process (e.g. '0,2-4') + output_format: Output format (json, html, markdown, chunks) + run_evals: Whether to run evaluation steps + skip_cache: Skip executor cache + webhook_url: URL to POST when complete + version: Pipeline version to execute (0=draft, omit=active) + max_polls: Maximum polling attempts after submission (default: 1) + poll_interval: Seconds between polls + """ + form_data = self.get_form_params(file_path=file_path, file_url=file_url) + if page_range is not None: + form_data.add_field("page_range", page_range) + if output_format is not None: + form_data.add_field("output_format", output_format) + if run_evals: + form_data.add_field("run_evals", str(run_evals)) + if skip_cache: + form_data.add_field("skip_cache", str(skip_cache)) + if webhook_url is not None: + form_data.add_field("webhook_url", webhook_url) + if version is not None: + form_data.add_field("version", str(version)) + + response = await self._submit_with_retry( + f"/api/v1/pipelines/{pipeline_id}/run", data=form_data + ) + execution = self._build_pipeline_execution(response) + + # Poll if requested + if max_polls > 1 and execution.status not in ("completed", "completed_with_errors", "failed"): + return await self.get_pipeline_execution( + execution.execution_id, max_polls=max_polls - 1, poll_interval=poll_interval + ) + return execution + + async def get_pipeline_execution( + self, + execution_id: str, + max_polls: int = 1, + poll_interval: int = 1, + ) -> PipelineExecution: + """ + Get the status of a pipeline execution, optionally polling until completion + + Args: + execution_id: Execution ID (pex_XXXXX) + max_polls: Maximum polling attempts (default: 1 for single check) + poll_interval: Seconds between polls + """ + for i in range(max_polls): + response = await self._make_request( + "GET", f"/api/v1/pipelines/executions/{execution_id}" + ) + execution = self._build_pipeline_execution(response) + + if execution.status in ("completed", "completed_with_errors", "failed"): + return execution + + if i < max_polls - 1: + await asyncio.sleep(poll_interval) + + return execution + + async def list_pipeline_executions( + self, + pipeline_id: str, + limit: int = 20, + offset: int = 0, + ) -> Dict[str, Any]: + """List recent executions for a pipeline""" + response = await self._make_request( + "GET", f"/api/v1/pipelines/{pipeline_id}/executions?limit={limit}&offset={offset}" + ) + return { + "executions": [ + self._build_pipeline_execution(e) for e in response.get("executions", []) + ], + "total": response.get("total", 0), + } + + async def get_step_result( + self, + execution_id: str, + step_index: int, + ) -> Dict[str, Any]: + """ + Fetch the result of a specific pipeline execution step + + Args: + execution_id: Execution ID (pex_XXXXX) + step_index: Zero-based step index + """ + return await self._make_request( + "GET", f"/api/v1/pipelines/executions/{execution_id}/steps/{step_index}/result" + ) + + @staticmethod + def _build_pipeline_config(data: Dict[str, Any]) -> PipelineConfig: + return PipelineConfig( + id=data.get("id"), + pipeline_id=data["pipeline_id"], + name=data.get("name"), + steps=data.get("steps", []), + is_saved=data.get("is_saved", False), + archived=data.get("archived", False), + active_version=data.get("active_version", 0), + created=data.get("created"), + updated=data.get("updated"), + ) + + @staticmethod + def _build_pipeline_execution(data: Dict[str, Any]) -> PipelineExecution: + steps = [ + PipelineExecutionStepResult( + step_index=s["step_index"], + step_type=s["step_type"], + status=s["status"], + lookup_key=s.get("lookup_key"), + result_url=s.get("result_url"), + started_at=s.get("started_at"), + finished_at=s.get("finished_at"), + error_message=s.get("error_message"), + checkpoint_id=s.get("checkpoint_id"), + ) + for s in data.get("steps", []) + ] + return PipelineExecution( + execution_id=data["execution_id"], + pipeline_id=data.get("pipeline_id", ""), + pipeline_version=data.get("pipeline_version", 0), + status=data.get("status", "pending"), + steps=steps, + started_at=data.get("started_at"), + completed_at=data.get("completed_at"), + created=data.get("created"), + config_snapshot=data.get("config_snapshot"), + input_config=data.get("input_config"), + rate_breakdown=data.get("rate_breakdown"), + ) + + # --- Custom Processor Management methods --- + + async def list_custom_processors( + self, + limit: int = 50, + offset: int = 0, + ) -> Dict[str, Any]: + """ + List custom processors for the authenticated user's team + + Args: + limit: Maximum number to return (default: 50) + offset: Offset for pagination (default: 0) + """ + response = await self._make_request( + "GET", f"/api/v1/custom_processors?limit={limit}&offset={offset}" + ) + processors = [ + CustomProcessor( + processor_id=p["processor_id"], + name=p.get("name"), + status=p.get("status", ""), + success=p.get("success"), + active_version=p.get("active_version", 0), + max_version=p.get("max_version", 0), + iteration_in_progress=p.get("iteration_in_progress", False), + pipeline_id=p.get("pipeline_id"), + created_at=p.get("created_at"), + completed_at=p.get("completed_at"), + error_message=p.get("error_message"), + eval_rubric_id=p.get("eval_rubric_id"), + ) + for p in response.get("pipelines", []) + ] + return {"processors": processors} + + async def get_custom_processor_status(self, lookup_key: str) -> Dict[str, Any]: + """ + Check the status of a custom processor request + + Args: + lookup_key: The lookup key returned when the processor was submitted + """ + return await self._make_request("GET", f"/api/v1/custom_processors/{lookup_key}") + + async def list_custom_processor_versions(self, processor_id: str) -> Dict[str, Any]: + """ + List versions of a custom processor + + Args: + processor_id: Processor ID (cp_XXXXX) + """ + response = await self._make_request( + "GET", f"/api/v1/custom_processors/{processor_id}/versions" + ) + versions = [ + CustomProcessorVersion( + version=v["version"], + request_description=v.get("request_description", ""), + created_at=v.get("created_at"), + runtime=v.get("runtime"), + is_active=v.get("is_active", False), + ) + for v in response.get("versions", []) + ] + return {"versions": versions} + + async def set_active_processor_version( + self, + processor_id: str, + version: int, + ) -> Dict[str, Any]: + """ + Set the active version of a custom processor + + Args: + processor_id: Processor ID (cp_XXXXX) + version: Version number to activate + """ + form_data = aiohttp.FormData() + form_data.add_field("version", str(version)) + return await self._submit_with_retry( + f"/api/v1/custom_processors/{processor_id}/set_active", data=form_data + ) + + async def archive_custom_processor(self, processor_id: str) -> Dict[str, Any]: + """ + Archive a custom processor + + Args: + processor_id: Processor ID (cp_XXXXX) + """ + return await self._make_request( + "POST", f"/api/v1/custom_processors/{processor_id}/archive" + ) + class DatalabClient: """Synchronous wrapper around AsyncDatalabClient""" @@ -1646,12 +2235,12 @@ def extract( poll_interval: int = 1, ) -> Union[ConversionResult, FileResult]: """ - Extract structured data from a document using a JSON schema (sync version) + Extract structured data using a JSON schema or saved extraction schema (sync version) Args: file_path: Path to the file to extract from file_url: URL of the file to extract from - options: Extraction options (must include page_schema) + options: Extraction options (must include page_schema or schema_id) save_output: Optional path to save output files stream_response_to: Optional path to stream raw JSON response to disk max_polls: Maximum number of polling attempts @@ -1703,30 +2292,30 @@ def segment( ) ) - def run_custom_pipeline( + def run_custom_processor( self, file_path: Optional[Union[str, Path]] = None, file_url: Optional[str] = None, - options: Optional[CustomPipelineOptions] = None, + options: Optional[CustomProcessorOptions] = None, save_output: Optional[Union[str, Path]] = None, stream_response_to: Optional[Union[str, Path]] = None, max_polls: int = 300, poll_interval: int = 1, ) -> Union[ConversionResult, FileResult]: """ - Execute a custom pipeline configuration (sync version) + Execute a custom processor on a document (sync version) Args: file_path: Path to the file to process file_url: URL of the file to process - options: Custom pipeline options (must include pipeline_id) + options: Custom processor options (must include pipeline_id) save_output: Optional path to save output files stream_response_to: Optional path to stream raw JSON response to disk max_polls: Maximum number of polling attempts poll_interval: Seconds between polling attempts """ return self._run_async( - self._async_client.run_custom_pipeline( + self._async_client.run_custom_processor( file_path=file_path, file_url=file_url, options=options, @@ -1737,6 +2326,36 @@ def run_custom_pipeline( ) ) + def run_custom_pipeline( + self, + file_path: Optional[Union[str, Path]] = None, + file_url: Optional[str] = None, + options: Optional[CustomProcessorOptions] = None, + save_output: Optional[Union[str, Path]] = None, + stream_response_to: Optional[Union[str, Path]] = None, + max_polls: int = 300, + poll_interval: int = 1, + ) -> Union[ConversionResult, FileResult]: + """Execute a custom processor on a document (sync version) + + .. deprecated:: + Use run_custom_processor() instead. + """ + warnings.warn( + "run_custom_pipeline() is deprecated. Use run_custom_processor() instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.run_custom_processor( + file_path=file_path, + file_url=file_url, + options=options, + save_output=save_output, + stream_response_to=stream_response_to, + max_polls=max_polls, + poll_interval=poll_interval, + ) + def track_changes( self, file_path: Optional[Union[str, Path]] = None, @@ -2051,3 +2670,208 @@ def delete_workflow(self, workflow_id: int) -> Dict[str, Any]: return self._run_async( self._async_client.delete_workflow(workflow_id=workflow_id) ) + + # --- Extraction Schema methods (sync) --- + + def create_extraction_schema( + self, name: str, schema_json: Dict[str, Any], description: Optional[str] = None, + ) -> ExtractionSchema: + """Create a new extraction schema (sync version)""" + return self._run_async( + self._async_client.create_extraction_schema( + name=name, schema_json=schema_json, description=description, + ) + ) + + def list_extraction_schemas( + self, limit: int = 50, offset: int = 0, include_archived: bool = False, + ) -> Dict[str, Any]: + """List extraction schemas (sync version)""" + return self._run_async( + self._async_client.list_extraction_schemas( + limit=limit, offset=offset, include_archived=include_archived, + ) + ) + + def get_extraction_schema(self, schema_id: str) -> ExtractionSchema: + """Get an extraction schema by ID (sync version)""" + return self._run_async(self._async_client.get_extraction_schema(schema_id=schema_id)) + + def update_extraction_schema( + self, + schema_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + schema_json: Optional[Dict[str, Any]] = None, + archived: Optional[bool] = None, + create_new_version: bool = False, + ) -> ExtractionSchema: + """Update an extraction schema (sync version)""" + return self._run_async( + self._async_client.update_extraction_schema( + schema_id=schema_id, name=name, description=description, + schema_json=schema_json, archived=archived, + create_new_version=create_new_version, + ) + ) + + def delete_extraction_schema(self, schema_id: str) -> ExtractionSchema: + """Delete (archive) an extraction schema (sync version)""" + return self._run_async(self._async_client.delete_extraction_schema(schema_id=schema_id)) + + # --- Pipeline methods (sync) --- + + def create_pipeline(self, steps: list[PipelineStep]) -> PipelineConfig: + """Create a new pipeline (sync version)""" + return self._run_async(self._async_client.create_pipeline(steps=steps)) + + def list_pipelines( + self, saved_only: bool = True, include_archived: bool = False, + limit: int = 50, offset: int = 0, + ) -> Dict[str, Any]: + """List pipelines (sync version)""" + return self._run_async( + self._async_client.list_pipelines( + saved_only=saved_only, include_archived=include_archived, + limit=limit, offset=offset, + ) + ) + + def get_pipeline(self, pipeline_id: str) -> PipelineConfig: + """Get a pipeline by ID (sync version)""" + return self._run_async(self._async_client.get_pipeline(pipeline_id=pipeline_id)) + + def update_pipeline(self, pipeline_id: str, steps: list[PipelineStep]) -> PipelineConfig: + """Update pipeline steps (sync version)""" + return self._run_async( + self._async_client.update_pipeline(pipeline_id=pipeline_id, steps=steps) + ) + + def save_pipeline(self, pipeline_id: str, name: str = "") -> PipelineConfig: + """Save and name a pipeline (sync version)""" + return self._run_async( + self._async_client.save_pipeline(pipeline_id=pipeline_id, name=name) + ) + + def archive_pipeline(self, pipeline_id: str) -> Dict[str, Any]: + """Archive a pipeline (sync version)""" + return self._run_async(self._async_client.archive_pipeline(pipeline_id=pipeline_id)) + + def unarchive_pipeline(self, pipeline_id: str) -> Dict[str, Any]: + """Unarchive a pipeline (sync version)""" + return self._run_async(self._async_client.unarchive_pipeline(pipeline_id=pipeline_id)) + + def create_pipeline_version( + self, pipeline_id: str, description: Optional[str] = None, + ) -> PipelineVersion: + """Create a pipeline version snapshot (sync version)""" + return self._run_async( + self._async_client.create_pipeline_version( + pipeline_id=pipeline_id, description=description, + ) + ) + + def list_pipeline_versions(self, pipeline_id: str) -> Dict[str, Any]: + """List pipeline versions (sync version)""" + return self._run_async(self._async_client.list_pipeline_versions(pipeline_id=pipeline_id)) + + def discard_pipeline_draft( + self, pipeline_id: str, version: Optional[int] = None, + ) -> PipelineConfig: + """Discard draft and revert to a published version (sync version)""" + return self._run_async( + self._async_client.discard_pipeline_draft( + pipeline_id=pipeline_id, version=version, + ) + ) + + def get_pipeline_rate(self, pipeline_id: str) -> Dict[str, Any]: + """Get pipeline per-page rate (sync version)""" + return self._run_async(self._async_client.get_pipeline_rate(pipeline_id=pipeline_id)) + + def run_pipeline( + self, + pipeline_id: str, + file_path: Optional[Union[str, Path]] = None, + file_url: Optional[str] = None, + page_range: Optional[str] = None, + output_format: Optional[str] = None, + run_evals: bool = False, + skip_cache: bool = False, + webhook_url: Optional[str] = None, + version: Optional[int] = None, + max_polls: int = 1, + poll_interval: int = 1, + ) -> PipelineExecution: + """Execute a pipeline on a file (sync version)""" + return self._run_async( + self._async_client.run_pipeline( + pipeline_id=pipeline_id, file_path=file_path, file_url=file_url, + page_range=page_range, output_format=output_format, + run_evals=run_evals, skip_cache=skip_cache, + webhook_url=webhook_url, version=version, + max_polls=max_polls, poll_interval=poll_interval, + ) + ) + + def get_pipeline_execution( + self, execution_id: str, max_polls: int = 1, poll_interval: int = 1, + ) -> PipelineExecution: + """Get pipeline execution status (sync version)""" + return self._run_async( + self._async_client.get_pipeline_execution( + execution_id=execution_id, max_polls=max_polls, poll_interval=poll_interval, + ) + ) + + def list_pipeline_executions( + self, pipeline_id: str, limit: int = 20, offset: int = 0, + ) -> Dict[str, Any]: + """List pipeline executions (sync version)""" + return self._run_async( + self._async_client.list_pipeline_executions( + pipeline_id=pipeline_id, limit=limit, offset=offset, + ) + ) + + def get_step_result(self, execution_id: str, step_index: int) -> Dict[str, Any]: + """Get a pipeline execution step result (sync version)""" + return self._run_async( + self._async_client.get_step_result( + execution_id=execution_id, step_index=step_index, + ) + ) + + # --- Custom Processor Management methods (sync) --- + + def list_custom_processors(self, limit: int = 50, offset: int = 0) -> Dict[str, Any]: + """List custom processors (sync version)""" + return self._run_async( + self._async_client.list_custom_processors(limit=limit, offset=offset) + ) + + def get_custom_processor_status(self, lookup_key: str) -> Dict[str, Any]: + """Check custom processor request status (sync version)""" + return self._run_async( + self._async_client.get_custom_processor_status(lookup_key=lookup_key) + ) + + def list_custom_processor_versions(self, processor_id: str) -> Dict[str, Any]: + """List custom processor versions (sync version)""" + return self._run_async( + self._async_client.list_custom_processor_versions(processor_id=processor_id) + ) + + def set_active_processor_version(self, processor_id: str, version: int) -> Dict[str, Any]: + """Set active processor version (sync version)""" + return self._run_async( + self._async_client.set_active_processor_version( + processor_id=processor_id, version=version, + ) + ) + + def archive_custom_processor(self, processor_id: str) -> Dict[str, Any]: + """Archive a custom processor (sync version)""" + return self._run_async( + self._async_client.archive_custom_processor(processor_id=processor_id) + ) diff --git a/datalab_sdk/models.py b/datalab_sdk/models.py index 147c036..aaf57de 100644 --- a/datalab_sdk/models.py +++ b/datalab_sdk/models.py @@ -53,6 +53,7 @@ class ConvertOptions(ProcessingOptions): add_block_ids: bool = False # add block IDs to HTML output include_markdown_in_chunks: bool = False # include markdown field in chunks/JSON output token_efficient_markdown: bool = False # optimize markdown for LLM token usage + eval_rubric_id: Optional[int] = None # run evaluation rubric after conversion def to_form_data(self) -> Dict[str, Any]: """Convert to form data format for API requests""" @@ -78,13 +79,23 @@ def to_form_data(self) -> Dict[str, Any]: class ExtractOptions(ProcessingOptions): """Options for structured data extraction via /extract endpoint""" - page_schema: str = "" # Required - JSON schema with 'properties' key + page_schema: str = "" # JSON schema with 'properties' key. Mutually exclusive with schema_id. + schema_id: Optional[str] = None # ID of a saved extraction schema (e.g. sch_k8Hx9mP2nQ4v). Mutually exclusive with page_schema. + schema_version: Optional[int] = None # Version of the schema. Only valid with schema_id. checkpoint_id: Optional[str] = None # From previous /convert with save_checkpoint=true mode: str = "fast" # fast, balanced, accurate output_format: str = "markdown" # markdown, json, html, chunks save_checkpoint: bool = False webhook_url: Optional[str] = None + def to_form_data(self) -> Dict[str, Any]: + """Convert to form data format for API requests""" + form_data = super().to_form_data() + # When using schema_id, suppress the empty default page_schema + if self.schema_id and not self.page_schema: + form_data.pop("page_schema", None) + return form_data + @dataclass class SegmentOptions(ProcessingOptions): @@ -98,15 +109,34 @@ class SegmentOptions(ProcessingOptions): @dataclass -class CustomPipelineOptions(ProcessingOptions): - """Options for running a custom pipeline via /custom-pipeline endpoint""" +class CustomProcessorOptions(ProcessingOptions): + """Options for running a custom processor via /custom-processor endpoint""" - pipeline_id: str = "" # Required - custom pipeline ID (cp_XXXXX format) - run_eval: bool = False # Run evaluation rules defined for the pipeline + pipeline_id: str = "" # Required - custom processor ID (cp_XXXXX format) + version: Optional[int] = None # Specify processor version to run (default: active version) + run_eval: bool = False # Run evaluation rules defined for the processor mode: str = "fast" # fast, balanced, accurate output_format: str = "markdown" # markdown, json, html, chunks + paginate: bool = False # Separate pages with horizontal rules containing page numbers + add_block_ids: bool = False # Add data-block-id attributes to HTML elements + include_markdown_in_chunks: bool = False # Include markdown field in chunks/JSON output + disable_image_extraction: bool = False # Disable image extraction from the document + disable_image_captions: bool = False # Disable synthetic image captions/descriptions webhook_url: Optional[str] = None + @property + def processor_id(self) -> str: + """Alias for pipeline_id (the API wire format uses pipeline_id)""" + return self.pipeline_id + + @processor_id.setter + def processor_id(self, value: str): + self.pipeline_id = value + + +# Backward-compatible alias +CustomPipelineOptions = CustomProcessorOptions + @dataclass class TrackChangesOptions(ProcessingOptions): @@ -509,3 +539,137 @@ class FileResult: status: str output_path: Path error: Optional[str] = None + + +# --- Extraction Schema models --- + + +@dataclass +class ExtractionSchema: + """Represents a saved extraction schema for structured extraction""" + + schema_id: str + name: str + schema_json: Dict[str, Any] + id: Optional[int] = None + description: Optional[str] = None + version: int = 1 + version_history: Optional[List[Dict[str, Any]]] = None + archived: bool = False + created: Optional[str] = None + updated: Optional[str] = None + + +# --- Pipeline models --- + + +@dataclass +class PipelineStep: + """Configuration for a single pipeline step""" + + type: str # convert, extract, segment, custom + settings: Dict[str, Any] = field(default_factory=dict) + custom_processor_id: Optional[str] = None # For custom steps (cp_XXXXX) + eval_rubric_id: Optional[int] = None + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for API requests""" + data: Dict[str, Any] = { + "type": self.type, + "settings": self.settings, + } + if self.custom_processor_id is not None: + data["custom_processor_id"] = self.custom_processor_id + if self.eval_rubric_id is not None: + data["eval_rubric_id"] = self.eval_rubric_id + return data + + +@dataclass +class PipelineConfig: + """Represents a pipeline definition""" + + pipeline_id: str # pl_XXXXX + steps: List[Dict[str, Any]] + name: Optional[str] = None + is_saved: bool = False + archived: bool = False + active_version: int = 0 + id: Optional[int] = None + created: Optional[str] = None + updated: Optional[str] = None + + +@dataclass +class PipelineVersion: + """Immutable version snapshot of a pipeline""" + + version: int + steps: List[Dict[str, Any]] + description: Optional[str] = None + id: Optional[int] = None + created: Optional[str] = None + + +@dataclass +class PipelineExecutionStepResult: + """Status of a single step within a pipeline execution""" + + step_index: int + step_type: str + status: str # pending, dispatched, running, completed, failed + lookup_key: Optional[str] = None + result_url: Optional[str] = None + started_at: Optional[str] = None + finished_at: Optional[str] = None + error_message: Optional[str] = None + checkpoint_id: Optional[str] = None + + +@dataclass +class PipelineExecution: + """Result from pipeline execution""" + + execution_id: str # pex_XXXXX + pipeline_id: str # pl_XXXXX + pipeline_version: int + status: str # pending, running, completed, completed_with_errors, failed + steps: List[PipelineExecutionStepResult] = field(default_factory=list) + started_at: Optional[str] = None + completed_at: Optional[str] = None + created: Optional[str] = None + config_snapshot: Optional[Dict[str, Any]] = None + input_config: Optional[Dict[str, Any]] = None + rate_breakdown: Optional[Dict[str, Any]] = None + + +# --- Custom Processor models --- + + +@dataclass +class CustomProcessor: + """Represents a custom processor (formerly custom pipeline)""" + + processor_id: str # cp_XXXXX + status: str # processing, completed, failed + name: Optional[str] = None + success: Optional[bool] = None + active_version: int = 0 + max_version: int = 0 + iteration_in_progress: bool = False + pipeline_id: Optional[str] = None # auto-created workspace pipeline (pl_XXXXX) + created_at: Optional[str] = None + completed_at: Optional[str] = None + error_message: Optional[str] = None + eval_rubric_id: Optional[int] = None + + +@dataclass +class CustomProcessorVersion: + """Version information for a custom processor""" + + version: int + request_description: str = "" + created_at: Optional[str] = None + runtime: Optional[float] = None + is_active: bool = False diff --git a/tests/test_client_methods.py b/tests/test_client_methods.py index a6038be..ed6cf7e 100644 --- a/tests/test_client_methods.py +++ b/tests/test_client_methods.py @@ -16,8 +16,16 @@ ExtractOptions, SegmentOptions, CustomPipelineOptions, + CustomProcessorOptions, TrackChangesOptions, OCROptions, + ExtractionSchema, + PipelineStep, + PipelineConfig, + PipelineExecution, + PipelineExecutionStepResult, + CustomProcessor, + CustomProcessorVersion, ) from datalab_sdk.exceptions import ( DatalabAPIError, @@ -789,3 +797,405 @@ async def test_poll_result_raises_on_failed_status(self): await client._poll_result( "https://api.example.com/check", max_polls=1, poll_interval=0 ) + + +class TestExtractWithSchemaId: + """Test extract with schema_id support""" + + @pytest.mark.asyncio + async def test_extract_with_schema_id(self, temp_dir): + """Test extraction using a saved schema ID""" + pdf_file = temp_dir / "test.pdf" + pdf_file.write_bytes(b"%PDF-1.4\n%Test PDF content\n%%EOF\n") + + mock_initial = { + "success": True, + "request_id": "ext-schema", + "request_check_url": "https://api.datalab.to/api/v1/extract/ext-schema", + } + mock_result = { + "success": True, + "status": "complete", + "output_format": "markdown", + "extraction_schema_json": '{"name": "John"}', + } + + options = ExtractOptions(schema_id="sch_k8Hx9mP2nQ4v") + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + with patch.object(client, "_poll_result", new_callable=AsyncMock) as mock_poll: + mock_req.return_value = mock_initial + mock_poll.return_value = mock_result + result = await client.extract(pdf_file, options=options) + assert result.success is True + + @pytest.mark.asyncio + async def test_extract_rejects_both_page_schema_and_schema_id(self): + options = ExtractOptions( + page_schema='{"properties": {"name": {"type": "string"}}}', + schema_id="sch_k8Hx9mP2nQ4v", + ) + async with AsyncDatalabClient(api_key="test-key") as client: + with pytest.raises(ValueError, match="mutually exclusive"): + await client.extract(file_path="test.pdf", options=options) + + @pytest.mark.asyncio + async def test_extract_rejects_neither_page_schema_nor_schema_id(self): + options = ExtractOptions() + async with AsyncDatalabClient(api_key="test-key") as client: + with pytest.raises(ValueError, match="Either page_schema or schema_id"): + await client.extract(file_path="test.pdf", options=options) + + @pytest.mark.asyncio + async def test_extract_rejects_schema_version_without_schema_id(self): + options = ExtractOptions( + page_schema='{"properties": {"name": {"type": "string"}}}', + schema_version=2, + ) + async with AsyncDatalabClient(api_key="test-key") as client: + with pytest.raises(ValueError, match="schema_version can only be used with schema_id"): + await client.extract(file_path="test.pdf", options=options) + + def test_extract_options_form_data_suppresses_empty_page_schema(self): + """When schema_id is set and page_schema is empty, page_schema should not be in form data""" + options = ExtractOptions(schema_id="sch_abc123") + form_data = options.to_form_data() + assert "schema_id" in form_data + assert "page_schema" not in form_data + + +class TestCustomProcessorOptions: + """Test CustomProcessorOptions new fields and alias""" + + def test_new_fields_have_defaults(self): + options = CustomProcessorOptions(pipeline_id="cp_abc12") + assert options.version is None + assert options.paginate is False + assert options.add_block_ids is False + assert options.include_markdown_in_chunks is False + assert options.disable_image_extraction is False + assert options.disable_image_captions is False + + def test_processor_id_alias(self): + options = CustomProcessorOptions(pipeline_id="cp_abc12") + assert options.processor_id == "cp_abc12" + options.processor_id = "cp_xyz99" + assert options.pipeline_id == "cp_xyz99" + + def test_backward_compatible_alias(self): + """CustomPipelineOptions should be the same class""" + assert CustomPipelineOptions is CustomProcessorOptions + + def test_form_data_includes_new_fields(self): + options = CustomProcessorOptions( + pipeline_id="cp_abc12", version=3, paginate=True, disable_image_extraction=True, + ) + form_data = options.to_form_data() + assert "version" in form_data + assert "paginate" in form_data + assert "disable_image_extraction" in form_data + + +class TestConvertEvalRubricId: + """Test eval_rubric_id on ConvertOptions""" + + def test_default_none(self): + options = ConvertOptions() + assert options.eval_rubric_id is None + form_data = options.to_form_data() + assert "eval_rubric_id" not in form_data + + def test_serialized_when_set(self): + options = ConvertOptions(eval_rubric_id=42) + form_data = options.to_form_data() + assert "eval_rubric_id" in form_data + + +class TestRunCustomProcessorMethod: + """Test run_custom_processor and deprecation of run_custom_pipeline""" + + @pytest.mark.asyncio + async def test_run_custom_processor(self, temp_dir): + pdf_file = temp_dir / "test.pdf" + pdf_file.write_bytes(b"%PDF-1.4\n%Test PDF content\n%%EOF\n") + + mock_initial = { + "success": True, + "request_id": "cp-id", + "request_check_url": "https://api.datalab.to/api/v1/custom-processor/cp-id", + } + mock_result = { + "success": True, + "status": "complete", + "output_format": "markdown", + "markdown": "# Output", + } + + options = CustomProcessorOptions(pipeline_id="cp_abc12") + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + with patch.object(client, "_poll_result", new_callable=AsyncMock) as mock_poll: + mock_req.return_value = mock_initial + mock_poll.return_value = mock_result + result = await client.run_custom_processor(pdf_file, options=options) + assert result.success is True + + @pytest.mark.asyncio + async def test_run_custom_pipeline_emits_deprecation(self, temp_dir): + pdf_file = temp_dir / "test.pdf" + pdf_file.write_bytes(b"%PDF-1.4\n%Test PDF content\n%%EOF\n") + + mock_initial = { + "success": True, + "request_id": "cp-id", + "request_check_url": "https://api.datalab.to/api/v1/custom-processor/cp-id", + } + mock_result = { + "success": True, + "status": "complete", + "output_format": "markdown", + "markdown": "# Output", + } + + options = CustomProcessorOptions(pipeline_id="cp_abc12") + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + with patch.object(client, "_poll_result", new_callable=AsyncMock) as mock_poll: + mock_req.return_value = mock_initial + mock_poll.return_value = mock_result + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = await client.run_custom_pipeline(pdf_file, options=options) + assert len(w) == 1 + assert issubclass(w[0].category, DeprecationWarning) + assert "run_custom_processor" in str(w[0].message) + assert result.success is True + + +class TestExtractionSchemaCRUD: + """Test extraction schema CRUD methods""" + + @pytest.mark.asyncio + async def test_create_extraction_schema(self): + mock_response = { + "id": 1, "schema_id": "sch_abc123", "name": "Invoice Schema", + "description": "Extract invoice fields", + "schema_json": {"properties": {"total": {"type": "number"}}}, + "version": 1, "version_history": None, "archived": False, + "created": "2026-01-01T00:00:00", "updated": "2026-01-01T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.create_extraction_schema( + name="Invoice Schema", + schema_json={"properties": {"total": {"type": "number"}}}, + description="Extract invoice fields", + ) + assert isinstance(result, ExtractionSchema) + assert result.schema_id == "sch_abc123" + + @pytest.mark.asyncio + async def test_list_extraction_schemas(self): + mock_response = { + "schemas": [{ + "id": 1, "schema_id": "sch_abc123", "name": "Schema 1", + "description": None, "schema_json": {"properties": {}}, + "version": 1, "version_history": None, "archived": False, + "created": "2026-01-01T00:00:00", "updated": "2026-01-01T00:00:00", + }], + "total": 1, + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.list_extraction_schemas(limit=10) + assert result["total"] == 1 + assert isinstance(result["schemas"][0], ExtractionSchema) + + @pytest.mark.asyncio + async def test_delete_extraction_schema(self): + mock_response = { + "id": 1, "schema_id": "sch_abc123", "name": "Archived", + "description": None, "schema_json": {"properties": {}}, + "version": 1, "version_history": None, "archived": True, + "created": "2026-01-01T00:00:00", "updated": "2026-01-04T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.delete_extraction_schema("sch_abc123") + assert isinstance(result, ExtractionSchema) + assert result.archived is True + + +class TestPipelineCRUD: + """Test pipeline CRUD methods""" + + @pytest.mark.asyncio + async def test_create_pipeline(self): + mock_response = { + "id": 1, "pipeline_id": "pl_abc123", "name": None, + "steps": [{"type": "convert", "settings": {}}], + "is_saved": False, "archived": False, "active_version": 0, + "created": "2026-01-01T00:00:00", "updated": "2026-01-01T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.create_pipeline( + steps=[PipelineStep(type="convert")] + ) + assert isinstance(result, PipelineConfig) + assert result.pipeline_id == "pl_abc123" + mock_req.assert_awaited_once() + + @pytest.mark.asyncio + async def test_list_pipelines(self): + mock_response = { + "pipelines": [{ + "id": 1, "pipeline_id": "pl_abc123", "name": "My Pipeline", + "steps": [], "is_saved": True, "archived": False, + "active_version": 1, + "created": "2026-01-01T00:00:00", "updated": "2026-01-01T00:00:00", + }], + "total": 1, + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.list_pipelines() + assert result["total"] == 1 + assert isinstance(result["pipelines"][0], PipelineConfig) + + @pytest.mark.asyncio + async def test_get_pipeline(self): + mock_response = { + "id": 1, "pipeline_id": "pl_abc123", "name": "Test", + "steps": [{"type": "convert", "settings": {}}], + "is_saved": True, "archived": False, "active_version": 1, + "created": "2026-01-01T00:00:00", "updated": "2026-01-01T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.get_pipeline("pl_abc123") + assert result.pipeline_id == "pl_abc123" + mock_req.assert_awaited_once_with("GET", "/api/v1/pipelines/pl_abc123") + + +class TestPipelineExecution: + """Test pipeline execution methods""" + + @pytest.mark.asyncio + async def test_run_pipeline(self, temp_dir): + pdf_file = temp_dir / "test.pdf" + pdf_file.write_bytes(b"%PDF-1.4\n%Test PDF content\n%%EOF\n") + + mock_response = { + "execution_id": "pex_abc123", "pipeline_id": "pl_abc123", + "pipeline_version": 1, "status": "pending", + "steps": [ + {"step_index": 0, "step_type": "convert", "status": "pending"}, + ], + "created": "2026-01-01T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_submit_with_retry", new_callable=AsyncMock) as mock_submit: + mock_submit.return_value = mock_response + result = await client.run_pipeline("pl_abc123", file_path=pdf_file) + assert isinstance(result, PipelineExecution) + assert result.execution_id == "pex_abc123" + assert len(result.steps) == 1 + assert isinstance(result.steps[0], PipelineExecutionStepResult) + + @pytest.mark.asyncio + async def test_get_pipeline_execution(self): + mock_response = { + "execution_id": "pex_abc123", "pipeline_id": "pl_abc123", + "pipeline_version": 1, "status": "completed", + "steps": [ + {"step_index": 0, "step_type": "convert", "status": "completed", + "result_url": "/api/v1/pipelines/executions/pex_abc123/steps/0/result"}, + ], + "created": "2026-01-01T00:00:00", + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.get_pipeline_execution("pex_abc123") + assert result.status == "completed" + assert result.steps[0].result_url is not None + + +class TestCustomProcessorManagement: + """Test custom processor management methods""" + + @pytest.mark.asyncio + async def test_list_custom_processors(self): + mock_response = { + "pipelines": [{ + "processor_id": "cp_abc12", "name": "My Processor", + "status": "completed", "success": True, + "active_version": 1, "max_version": 2, + "iteration_in_progress": False, + "created_at": "2026-01-01T00:00:00", + }], + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.list_custom_processors() + assert len(result["processors"]) == 1 + assert isinstance(result["processors"][0], CustomProcessor) + assert result["processors"][0].processor_id == "cp_abc12" + + @pytest.mark.asyncio + async def test_list_custom_processor_versions(self): + mock_response = { + "versions": [ + {"version": 2, "request_description": "Add totals", "created_at": "2026-01-02T00:00:00", + "runtime": 45.2, "is_active": True}, + {"version": 1, "request_description": "Initial", "created_at": "2026-01-01T00:00:00", + "runtime": 30.0, "is_active": False}, + ], + } + + async with AsyncDatalabClient(api_key="test-key") as client: + with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: + mock_req.return_value = mock_response + result = await client.list_custom_processor_versions("cp_abc12") + assert len(result["versions"]) == 2 + assert isinstance(result["versions"][0], CustomProcessorVersion) + assert result["versions"][0].is_active is True + + +class TestPipelineStepModel: + """Test PipelineStep model""" + + def test_to_dict_minimal(self): + step = PipelineStep(type="convert") + d = step.to_dict() + assert d == {"type": "convert", "settings": {}} + + def test_to_dict_with_custom_processor(self): + step = PipelineStep( + type="custom", settings={"mode": "fast"}, + custom_processor_id="cp_abc12", eval_rubric_id=5, + ) + d = step.to_dict() + assert d["custom_processor_id"] == "cp_abc12" + assert d["eval_rubric_id"] == 5 From 6df52c159f48345b242dcb9dc0801ac2c786361a Mon Sep 17 00:00:00 2001 From: ashish uppala Date: Mon, 6 Apr 2026 12:02:59 -0400 Subject: [PATCH 2/3] PipelineStep -> PipelineProcessor --- datalab_sdk/__init__.py | 4 ++-- datalab_sdk/client.py | 14 +++++++------- datalab_sdk/models.py | 4 ++-- tests/test_client_methods.py | 12 ++++++------ 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/datalab_sdk/__init__.py b/datalab_sdk/__init__.py index eb148c0..57cab76 100644 --- a/datalab_sdk/__init__.py +++ b/datalab_sdk/__init__.py @@ -27,7 +27,7 @@ InputConfig, UploadedFileMetadata, ExtractionSchema, - PipelineStep, + PipelineProcessor, PipelineConfig, PipelineVersion, PipelineExecution, @@ -63,7 +63,7 @@ "InputConfig", "UploadedFileMetadata", "ExtractionSchema", - "PipelineStep", + "PipelineProcessor", "PipelineConfig", "PipelineVersion", "PipelineExecution", diff --git a/datalab_sdk/client.py b/datalab_sdk/client.py index 72df01a..f57e1bc 100644 --- a/datalab_sdk/client.py +++ b/datalab_sdk/client.py @@ -48,7 +48,7 @@ InputConfig, UploadedFileMetadata, ExtractionSchema, - PipelineStep, + PipelineProcessor, PipelineConfig, PipelineVersion, PipelineExecution, @@ -1734,13 +1734,13 @@ def _build_extraction_schema(data: Dict[str, Any]) -> ExtractionSchema: async def create_pipeline( self, - steps: list[PipelineStep], + steps: list[PipelineProcessor], ) -> PipelineConfig: """ Create a new pipeline Args: - steps: Ordered list of PipelineStep objects + steps: Ordered list of PipelineProcessor objects """ payload = {"steps": [s.to_dict() for s in steps]} response = await self._make_request("POST", "/api/v1/pipelines", json=payload) @@ -1786,14 +1786,14 @@ async def get_pipeline(self, pipeline_id: str) -> PipelineConfig: async def update_pipeline( self, pipeline_id: str, - steps: list[PipelineStep], + steps: list[PipelineProcessor], ) -> PipelineConfig: """ Update pipeline steps (auto-save path for draft edits) Args: pipeline_id: Pipeline ID string - steps: New ordered list of PipelineStep objects + steps: New ordered list of PipelineProcessor objects """ payload = {"steps": [s.to_dict() for s in steps]} response = await self._make_request("PUT", f"/api/v1/pipelines/{pipeline_id}", json=payload) @@ -2721,7 +2721,7 @@ def delete_extraction_schema(self, schema_id: str) -> ExtractionSchema: # --- Pipeline methods (sync) --- - def create_pipeline(self, steps: list[PipelineStep]) -> PipelineConfig: + def create_pipeline(self, steps: list[PipelineProcessor]) -> PipelineConfig: """Create a new pipeline (sync version)""" return self._run_async(self._async_client.create_pipeline(steps=steps)) @@ -2741,7 +2741,7 @@ def get_pipeline(self, pipeline_id: str) -> PipelineConfig: """Get a pipeline by ID (sync version)""" return self._run_async(self._async_client.get_pipeline(pipeline_id=pipeline_id)) - def update_pipeline(self, pipeline_id: str, steps: list[PipelineStep]) -> PipelineConfig: + def update_pipeline(self, pipeline_id: str, steps: list[PipelineProcessor]) -> PipelineConfig: """Update pipeline steps (sync version)""" return self._run_async( self._async_client.update_pipeline(pipeline_id=pipeline_id, steps=steps) diff --git a/datalab_sdk/models.py b/datalab_sdk/models.py index aaf57de..1757b7c 100644 --- a/datalab_sdk/models.py +++ b/datalab_sdk/models.py @@ -564,8 +564,8 @@ class ExtractionSchema: @dataclass -class PipelineStep: - """Configuration for a single pipeline step""" +class PipelineProcessor: + """Configuration for a single processor within a pipeline""" type: str # convert, extract, segment, custom settings: Dict[str, Any] = field(default_factory=dict) diff --git a/tests/test_client_methods.py b/tests/test_client_methods.py index ed6cf7e..d5896ba 100644 --- a/tests/test_client_methods.py +++ b/tests/test_client_methods.py @@ -20,7 +20,7 @@ TrackChangesOptions, OCROptions, ExtractionSchema, - PipelineStep, + PipelineProcessor, PipelineConfig, PipelineExecution, PipelineExecutionStepResult, @@ -1052,7 +1052,7 @@ async def test_create_pipeline(self): with patch.object(client, "_make_request", new_callable=AsyncMock) as mock_req: mock_req.return_value = mock_response result = await client.create_pipeline( - steps=[PipelineStep(type="convert")] + steps=[PipelineProcessor(type="convert")] ) assert isinstance(result, PipelineConfig) assert result.pipeline_id == "pl_abc123" @@ -1183,16 +1183,16 @@ async def test_list_custom_processor_versions(self): assert result["versions"][0].is_active is True -class TestPipelineStepModel: - """Test PipelineStep model""" +class TestPipelineProcessorModel: + """Test PipelineProcessor model""" def test_to_dict_minimal(self): - step = PipelineStep(type="convert") + step = PipelineProcessor(type="convert") d = step.to_dict() assert d == {"type": "convert", "settings": {}} def test_to_dict_with_custom_processor(self): - step = PipelineStep( + step = PipelineProcessor( type="custom", settings={"mode": "fast"}, custom_processor_id="cp_abc12", eval_rubric_id=5, ) From 399a5b646be58c5841a73b177da6eae9a2e7028b Mon Sep 17 00:00:00 2001 From: ashish uppala Date: Mon, 6 Apr 2026 12:55:54 -0400 Subject: [PATCH 3/3] bump version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 391547f..bd09d83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ readme = "README.md" license = "MIT" repository = "https://github.com/datalab-to/sdk" keywords = ["datalab", "sdk", "document-intelligence", "api"] -version = "0.4.0" +version = "0.5.0" description = "SDK for the Datalab document intelligence API" requires-python = ">=3.10" dependencies = [