Remove and replace ExtractImagesResponse

Assisted-by: Codex

Author: datalogics-cgreen

src/pdfrest/client.py

@@ -59,7 +59,6 @@
 )
 from .models import (
     ConvertToMarkdownResponse,
-    ExtractImagesResponse,
     ExtractTextResponse,
     PdfRestErrorResponse,
     PdfRestFile,
@@ -2214,7 +2213,7 @@ def extract_images(
         extra_headers: AnyMapping | None = None,
         extra_body: Body | None = None,
         timeout: TimeoutTypes | None = None,
-    ) -> ExtractImagesResponse:
+    ) -> PdfRestFileBasedResponse:
         """Extract embedded images from a PDF."""
 
         payload: dict[str, Any] = {"files": file}
@@ -2223,40 +2222,15 @@ def extract_images(
         if output is not None:
             payload["output"] = output
 
-        validated_payload = ExtractImagesPayload.model_validate(payload)
-        request = self.prepare_request(
-            "POST",
-            "/extracted-images",
-            json_body=validated_payload.model_dump(
-                mode="json", by_alias=True, exclude_none=True, exclude_unset=True
-            ),
+        return self._post_file_operation(
+            endpoint="/extracted-images",
+            payload=payload,
+            payload_model=ExtractImagesPayload,
             extra_query=extra_query,
             extra_headers=extra_headers,
             extra_body=extra_body,
             timeout=timeout,
         )
-        raw_payload = self._send_request(request)
-        raw_response = PdfRestRawFileResponse.model_validate(raw_payload)
-        output_ids = raw_response.ids or []
-        output_files = [
-            self.fetch_file_info(
-                str(file_id),
-                extra_query=extra_query,
-                extra_headers=extra_headers,
-                timeout=timeout,
-            )
-            for file_id in output_ids
-        ]
-        input_id = raw_response.input_id[0] if raw_response.input_id else ""
-        return ExtractImagesResponse.model_validate(
-            {
-                "input_id": input_id,
-                "output_files": [
-                    file.model_dump(mode="json", by_alias=True) for file in output_files
-                ],
-                "warning": raw_response.warning,
-            }
-        )
 
     def extract_text(
         self,
@@ -2954,7 +2928,7 @@ async def extract_images(
         extra_headers: AnyMapping | None = None,
         extra_body: Body | None = None,
         timeout: TimeoutTypes | None = None,
-    ) -> ExtractImagesResponse:
+    ) -> PdfRestFileBasedResponse:
         """Extract embedded images from a PDF."""
 
         payload: dict[str, Any] = {"files": file}
@@ -2963,47 +2937,15 @@ async def extract_images(
         if output is not None:
             payload["output"] = output
 
-        validated_payload = ExtractImagesPayload.model_validate(payload)
-        request = self.prepare_request(
-            "POST",
-            "/extracted-images",
-            json_body=validated_payload.model_dump(
-                mode="json", by_alias=True, exclude_none=True, exclude_unset=True
-            ),
+        return await self._post_file_operation(
+            endpoint="/extracted-images",
+            payload=payload,
+            payload_model=ExtractImagesPayload,
             extra_query=extra_query,
             extra_headers=extra_headers,
             extra_body=extra_body,
             timeout=timeout,
         )
-        raw_payload = await self._send_request(request)
-        raw_response = PdfRestRawFileResponse.model_validate(raw_payload)
-        output_ids = raw_response.ids or []
-        semaphore = asyncio.Semaphore(DEFAULT_FILE_INFO_CONCURRENCY)
-
-        async def fetch(file_id: str) -> PdfRestFile:
-            async with semaphore:
-                return await self.fetch_file_info(
-                    file_id,
-                    extra_query=extra_query,
-                    extra_headers=extra_headers,
-                    timeout=timeout,
-                )
-
-        output_files: list[PdfRestFile] = []
-        if output_ids:
-            output_files = list(
-                await asyncio.gather(*(fetch(fid) for fid in output_ids))
-            )
-        input_id = raw_response.input_id[0] if raw_response.input_id else ""
-        return ExtractImagesResponse.model_validate(
-            {
-                "input_id": input_id,
-                "output_files": [
-                    file.model_dump(mode="json", by_alias=True) for file in output_files
-                ],
-                "warning": raw_response.warning,
-            }
-        )
 
     async def extract_text(
         self,

src/pdfrest/models/__init__.py

@@ -1,6 +1,5 @@
 from .public import (
     ConvertToMarkdownResponse,
-    ExtractImagesResponse,
     ExtractTextResponse,
     PdfRestErrorResponse,
     PdfRestFile,
@@ -14,7 +13,6 @@
 
 __all__ = [
     "ConvertToMarkdownResponse",
-    "ExtractImagesResponse",
     "ExtractTextResponse",
     "PdfRestErrorResponse",
     "PdfRestFile",

src/pdfrest/models/public.py

@@ -21,7 +21,6 @@
 
 __all__ = (
     "ConvertToMarkdownResponse",
-    "ExtractImagesResponse",
     "ExtractTextResponse",
     "PdfRestErrorResponse",
     "PdfRestFile",
@@ -379,31 +378,6 @@ class TranslatePdfTextResponse(BaseModel):
     ] = None
 
 
-class ExtractImagesResponse(BaseModel):
-    """Response returned by the extracted-images tool."""
-
-    model_config = ConfigDict(extra="allow")
-
-    input_id: Annotated[
-        PdfRestFileID,
-        Field(
-            validation_alias=AliasChoices("input_id", "inputId"),
-            description="The id of the input file.",
-        ),
-    ]
-    output_files: Annotated[
-        list[PdfRestFile],
-        Field(
-            description="The list of extracted image files.",
-            validation_alias=AliasChoices("output_files", "outputFiles"),
-        ),
-    ]
-    warning: Annotated[
-        str | None,
-        Field(description="A warning that was generated during extraction."),
-    ] = None
-
-
 class ExtractTextResponse(BaseModel):
     """Response returned by the extracted-text tool."""
 

tests/live/test_live_extract_images.py

@@ -3,7 +3,7 @@
 import pytest
 
 from pdfrest import PdfRestApiError, PdfRestClient
-from pdfrest.models import ExtractImagesResponse
+from pdfrest.models import PdfRestFileBasedResponse
 
 from ..resources import get_test_resource_path
 
@@ -20,7 +20,7 @@ def test_live_extract_images_success(
         uploaded = client.files.create_from_paths([resource])[0]
         response = client.extract_images(uploaded)
 
-    assert isinstance(response, ExtractImagesResponse)
+    assert isinstance(response, PdfRestFileBasedResponse)
     assert response.output_files
     assert response.input_id == uploaded.id
 

tests/test_extract_images.py

@@ -7,7 +7,7 @@
 from pydantic import ValidationError
 
 from pdfrest import AsyncPdfRestClient, PdfRestClient
-from pdfrest.models import ExtractImagesResponse, PdfRestFile, PdfRestFileID
+from pdfrest.models import PdfRestFile, PdfRestFileBasedResponse, PdfRestFileID
 from pdfrest.models._internal import ExtractImagesPayload
 
 from .graphics_test_helpers import ASYNC_API_KEY, VALID_API_KEY, make_pdf_file
@@ -98,7 +98,7 @@ def handler(request: httpx.Request) -> httpx.Response:
         response = client.extract_images(input_file, pages=["1-3"], output="images")
 
     assert seen == {"post": 1, "get": 2}
-    assert isinstance(response, ExtractImagesResponse)
+    assert isinstance(response, PdfRestFileBasedResponse)
     assert len(response.output_files) == 2
     assert response.input_id == input_file.id
 
@@ -152,7 +152,7 @@ def handler(request: httpx.Request) -> httpx.Response:
             timeout=0.3,
         )
 
-    assert isinstance(response, ExtractImagesResponse)
+    assert isinstance(response, PdfRestFileBasedResponse)
     assert len(response.output_files) == 1
     timeout_value = captured_timeout["value"]
     assert timeout_value is not None
@@ -206,6 +206,6 @@ def handler(request: httpx.Request) -> httpx.Response:
         response = await client.extract_images(input_file)
 
     assert seen == {"post": 1, "get": 1}
-    assert isinstance(response, ExtractImagesResponse)
+    assert isinstance(response, PdfRestFileBasedResponse)
     assert len(response.output_files) == 1
     assert response.input_id == input_file.id