From f15cc3c2766704be461fd5febbeaf12ea7b01ec4 Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Thu, 23 Jul 2026 11:26:12 +0200 Subject: [PATCH] Add PDF annotation services (referenceAnnotations, annotatePDF, citationPatentAnnotations) Wire the GROBID PDF annotation endpoints into the CLI and batch pipeline. Each service declares its own Accept header, output suffix and write mode: - referenceAnnotations -> application/json, *.references.json - citationPatentAnnotations -> application/json, *.patent-citations.json - annotatePDF -> application/pdf (binary), *.annotated.pdf The TEI -> JSON/Markdown conversion is skipped for these services. Documented in the Readme and covered by unit tests. Closes #79 --- Readme.md | 16 ++++++ grobid_client/grobid_client.py | 93 ++++++++++++++++++++++++++-------- tests/test_grobid_client.py | 89 ++++++++++++++++++++++++++++++++ 3 files changed, 177 insertions(+), 21 deletions(-) diff --git a/Readme.md b/Readme.md index 4a40f99..9adacad 100644 --- a/Readme.md +++ b/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 diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index 1af79d1..6985cfb 100644 --- a/grobid_client/grobid_client.py +++ b/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,7 +326,7 @@ 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) @@ -309,10 +334,10 @@ def _output_file_name(self, input_file, input_path, output): # 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,7 +588,7 @@ 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}") @@ -565,10 +596,14 @@ def process_batch( # 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}") @@ -576,15 +611,21 @@ def process_batch( 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") diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index b6ab324..83c7d3a 100644 --- a/tests/test_grobid_client.py +++ b/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')