Add PDF annotation services in CLI

Readme.md

@@ -117,6 +117,16 @@ grobid_client [OPTIONS] SERVICE
 | `processCitationList`       | Parse citation strings            | Text files (one citation per line) |
 | `processCitationPatentST36` | Process patent citations          | XML ST36 format                    |
 | `processCitationPatentPDF`  | Process patent PDFs               | PDF files                          |
+| `referenceAnnotations`      | Reference annotations (JSON coordinates) | PDF files                   |
+| `annotatePDF`               | Annotated PDF with reference/citation annotations | PDF files          |
+| `citationPatentAnnotations` | Patent citation annotations (JSON coordinates) | Patent PDF files      |
+
+> [!NOTE]
+> The annotation services (`referenceAnnotations`, `annotatePDF`, `citationPatentAnnotations`) do not produce TEI XML.
+> `referenceAnnotations` and `citationPatentAnnotations` write JSON coordinate files (`*.references.json`,
+> `*.patent-citations.json`), while `annotatePDF` writes an annotated PDF (`*.annotated.pdf`). The `--json` and
+> `--markdown` conversion options do not apply to these services.
+> See [issue #79](https://github.com/grobidOrg/grobid-client-python/issues/79).
 
 #### Common Options
 
@@ -166,6 +176,12 @@ grobid_client --server https://grobid.example.com --input ~/citations.txt proces
 
 # Force reprocessing with sentence segmentation and JSON output
 grobid_client --input ~/docs --force --segmentSentences --json processFulltextDocument
+
+# Reference annotations as JSON coordinates (writes *.references.json)
+grobid_client --input ~/pdfs --output ~/annotations referenceAnnotations
+
+# Annotated PDF with reference/citation annotations (writes *.annotated.pdf)
+grobid_client --input ~/pdfs --output ~/annotated annotatePDF
 ```
 
 ### Python Library

grobid_client/grobid_client.py

@@ -39,6 +39,21 @@ def __init__(self, message="GROBID server is not available"):
 
 
 class GrobidClient(ApiClient):
+    # Default output descriptor for TEI-producing services: (Accept header,
+    # output file suffix, binary output). All the "process*" services return
+    # TEI XML as text.
+    DEFAULT_SERVICE_OUTPUT = ("text/plain", ".grobid.tei.xml", False)
+
+    # PDF annotation services return either JSON coordinates or an annotated
+    # (binary) PDF instead of TEI XML, so they need their own Accept header,
+    # output suffix and write mode.
+    # See https://github.com/grobidOrg/grobid-client-python/issues/79
+    SERVICE_OUTPUTS = {
+        "referenceAnnotations": ("application/json", ".references.json", False),
+        "citationPatentAnnotations": ("application/json", ".patent-citations.json", False),
+        "annotatePDF": ("application/pdf", ".annotated.pdf", True),
+    }
+
     # Default configuration values
     DEFAULT_CONFIG = {
         'grobid_server': 'http://localhost:8070',
@@ -112,6 +127,16 @@ def _set_config_params(self, params):
             if value is not None:
                 self.config[key] = value
 
+    def _service_output(self, service):
+        """Return the (Accept header, output suffix, is_binary) tuple for a service.
+
+        Annotation services (referenceAnnotations, annotatePDF,
+        citationPatentAnnotations) produce JSON or binary PDF output; every
+        other service produces TEI XML.
+        See https://github.com/grobidOrg/grobid-client-python/issues/79
+        """
+        return self.SERVICE_OUTPUTS.get(service, self.DEFAULT_SERVICE_OUTPUT)
+
     def _handle_server_busy_retry(self, file_path, retry_func, *args, **kwargs):
         """Handle server busy (503) retry logic."""
         self.logger.warning(f"Server busy (503), retrying {file_path} after {self.config['sleep_time']} seconds")
@@ -301,18 +326,18 @@ def _test_server_connection(self) -> Tuple[bool, int]:
             self.logger.error(error_msg)
             raise ServerUnavailableException(error_msg) from e
 
-    def _output_file_name(self, input_file, input_path, output):
+    def _output_file_name(self, input_file, input_path, output, suffix=".grobid.tei.xml"):
         # Use pathlib for consistent cross-platform path handling
         input_file_path = pathlib.Path(input_file)
 
         if output is not None:
             # Calculate relative path from input_path, then join with output directory
             input_path_abs = pathlib.Path(input_path).resolve()
             input_file_rel = input_file_path.resolve().relative_to(input_path_abs)
-            filename = pathlib.Path(output) / f"{input_file_rel.stem}.grobid.tei.xml"
+            filename = pathlib.Path(output) / f"{input_file_rel.stem}{suffix}"
         else:
             # Use the same directory as the input file
-            filename = input_file_path.parent / f"{input_file_path.stem}.grobid.tei.xml"
+            filename = input_file_path.parent / f"{input_file_path.stem}{suffix}"
 
         return str(filename)
 
@@ -476,20 +501,26 @@ def process_batch(
         error_count = 0
         skipped_count = 0
 
+        # Determine the output format for this service. Annotation services
+        # produce JSON or a binary PDF instead of TEI XML.
+        accept_header, output_suffix, binary_output = self._service_output(service)
+        # TEI -> JSON/Markdown conversion only makes sense for TEI services.
+        tei_service = service not in self.SERVICE_OUTPUTS
+
         # we use ThreadPoolExecutor and not ProcessPoolExecutor because it is an I/O intensive process
         with concurrent.futures.ThreadPoolExecutor(max_workers=n) as executor:
             # with concurrent.futures.ProcessPoolExecutor(max_workers=n) as executor:
             results = []
             for input_file in input_files:
-                # check if TEI file is already produced
-                filename = self._output_file_name(input_file, input_path, output)
+                # check if the output file is already produced
+                filename = self._output_file_name(input_file, input_path, output, output_suffix)
                 if not force and os.path.isfile(filename):
                     self.logger.info(
                         f"{filename} already exists, skipping... (use --force to reprocess pdf input files)")
                     skipped_count += 1
 
                     # Check if JSON output is needed but JSON file doesn't exist
-                    if json_output:
+                    if tei_service and json_output:
                         json_filename = filename.replace('.grobid.tei.xml', '.json')
                         # Expand ~ to home directory before checking file existence
                         json_filename_expanded = os.path.expanduser(json_filename)
@@ -509,7 +540,7 @@ def process_batch(
                                 self.logger.error(f"Failed to convert TEI to JSON for {filename}: {str(e)}")
 
                     # Check if Markdown output is needed but Markdown file doesn't exist
-                    if markdown_output:
+                    if tei_service and markdown_output:
                         markdown_filename = filename.replace('.grobid.tei.xml', '.md')
                         # Expand ~ to home directory before checking file existence
                         markdown_filename_expanded = os.path.expanduser(markdown_filename)
@@ -557,34 +588,44 @@ def process_batch(
 
         for r in concurrent.futures.as_completed(results):
             input_file, status, text = r.result()
-            filename = self._output_file_name(input_file, input_path, output)
+            filename = self._output_file_name(input_file, input_path, output, output_suffix)
 
             if status != 200 or text is None:
                 self.logger.error(f"Processing of {input_file} failed with error {status}: {text}")
                 error_count += 1
                 # writing error file with suffixed error code
                 try:
                     pathlib.Path(os.path.dirname(filename)).mkdir(parents=True, exist_ok=True)
-                    error_filename = filename.replace(".grobid.tei.xml", f"_{status}.txt")
+                    if filename.endswith(output_suffix):
+                        error_filename = filename[:-len(output_suffix)] + f"_{status}.txt"
+                    else:
+                        error_filename = filename + f"_{status}.txt"
                     with open(error_filename, 'w', encoding='utf8') as error_file:
+                        # Error responses are always text, even for binary services
                         if text is not None:
-                            error_file.write(text)
+                            error_file.write(text if isinstance(text, str) else text.decode('utf-8', 'replace'))
                         else:
                             error_file.write("")
                     self.logger.info(f"Error details written to {error_filename}")
                 except OSError as e:
                     self.logger.error(f"Failed to write error file {filename}: {str(e)}")
             else:
                 processed_count += 1
-                # writing TEI file
+                # writing output file
                 try:
                     pathlib.Path(os.path.dirname(filename)).mkdir(parents=True, exist_ok=True)
-                    with open(filename, 'w', encoding='utf8') as tei_file:
-                        tei_file.write(text)
-                    self.logger.debug(f"Successfully wrote TEI file: {filename}")
-
+                    if binary_output:
+                        # e.g. annotatePDF returns an annotated PDF (bytes)
+                        payload = text if isinstance(text, (bytes, bytearray)) else str(text).encode('utf-8')
+                        with open(filename, 'wb') as output_file_handle:
+                            output_file_handle.write(payload)
+                    else:
+                        with open(filename, 'w', encoding='utf8') as tei_file:
+                            tei_file.write(text)
+                    self.logger.debug(f"Successfully wrote output file: {filename}")
+
                     # Convert to JSON if requested
-                    if json_output:
+                    if tei_service and json_output:
                         try:
                             converter = TEI2LossyJSONConverter()
                             json_data = converter.convert_tei_file(filename, stream=False)
@@ -602,7 +643,7 @@ def process_batch(
                             self.logger.error(f"Failed to convert TEI to JSON for {filename}: {str(e)}")
 
                     # Convert to Markdown if requested
-                    if markdown_output:
+                    if tei_service and markdown_output:
                         try:
                             from .format.TEI2Markdown import TEI2MarkdownConverter
                             converter = TEI2MarkdownConverter()
@@ -621,7 +662,7 @@ def process_batch(
                             self.logger.error(f"Failed to convert TEI to Markdown for {filename}: {str(e)}")
 
                 except OSError as e:
-                    self.logger.error(f"Failed to write TEI XML file {filename}: {str(e)}")
+                    self.logger.error(f"Failed to write output file {filename}: {str(e)}")
 
         # Calculate batch statistics
         batch_runtime = time.time() - batch_start_time
@@ -688,8 +729,9 @@ def process_pdf(
             if end and end > 0:
                 the_data["end"] = str(end)
 
+            accept_header, _, binary_output = self._service_output(service)
             res, status = self.post(
-                url=the_url, files=files, data=the_data, headers={"Accept": "text/plain"},
+                url=the_url, files=files, data=the_data, headers={"Accept": accept_header},
                 timeout=self.config['timeout']
             )
 
@@ -711,8 +753,12 @@ def process_pdf(
                     end
                 )
 
+            # Binary services (e.g. annotatePDF) return raw bytes on success;
+            # error responses are always text.
+            if binary_output and status == 200:
+                return (pdf_file, status, res.content)
             return (pdf_file, status, res.text)
-        
+
         except IOError as e:
             self.logger.error(f"Failed to open PDF file {pdf_file}: {str(e)}")
             return (pdf_file, 400, f"Failed to open file: {str(e)}")
@@ -810,7 +856,12 @@ def main():
         "processReferences",
         "processCitationList",
         "processCitationPatentST36",
-        "processCitationPatentPDF"
+        "processCitationPatentPDF",
+        # PDF annotation services (see issue #79). These return JSON coordinates
+        # or an annotated (binary) PDF rather than TEI XML.
+        "referenceAnnotations",
+        "annotatePDF",
+        "citationPatentAnnotations"
     ]
 
     parser = argparse.ArgumentParser(description="Client for GROBID services")

tests/test_grobid_client.py

@@ -671,3 +671,92 @@ def test_get_server_url_edge_cases(self, mock_configure_logging, mock_test_serve
         result = client.get_server_url(service)
         expected = 'http://localhost:8070/api/processCitationPatentST36'
         assert result == expected
+
+
+class TestAnnotationServices:
+    """Tests for the PDF annotation services (issue #79)."""
+
+    def _client(self):
+        with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'):
+            with patch('grobid_client.grobid_client.GrobidClient._configure_logging'):
+                client = GrobidClient(check_server=False)
+        client.logger = Mock()
+        return client
+
+    def test_service_output_defaults_to_tei(self):
+        """Unknown/regular services fall back to the TEI output descriptor."""
+        client = self._client()
+        assert client._service_output('processFulltextDocument') == \
+            ('text/plain', '.grobid.tei.xml', False)
+
+    def test_service_output_annotation_services(self):
+        """Annotation services expose their own Accept header/suffix/binary flag."""
+        client = self._client()
+        assert client._service_output('referenceAnnotations') == \
+            ('application/json', '.references.json', False)
+        assert client._service_output('citationPatentAnnotations') == \
+            ('application/json', '.patent-citations.json', False)
+        assert client._service_output('annotatePDF') == \
+            ('application/pdf', '.annotated.pdf', True)
+
+    def test_output_file_name_with_custom_suffix(self):
+        """The output suffix is honoured when building the output file name."""
+        client = self._client()
+        result = client._output_file_name(
+            '/input/document.pdf', '/input', '/output', '.references.json')
+        assert result == '/output/document.references.json'
+
+    @patch('builtins.open', new_callable=mock_open)
+    @patch('grobid_client.grobid_client.GrobidClient.post')
+    def test_process_pdf_json_annotation_uses_json_accept(self, mock_post, mock_file):
+        """referenceAnnotations requests JSON and returns the response text."""
+        mock_response = Mock()
+        mock_response.text = '{"refs": []}'
+        mock_post.return_value = (mock_response, 200)
+
+        client = self._client()
+        result = client.process_pdf(
+            'referenceAnnotations', '/test/document.pdf',
+            generateIDs=False, consolidate_header=False, consolidate_citations=False,
+            include_raw_citations=False, include_raw_affiliations=False,
+            tei_coordinates=False, segment_sentences=False)
+
+        assert mock_post.call_args.kwargs['headers']['Accept'] == 'application/json'
+        assert result == ('/test/document.pdf', 200, '{"refs": []}')
+
+    @patch('builtins.open', new_callable=mock_open)
+    @patch('grobid_client.grobid_client.GrobidClient.post')
+    def test_process_pdf_annotate_returns_binary_content(self, mock_post, mock_file):
+        """annotatePDF requests a PDF and returns raw bytes on success."""
+        mock_response = Mock()
+        mock_response.content = b'%PDF-annotated'
+        mock_response.text = 'should not be used'
+        mock_post.return_value = (mock_response, 200)
+
+        client = self._client()
+        result = client.process_pdf(
+            'annotatePDF', '/test/document.pdf',
+            generateIDs=False, consolidate_header=False, consolidate_citations=False,
+            include_raw_citations=False, include_raw_affiliations=False,
+            tei_coordinates=False, segment_sentences=False)
+
+        assert mock_post.call_args.kwargs['headers']['Accept'] == 'application/pdf'
+        assert result == ('/test/document.pdf', 200, b'%PDF-annotated')
+
+    @patch('builtins.open', new_callable=mock_open)
+    @patch('grobid_client.grobid_client.GrobidClient.post')
+    def test_process_pdf_annotate_error_returns_text(self, mock_post, mock_file):
+        """On error, annotatePDF returns the response text (not bytes)."""
+        mock_response = Mock()
+        mock_response.content = b'binary'
+        mock_response.text = 'error detail'
+        mock_post.return_value = (mock_response, 500)
+
+        client = self._client()
+        result = client.process_pdf(
+            'annotatePDF', '/test/document.pdf',
+            generateIDs=False, consolidate_header=False, consolidate_citations=False,
+            include_raw_citations=False, include_raw_affiliations=False,
+            tei_coordinates=False, segment_sentences=False)
+
+        assert result == ('/test/document.pdf', 500, 'error detail')