Fix ParagraphPdfDocumentReader to reliably extract text from PDFs with imperfect outlines and coordinate edge cases

Signed-off-by: WOONBE <[email protected]>

Author: WOONBE

document-readers/pdf-reader/src/main/java/org/springframework/ai/reader/pdf/ParagraphPdfDocumentReader.java

@@ -128,29 +128,18 @@ public ParagraphPdfDocumentReader(Resource pdfResource, PdfDocumentReaderConfig
 	 */
 	@Override
 	public List<Document> get() {
-
 		var paragraphs = this.paragraphTextExtractor.flatten();
-
-		List<Document> documents = new ArrayList<>(paragraphs.size());
-
-		if (!CollectionUtils.isEmpty(paragraphs)) {
-			logger.info("Start processing paragraphs from PDF");
-			Iterator<Paragraph> itr = paragraphs.iterator();
-
-			var current = itr.next();
-
-			if (!itr.hasNext()) {
-				documents.add(toDocument(current, current));
-			}
-			else {
-				while (itr.hasNext()) {
-					var next = itr.next();
-					Document document = toDocument(current, next);
-					if (document != null && StringUtils.hasText(document.getText())) {
-						documents.add(toDocument(current, next));
-					}
-					current = next;
-				}
+		List<Document> documents = new ArrayList<>();
+		if (CollectionUtils.isEmpty(paragraphs)) {
+			return documents;
+		}
+		logger.info("Start processing paragraphs from PDF");
+		for (int i = 0; i < paragraphs.size(); i++) {
+			Paragraph from = paragraphs.get(i);
+			Paragraph to = (i + 1 < paragraphs.size()) ? paragraphs.get(i + 1) : from;
+			Document document = toDocument(from, to);
+			if (document != null && StringUtils.hasText(document.getText())) {
+				documents.add(document);
 			}
 		}
 		logger.info("End processing paragraphs from PDF");
@@ -174,7 +163,7 @@ protected Document toDocument(Paragraph from, Paragraph to) {
 	protected void addMetadata(Paragraph from, Paragraph to, Document document) {
 		document.getMetadata().put(METADATA_TITLE, from.title());
 		document.getMetadata().put(METADATA_START_PAGE, from.startPageNumber());
-		document.getMetadata().put(METADATA_END_PAGE, to.startPageNumber());
+		document.getMetadata().put(METADATA_END_PAGE, from.endPageNumber());
 		document.getMetadata().put(METADATA_LEVEL, from.level());
 		document.getMetadata().put(METADATA_FILE_NAME, this.resourceFileName);
 	}
@@ -192,7 +181,7 @@ public String getTextBetweenParagraphs(Paragraph fromParagraph, Paragraph toPara
 		int startPage = fromParagraph.startPageNumber() - 1;
 		int endPage = toParagraph.startPageNumber() - 1;
 
-		if (endPage < 0) {
+		if (fromParagraph == toParagraph || endPage < startPage) {
 			endPage = startPage;
 		}
 
@@ -206,39 +195,32 @@ public String getTextBetweenParagraphs(Paragraph fromParagraph, Paragraph toPara
 			for (int pageNumber = startPage; pageNumber <= endPage; pageNumber++) {
 
 				var page = this.document.getPage(pageNumber);
-
-				int fromPosition = fromParagraph.position();
-				int toPosition = toParagraph.position();
-
-				if (this.config.reversedParagraphPosition) {
-					fromPosition = (int) (page.getMediaBox().getHeight() - fromPosition);
-					toPosition = (int) (page.getMediaBox().getHeight() - toPosition);
+				float pageHeight = page.getMediaBox().getHeight();
+
+				int fromPos = fromParagraph.position();
+				int toPos = (fromParagraph != toParagraph) ? toParagraph.position() : 0;
+
+				int x = (int) page.getMediaBox().getLowerLeftX();
+				int w = (int) page.getMediaBox().getWidth();
+				int y, h;
+
+				if (pageNumber == startPage && pageNumber == endPage) {
+					y = toPos;
+					h = fromPos - toPos;
+				} else if (pageNumber == startPage) {
+					y = 0;
+					h = fromPos;
+				} else if (pageNumber == endPage) {
+					y = toPos;
+					h = (int) pageHeight - toPos;
+				} else {
+					y = 0;
+					h = (int) pageHeight;
 				}
 
-				int x0 = (int) page.getMediaBox().getLowerLeftX();
-				int xW = (int) page.getMediaBox().getWidth();
-
-				int y0 = (int) page.getMediaBox().getLowerLeftY();
-				int yW = (int) page.getMediaBox().getHeight();
-
-				if (pageNumber == startPage) {
-					y0 = fromPosition;
-					yW = (int) page.getMediaBox().getHeight() - y0;
-				}
-				if (pageNumber == endPage) {
-					yW = toPosition - y0;
-				}
-
-				if ((y0 + yW) == (int) page.getMediaBox().getHeight()) {
-					yW = yW - this.config.pageBottomMargin;
-				}
-
-				if (y0 == 0) {
-					y0 = y0 + this.config.pageTopMargin;
-					yW = yW - this.config.pageTopMargin;
-				}
+				if (h < 0) h = 0;
 
-				pdfTextStripper.addRegion("pdfPageRegion", new Rectangle(x0, y0, xW, yW));
+				pdfTextStripper.addRegion("pdfPageRegion", new Rectangle(x, y, w, h));
 				pdfTextStripper.extractRegions(page);
 				var text = pdfTextStripper.getTextForRegion("pdfPageRegion");
 				if (StringUtils.hasText(text)) {

document-readers/pdf-reader/src/main/java/org/springframework/ai/reader/pdf/config/ParagraphManager.java

@@ -117,7 +117,6 @@ protected Paragraph generateParagraphs(Paragraph parentParagraph, PDOutlineNode
 			if (nextSiblingNumber < 0) {
 				nextSiblingNumber = getPageNumber(current.getLastChild());
 			}
-
 			var paragraphPosition = (current.getDestination() instanceof PDPageXYZDestination)
 					? ((PDPageXYZDestination) current.getDestination()).getTop() : 0;
 

document-readers/pdf-reader/src/test/java/org/springframework/ai/reader/pdf/ParagraphPdfDocumentReaderTests.java

@@ -16,15 +16,43 @@
 
 package org.springframework.ai.reader.pdf;
 
+
+import org.apache.pdfbox.Loader;
+import org.apache.pdfbox.io.RandomAccessRead;
+import org.apache.pdfbox.pdmodel.PDDocument;
+import org.apache.pdfbox.pdmodel.interactive.documentnavigation.destination.PDDestination;
+import org.apache.pdfbox.pdmodel.interactive.documentnavigation.outline.PDDocumentOutline;
+import org.apache.pdfbox.pdmodel.interactive.documentnavigation.outline.PDOutlineItem;
 import org.junit.jupiter.api.Test;
 
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mockito;
+import org.mockito.Spy;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.ai.document.Document;
 import org.springframework.ai.reader.ExtractedTextFormatter;
+
+import org.springframework.ai.reader.pdf.config.ParagraphManager;
 import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.ClassPathResource;
+import org.springframework.core.io.Resource;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.List;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
 
 /**
  * @author Christian Tzolov
+ * @author Heonwoo Kim
  */
 public class ParagraphPdfDocumentReaderTests {
 
@@ -50,4 +78,7 @@ public void testPdfWithoutToc() {
 
 	}
 
+
+
+
 }