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')