From 0e3a68c3c1e9da1c1f38f3c1dec58e9220eb0f4a Mon Sep 17 00:00:00 2001 From: let5sne Date: Sun, 12 Jul 2026 21:33:18 +0800 Subject: [PATCH] feat(personal): process scanned PDFs in OCR batches --- .../service/PersonalIngestionWorker.java | 54 ++- .../service/PersonalPdfOcrService.java | 459 ++++++++++++++++++ .../personal/PersonalIngestionWorkerTest.java | 22 + .../personal/PersonalPdfOcrServiceTest.java | 199 ++++++++ 4 files changed, 729 insertions(+), 5 deletions(-) create mode 100644 backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalPdfOcrService.java create mode 100644 backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalPdfOcrServiceTest.java diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalIngestionWorker.java b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalIngestionWorker.java index 5f8edf24..67106064 100644 --- a/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalIngestionWorker.java +++ b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalIngestionWorker.java @@ -35,6 +35,7 @@ public class PersonalIngestionWorker { private final TransactionTemplate transactionTemplate; private final StoredObjectReader objectReader; private final TenantRunner tenantRunner; + private final PdfOcrHandler pdfOcrHandler; private final long maxInputBytes; private final int chunkSize; private final int chunkOverlap; @@ -44,22 +45,25 @@ public class PersonalIngestionWorker { @Autowired public PersonalIngestionWorker(JdbcTemplate jdbcTemplate, ISysOssService ossService, KnowledgeDocumentParser parser, PersonalKnowledgeProperties properties, - PlatformTransactionManager transactionManager) { + PlatformTransactionManager transactionManager, + PersonalPdfOcrService pdfOcrService) { this(jdbcTemplate, parser, new TransactionTemplate(transactionManager), defaultReader(ossService, PersonalIngestionWorker::ossClient), configuredMaxBytes(properties), properties.getChunkSize(), properties.getChunkOverlap(), properties.getParsingLeaseMinutes(), - properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant); + properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant, pdfOcrService::enqueue); } private PersonalIngestionWorker(JdbcTemplate jdbcTemplate, KnowledgeDocumentParser parser, TransactionTemplate transactionTemplate, StoredObjectReader objectReader, long maxInputBytes, int chunkSize, int chunkOverlap, - int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner) { + int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner, + PdfOcrHandler pdfOcrHandler) { this.jdbcTemplate = jdbcTemplate; this.parser = parser; this.transactionTemplate = transactionTemplate; this.objectReader = objectReader; this.tenantRunner = tenantRunner; + this.pdfOcrHandler = pdfOcrHandler; this.maxInputBytes = maxInputBytes; if (chunkSize <= 0 || chunkOverlap < 0 || chunkOverlap >= chunkSize || parsingLeaseMinutes <= 0 || maxParseAttempts <= 0) { @@ -76,7 +80,8 @@ public class PersonalIngestionWorker { TransactionTemplate transactionTemplate, StoredObjectReader objectReader) { return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader, - 20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute()); + 20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute(), + unsupportedPdfOcr()); } public static PersonalIngestionWorker forTest(JdbcTemplate jdbcTemplate, ISysOssService ossService, @@ -84,7 +89,17 @@ public class PersonalIngestionWorker { TransactionTemplate transactionTemplate, StoredObjectReader objectReader, TenantRunner tenantRunner) { return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader, - 20L * 1024 * 1024, 800, 120, 15, 3, tenantRunner); + 20L * 1024 * 1024, 800, 120, 15, 3, tenantRunner, unsupportedPdfOcr()); + } + + public static PersonalIngestionWorker forTestWithPdfOcr(JdbcTemplate jdbcTemplate, ISysOssService ossService, + KnowledgeDocumentParser parser, + TransactionTemplate transactionTemplate, + StoredObjectReader objectReader, + PdfOcrHandler pdfOcrHandler) { + return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader, + 20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute(), + pdfOcrHandler); } public static StoredObjectReader objectReaderForTest(ISysOssService ossService, @@ -173,6 +188,12 @@ public class PersonalIngestionWorker { ParsedDocument document = parser.parse(stored.fileName(), item.mimeType(), stored.bytes()); List chunks = document.chunks(chunkSize, chunkOverlap); if (chunks.isEmpty()) { + if (isPdf(item, stored)) { + pdfOcrHandler.enqueue(new PersonalPdfOcrService.OcrItem(item.id(), item.tenantId(), + item.spaceId(), item.ownerUserId(), item.ossId(), item.tagsJson(), attemptVersion), + stored.bytes()); + return true; + } throw new KnowledgeDocumentParser.ParseException( KnowledgeDocumentParser.Failure.EMPTY, "document contains no text"); } @@ -308,6 +329,12 @@ public class PersonalIngestionWorker { && "PERSONAL_OSS_NOT_PRIVATE".equals(serviceError.getMessage())) { return new Failure("PERSONAL_OSS_NOT_PRIVATE", "个人资料存储策略不可用"); } + if (candidate instanceof PersonalPdfPageRenderer.PdfPageLimitException) { + return new Failure("PERSONAL_PDF_PAGE_LIMIT", "文件超过200页，请拆分后重新上传"); + } + if (candidate instanceof PersonalPdfPageRenderer.PdfRenderException renderError) { + return new Failure(renderError.code(), "扫描PDF预处理失败，请检查文件后重试"); + } candidate = candidate.getCause(); } return new Failure("PERSONAL_PARSE_FAILED", "资料处理失败，请稍后重试"); @@ -351,6 +378,18 @@ public class PersonalIngestionWorker { return "personal/" + item.tenantId() + "/" + item.ownerUserId() + "/" + item.id() + "/"; } + private static boolean isPdf(Item item, StoredObject stored) { + return "application/pdf".equalsIgnoreCase(item.mimeType()) + || stored.fileName().toLowerCase(java.util.Locale.ROOT).endsWith(".pdf"); + } + + private static PdfOcrHandler unsupportedPdfOcr() { + return (item, bytes) -> { + throw new KnowledgeDocumentParser.ParseException( + KnowledgeDocumentParser.Failure.EMPTY, "document contains no text"); + }; + } + @FunctionalInterface public interface StoredObjectReader { StoredObject read(long ossId, String expectedPrefix, long ownerUserId, long maxBytes) throws Exception; @@ -371,6 +410,11 @@ public class PersonalIngestionWorker { OssClient get(String configKey); } + @FunctionalInterface + public interface PdfOcrHandler { + void enqueue(PersonalPdfOcrService.OcrItem item, byte[] pdfBytes); + } + public record StoredObject(String fileName, byte[] bytes) { public StoredObject { bytes = bytes == null ? new byte[0] : bytes.clone(); diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalPdfOcrService.java b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalPdfOcrService.java new file mode 100644 index 00000000..72208d3e --- /dev/null +++ b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/personal/service/PersonalPdfOcrService.java @@ -0,0 +1,459 @@ +package org.dromara.aihr.personal.service; + +import org.dromara.aihr.knowledge.parse.ParsedDocument; +import org.dromara.aihr.personal.support.PersonalKnowledgeProperties; +import org.dromara.common.core.exception.ServiceException; +import org.dromara.common.oss.core.OssClient; +import org.dromara.common.oss.enums.AccessPolicyType; +import org.dromara.common.oss.factory.OssFactory; +import org.dromara.common.tenant.helper.TenantHelper; +import org.dromara.system.domain.vo.SysOssVo; +import org.dromara.system.service.ISysOssService; +import org.springframework.jdbc.core.BatchPreparedStatementSetter; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.scheduling.annotation.Scheduled; +import org.springframework.stereotype.Service; +import org.springframework.transaction.PlatformTransactionManager; +import org.springframework.transaction.support.TransactionTemplate; + +import java.io.InputStream; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.StringJoiner; + +@Service +public class PersonalPdfOcrService { + + private final JdbcTemplate jdbcTemplate; + private final PersonalPdfPageRenderer renderer; + private final PersonalVisionOcrService visionOcr; + private final TransactionTemplate transactions; + private final OcrSourceReader sourceReader; + + public PersonalPdfOcrService(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer, + PersonalVisionOcrService visionOcr, PlatformTransactionManager transactionManager, + ISysOssService ossService, PersonalKnowledgeProperties properties) { + this(jdbcTemplate, renderer, visionOcr, new TransactionTemplate(transactionManager), + defaultSourceReader(ossService, maxBytes(properties))); + } + + private PersonalPdfOcrService(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer, + PersonalVisionOcrService visionOcr, TransactionTemplate transactions, + OcrSourceReader sourceReader) { + this.jdbcTemplate = jdbcTemplate; + this.renderer = renderer; + this.visionOcr = visionOcr; + this.transactions = transactions; + this.sourceReader = sourceReader; + } + + public static PersonalPdfOcrService forTest(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer, + PersonalVisionOcrService visionOcr, + TransactionTemplate transactions, + OcrSourceReader sourceReader) { + return new PersonalPdfOcrService(jdbcTemplate, renderer, visionOcr, transactions, sourceReader); + } + + public void enqueue(OcrItem item, byte[] pdfBytes) { + int totalPages = renderer.pageCount(pdfBytes); + transactions.execute(status -> { + jdbcTemplate.update(""" + insert into aihr_personal_ocr_job + (tenant_id, owner_user_id, item_id, status, total_pages, processed_pages, + success_pages, failed_pages, failed_page_numbers_json, next_page, + attempt_count, create_time, update_time) + values (?, ?, ?, 'PENDING', ?, 0, 0, 0, json_array(), 1, 0, now(), now()) + on duplicate key update status = 'PENDING', total_pages = values(total_pages), + processed_pages = 0, success_pages = 0, failed_pages = 0, + failed_page_numbers_json = json_array(), next_page = 1, + attempt_count = 0, last_error_code = null, lease_until = null, update_time = now() + """, item.tenantId(), item.ownerUserId(), item.id(), totalPages); + jdbcTemplate.update(""" + delete from aihr_personal_ocr_page + where tenant_id = ? and owner_user_id = ? and item_id = ? + """, item.tenantId(), item.ownerUserId(), item.id()); + jdbcTemplate.batchUpdate(""" + insert into aihr_personal_ocr_page + (tenant_id, owner_user_id, item_id, page_number, status, + attempt_count, create_time, update_time) + values (?, ?, ?, ?, 'PENDING', 0, now(), now()) + """, new BatchPreparedStatementSetter() { + @Override + public void setValues(PreparedStatement statement, int index) throws SQLException { + statement.setString(1, item.tenantId()); + statement.setLong(2, item.ownerUserId()); + statement.setLong(3, item.id()); + statement.setInt(4, index + 1); + } + + @Override + public int getBatchSize() { + return totalPages; + } + }); + return null; + }); + } + + public static List batchPageNumbers(int firstPage, int totalPages) { + if (firstPage < 1 || totalPages < firstPage) { + return List.of(); + } + int end = Math.min(totalPages, firstPage + PersonalPdfPageRenderer.BATCH_SIZE - 1); + List pages = new ArrayList<>(end - firstPage + 1); + for (int page = firstPage; page <= end; page++) { + pages.add(page); + } + return List.copyOf(pages); + } + + @Scheduled(fixedDelayString = "${aihr.personal.ocr-delay-ms:2000}", scheduler = "personalTaskScheduler") + public void poll() { + recoverExpiredJobs(); + processNextBatch(); + } + + public void recoverExpiredJobs() { + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set status = 'RETRY', lease_until = null, last_error_code = 'PERSONAL_OCR_LEASE_EXPIRED', + update_time = now() + where status = 'RUNNING' and lease_until < now() + """); + } + + public boolean processNextBatch() { + List> rows = jdbcTemplate.queryForList(""" + select j.id, j.tenant_id, j.owner_user_id, j.item_id, j.status, j.attempt_count, + j.total_pages, j.next_page, + i.space_id, i.oss_id, i.tags_json + from aihr_personal_ocr_job j + join aihr_personal_item i on i.id = j.item_id + and binary i.tenant_id = binary j.tenant_id and i.owner_user_id = j.owner_user_id + where j.status in ('PENDING','RETRY') and i.status = 'PARSING' + order by j.update_time, j.id + limit 1 + """); + if (rows.isEmpty()) { + return false; + } + OcrJob job = job(rows.get(0)); + int claimed = jdbcTemplate.update(""" + update aihr_personal_ocr_job + set status = 'RUNNING', attempt_count = attempt_count + 1, + lease_until = date_add(now(), interval 15 minute), update_time = now() + where id = ? and status in ('PENDING','RETRY') + """, job.id()); + if (claimed != 1) { + return false; + } + List pageNumbers = retryJob(job) ? retryPageNumbers(job) + : batchPageNumbers(job.nextPage(), job.totalPages()); + int success = 0; + int failed = 0; + try { + byte[] pdfBytes = sourceReader.read(job); + List rendered = renderPages(pdfBytes, pageNumbers, retryJob(job)); + for (PersonalPdfPageRenderer.RenderedPage page : rendered) { + int pageNumber = page.pageIndex() + 1; + try { + String text = visionOcr.recognize(page.bytes(), page.mimeType(), pageNumber); + if (text.isBlank()) { + throw new PersonalVisionOcrService.OcrUnavailableException("PERSONAL_OCR_PAGE_EMPTY"); + } + jdbcTemplate.update(""" + update aihr_personal_ocr_page + set status = 'SUCCESS', ocr_text = ?, attempt_count = attempt_count + 1, + last_error_code = null, update_time = now() + where tenant_id = ? and owner_user_id = ? and item_id = ? and page_number = ? + """, text, job.tenantId(), job.ownerUserId(), job.itemId(), pageNumber); + success++; + } catch (PersonalVisionOcrService.OcrUnavailableException exception) { + jdbcTemplate.update(""" + update aihr_personal_ocr_page + set status = 'FAILED', ocr_text = null, attempt_count = attempt_count + 1, + last_error_code = ?, update_time = now() + where tenant_id = ? and owner_user_id = ? and item_id = ? and page_number = ? + """, exception.code(), job.tenantId(), job.ownerUserId(), job.itemId(), pageNumber); + failed++; + } + } + int nextPage = job.nextPage() + rendered.size(); + if (retryJob(job)) { + Integer remaining = jdbcTemplate.queryForObject(""" + select count(*) from aihr_personal_ocr_page + where tenant_id = ? and owner_user_id = ? and item_id = ? + and status = 'FAILED' and attempt_count < 3 + """, Integer.class, job.tenantId(), job.ownerUserId(), job.itemId()); + if (remaining != null && remaining > 0) { + refreshRetryState(job); + } else { + publishTerminal(job); + } + } else if (nextPage <= job.totalPages()) { + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set processed_pages = processed_pages + ?, success_pages = success_pages + ?, + failed_pages = failed_pages + ?, next_page = ?, status = 'PENDING', + lease_until = null, update_time = now() + where id = ? and status = 'RUNNING' + """, rendered.size(), success, failed, nextPage, job.id()); + } else if (failed > 0) { + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set processed_pages = processed_pages + ?, success_pages = success_pages + ?, + failed_pages = failed_pages + ?, next_page = total_pages + 1, + status = 'RETRY', lease_until = null, update_time = now() + where id = ? and status = 'RUNNING' + """, rendered.size(), success, failed, job.id()); + } else { + publishTerminal(job); + } + } catch (Exception exception) { + if (job.attemptCount() + 1 >= 3) { + transactions.execute(status -> { + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set status = 'FAILED', last_error_code = 'PERSONAL_OCR_BATCH_FAILED', + lease_until = null, update_time = now() + where id = ? and status = 'RUNNING' + """, job.id()); + jdbcTemplate.update(""" + update aihr_personal_item + set status = 'FAILED', error_code = 'PERSONAL_OCR_BATCH_FAILED', + error_message = '扫描PDF识别任务连续失败，请稍后重试', update_time = now() + where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING' + """, job.tenantId(), job.ownerUserId(), job.itemId()); + return null; + }); + } else { + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set status = 'RETRY', last_error_code = 'PERSONAL_OCR_BATCH_FAILED', + lease_until = null, update_time = now() + where id = ? and status = 'RUNNING' + """, job.id()); + } + } + return true; + } + + private List retryPageNumbers(OcrJob job) { + return jdbcTemplate.queryForList(""" + select page_number from aihr_personal_ocr_page + where tenant_id = ? and owner_user_id = ? and item_id = ? + and status = 'FAILED' and attempt_count < 3 + order by page_number limit ? + """, Integer.class, job.tenantId(), job.ownerUserId(), job.itemId(), + PersonalPdfPageRenderer.BATCH_SIZE); + } + + private List renderPages(byte[] pdfBytes, List pageNumbers, + boolean retry) { + if (!retry) { + if (pageNumbers.isEmpty()) return List.of(); + return renderer.render(pdfBytes, pageNumbers.get(0) - 1, pageNumbers.size()); + } + List rendered = new ArrayList<>(pageNumbers.size()); + for (int pageNumber : pageNumbers) { + List page = renderer.render(pdfBytes, pageNumber - 1, 1); + if (!page.isEmpty()) rendered.add(page.get(0)); + } + return List.copyOf(rendered); + } + + private void refreshRetryState(OcrJob job) { + jdbcTemplate.update(""" + update aihr_personal_ocr_job j + set j.processed_pages = (select count(*) from aihr_personal_ocr_page p + where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id + and p.item_id = j.item_id and (p.status = 'SUCCESS' or p.attempt_count >= 3)), + j.success_pages = (select count(*) from aihr_personal_ocr_page p + where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id + and p.item_id = j.item_id and p.status = 'SUCCESS'), + j.failed_pages = (select count(*) from aihr_personal_ocr_page p + where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id + and p.item_id = j.item_id and p.status = 'FAILED'), + j.status = 'RETRY', j.lease_until = null, j.update_time = now() + where j.id = ? and j.status = 'RUNNING' + """, job.id()); + } + + private static boolean retryJob(OcrJob job) { + return "RETRY".equals(job.status()); + } + + private void publishTerminal(OcrJob job) { + List pageTexts = jdbcTemplate.queryForList(""" + select ocr_text from aihr_personal_ocr_page + where tenant_id = ? and owner_user_id = ? and item_id = ? and status = 'SUCCESS' + order by page_number + """, String.class, job.tenantId(), job.ownerUserId(), job.itemId()); + List failedPages = jdbcTemplate.queryForList(""" + select page_number from aihr_personal_ocr_page + where tenant_id = ? and owner_user_id = ? and item_id = ? and status = 'FAILED' + order by page_number + """, Integer.class, job.tenantId(), job.ownerUserId(), job.itemId()); + transactions.execute(status -> { + if (pageTexts.isEmpty()) { + jdbcTemplate.update(""" + update aihr_personal_item + set status = 'FAILED', error_code = 'PERSONAL_OCR_EMPTY', + error_message = '扫描PDF未识别到可用文字', update_time = now() + where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING' + """, job.tenantId(), job.ownerUserId(), job.itemId()); + } else { + publishFragments(job, String.join("\n\n", pageTexts), failedPages); + } + jdbcTemplate.update(""" + update aihr_personal_ocr_job + set processed_pages = ?, success_pages = ?, failed_pages = ?, + failed_page_numbers_json = ?, next_page = total_pages + 1, + status = 'COMPLETED', lease_until = null, update_time = now() + where id = ? and status = 'RUNNING' + """, job.totalPages(), pageTexts.size(), failedPages.size(), jsonNumbers(failedPages), job.id()); + return null; + }); + } + + private void publishFragments(OcrJob job, String text, List failedPages) { + List chunks = new ParsedDocument(text, "application/pdf", Map.of()).chunks(800, 120); + jdbcTemplate.update(""" + delete from aihr_personal_fragment + where tenant_id = ? and owner_user_id = ? and item_id = ? + """, job.tenantId(), job.ownerUserId(), job.itemId()); + jdbcTemplate.batchUpdate(""" + insert into aihr_personal_fragment + (tenant_id, space_id, owner_user_id, item_id, idx, content, token_count, + embedding_json, embedding_model, embedding_time, create_time) + values (?, ?, ?, ?, ?, ?, ?, null, null, null, now()) + """, new BatchPreparedStatementSetter() { + @Override + public void setValues(PreparedStatement statement, int index) throws SQLException { + String chunk = chunks.get(index); + statement.setString(1, job.tenantId()); + statement.setLong(2, job.spaceId()); + statement.setLong(3, job.ownerUserId()); + statement.setLong(4, job.itemId()); + statement.setInt(5, index); + statement.setString(6, chunk); + statement.setInt(7, Math.max(1, (chunk.codePointCount(0, chunk.length()) + 1) / 2)); + } + + @Override + public int getBatchSize() { + return chunks.size(); + } + }); + if (failedPages.isEmpty()) { + jdbcTemplate.update(""" + update aihr_personal_item + set status = 'READY', parsed_at = now(), summary = ?, tags_json = ?, + error_code = null, error_message = null, update_time = now() + where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING' + """, summary(text), job.tagsJson(), job.tenantId(), job.ownerUserId(), job.itemId()); + } else { + jdbcTemplate.update(""" + update aihr_personal_item + set status = 'READY', parsed_at = now(), summary = ?, tags_json = ?, + error_code = 'PERSONAL_OCR_PARTIAL', error_message = ?, update_time = now() + where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING' + """, summary(text), job.tagsJson(), "已收录，" + failedPages.size() + "页识别失败", + job.tenantId(), job.ownerUserId(), job.itemId()); + } + } + + private static String jsonNumbers(List values) { + StringJoiner joiner = new StringJoiner(",", "[", "]"); + values.forEach(value -> joiner.add(String.valueOf(value))); + return joiner.toString(); + } + + private static String summary(String text) { + int[] codePoints = text.codePoints().limit(300).toArray(); + return new String(codePoints, 0, codePoints.length).trim(); + } + + private static OcrJob job(Map row) { + return new OcrJob(number(row, "id"), String.valueOf(row.get("tenant_id")), + number(row, "owner_user_id"), number(row, "item_id"), number(row, "space_id"), + number(row, "oss_id"), row.get("tags_json") == null ? "[]" : String.valueOf(row.get("tags_json")), + String.valueOf(row.get("status")), intValue(row, "attempt_count", 0), + Math.toIntExact(number(row, "total_pages")), + Math.toIntExact(number(row, "next_page"))); + } + + private static long number(Map row, String key) { + Object value = row.get(key); + if (!(value instanceof Number number)) { + throw new IllegalStateException("personal OCR metadata is incomplete"); + } + return number.longValue(); + } + + private static int intValue(Map row, String key, int fallback) { + Object value = row.get(key); + return value instanceof Number number ? Math.toIntExact(number.longValue()) : fallback; + } + + private static OcrSourceReader defaultSourceReader(ISysOssService ossService, long maxBytes) { + return job -> runInTenant(job.tenantId(), () -> { + SysOssVo object = ossService.getById(job.ossId()); + if (object == null || object.getCreateBy() == null || object.getCreateBy() != job.ownerUserId() + || object.getService() == null || object.getFileName() == null + || !object.getFileName().startsWith(ownerPrefix(job))) { + throw new ServiceException("PERSONAL_OCR_SOURCE_UNAVAILABLE"); + } + OssClient client = OssFactory.instance(object.getService()); + if (client.getAccessPolicy() != AccessPolicyType.PRIVATE) { + throw new ServiceException("PERSONAL_OSS_NOT_PRIVATE"); + } + try (InputStream input = client.getObjectContent(object.getFileName())) { + int limit = (int) Math.min(Integer.MAX_VALUE - 1L, maxBytes); + byte[] bytes = input.readNBytes(limit + 1); + if (bytes.length > maxBytes) throw new ServiceException("PERSONAL_PARSE_TOO_LARGE"); + return bytes; + } + }); + } + + private static byte[] runInTenant(String tenantId, OcrSourceOperation operation) throws Exception { + String previous = TenantHelper.getDynamic(); + TenantHelper.setDynamic(tenantId); + try { + return operation.execute(); + } finally { + TenantHelper.clearDynamic(); + if (previous != null && !previous.isBlank()) TenantHelper.setDynamic(previous); + } + } + + private static long maxBytes(PersonalKnowledgeProperties properties) { + return Math.multiplyExact(properties.getMaxFileSizeMb(), 1024L * 1024L); + } + + private static String ownerPrefix(OcrJob job) { + return "personal/" + job.tenantId() + "/" + job.ownerUserId() + "/" + job.itemId() + "/"; + } + + @FunctionalInterface + public interface OcrSourceReader { + byte[] read(OcrJob job) throws Exception; + } + + @FunctionalInterface + private interface OcrSourceOperation { + byte[] execute() throws Exception; + } + + public record OcrItem(long id, String tenantId, long spaceId, long ownerUserId, long ossId, + String tagsJson, int attemptVersion) { + } + + public record OcrJob(long id, String tenantId, long ownerUserId, long itemId, long spaceId, + long ossId, String tagsJson, String status, int attemptCount, int totalPages, int nextPage) { + } +} diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalIngestionWorkerTest.java b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalIngestionWorkerTest.java index cfadff59..f0ad1fea 100644 --- a/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalIngestionWorkerTest.java +++ b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalIngestionWorkerTest.java @@ -37,6 +37,28 @@ import static org.mockito.Mockito.when; @Tag("dev") class PersonalIngestionWorkerTest { + @Test + void emptyPdfHandsOffToOcrInsteadOfMarkingParseEmpty() throws Exception { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + KnowledgeDocumentParser parser = mock(KnowledgeDocumentParser.class); + when(jdbc.queryForList(contains("status = 'QUEUED'"))).thenReturn(List.of(item())); + when(jdbc.update(contains("status = 'PARSING'"), eq("000000"), eq(101L), eq(9L), eq(0))).thenReturn(1); + when(parser.parse(eq("sample.pdf"), eq("text/plain"), any(byte[].class))) + .thenReturn(new ParsedDocument("", "application/pdf", Map.of())); + AtomicReference handedOff = + new AtomicReference<>(); + PersonalIngestionWorker worker = PersonalIngestionWorker.forTestWithPdfOcr( + jdbc, mock(ISysOssService.class), parser, immediateTransactions(), + (ossId, prefix, ownerUserId, maxBytes) -> + new PersonalIngestionWorker.StoredObject("sample.pdf", new byte[] {1, 2, 3}), + (item, bytes) -> handedOff.set(item)); + + assertTrue(worker.processNext()); + + assertEquals(9L, handedOff.get().id()); + verify(jdbc, never()).update(contains("PERSONAL_PARSE_EMPTY"), any(), any(), any(), any(), any()); + } + @Test void workerReadsPrivateObjectInsideItemTenantScope() throws Exception { JdbcTemplate jdbc = mock(JdbcTemplate.class); diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalPdfOcrServiceTest.java b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalPdfOcrServiceTest.java new file mode 100644 index 00000000..6241a053 --- /dev/null +++ b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/personal/PersonalPdfOcrServiceTest.java @@ -0,0 +1,199 @@ +package org.dromara.aihr.personal; + +import org.dromara.aihr.personal.service.PersonalPdfOcrService; +import org.dromara.aihr.personal.service.PersonalPdfOcrService.OcrItem; +import org.dromara.aihr.personal.service.PersonalPdfPageRenderer; +import org.dromara.aihr.personal.service.PersonalVisionOcrService; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.jdbc.core.BatchPreparedStatementSetter; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.transaction.support.TransactionCallback; +import org.springframework.transaction.support.TransactionTemplate; + +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Tag("dev") +class PersonalPdfOcrServiceTest { + + @Test + void enqueueCreatesOwnerScopedJobAndOnePendingRowPerPage() { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class); + when(renderer.pageCount(any(byte[].class))).thenReturn(8); + TransactionTemplate transactions = immediateTransactions(); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, renderer, mock(PersonalVisionOcrService.class), transactions, + job -> new byte[] {1}); + OcrItem item = new OcrItem(9L, "000000", 7L, 101L, 88L, "[]", 1); + + service.enqueue(item, new byte[] {1, 2, 3}); + + verify(jdbc).update(contains("insert into aihr_personal_ocr_job"), + eq("000000"), eq(101L), eq(9L), eq(8)); + ArgumentCaptor pages = + ArgumentCaptor.forClass(BatchPreparedStatementSetter.class); + verify(jdbc).batchUpdate(contains("insert into aihr_personal_ocr_page"), pages.capture()); + assertEquals(8, pages.getValue().getBatchSize()); + } + + @Test + void batchPageNumbersNeverExceedTwentyAndContinueWherePreviousBatchEnded() { + assertEquals(List.of(1, 2, 3), PersonalPdfOcrService.batchPageNumbers(1, 3)); + assertEquals(20, PersonalPdfOcrService.batchPageNumbers(1, 86).size()); + assertEquals(21, PersonalPdfOcrService.batchPageNumbers(21, 86).get(0)); + assertEquals(40, PersonalPdfOcrService.batchPageNumbers(21, 86).get(19)); + assertEquals(List.of(81, 82, 83, 84, 85, 86), PersonalPdfOcrService.batchPageNumbers(81, 86)); + } + + @Test + void processNextBatchRendersAndRecognizesAtMostTwentyPages() throws Exception { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class); + PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class); + Map job = Map.of( + "id", 5L, "tenant_id", "000000", "owner_user_id", 101L, "item_id", 9L, + "space_id", 7L, "oss_id", 88L, "tags_json", "[]", "total_pages", 86, "next_page", 1); + when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job)); + when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1); + List rendered = java.util.stream.IntStream.range(0, 20) + .mapToObj(index -> new PersonalPdfPageRenderer.RenderedPage(index, new byte[] {1}, "image/jpeg")) + .toList(); + when(renderer.render(any(byte[].class), eq(0), eq(20))).thenReturn(rendered); + when(vision.recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class))) + .thenAnswer(invocation -> "第" + invocation.getArgument(2) + "页"); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9, 8, 7}); + + assertEquals(true, service.processNextBatch()); + + verify(renderer).render(any(byte[].class), eq(0), eq(20)); + verify(vision, times(20)).recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class)); + verify(jdbc).update(contains("next_page = ?"), eq(20), eq(20), eq(0), eq(21), eq(5L)); + } + + @Test + void finalBatchPublishesFragmentsAndMarksItemReady() throws Exception { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class); + PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class); + Map job = Map.of( + "id", 5L, "tenant_id", "000000", "owner_user_id", 101L, "item_id", 9L, + "space_id", 7L, "oss_id", 88L, "tags_json", "[]", "total_pages", 2, "next_page", 1); + when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job)); + when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1); + when(renderer.render(any(byte[].class), eq(0), eq(2))).thenReturn(List.of( + new PersonalPdfPageRenderer.RenderedPage(0, new byte[] {1}, "image/jpeg"), + new PersonalPdfPageRenderer.RenderedPage(1, new byte[] {2}, "image/jpeg"))); + when(vision.recognize(any(byte[].class), eq("image/jpeg"), eq(1))).thenReturn("第一页内容"); + when(vision.recognize(any(byte[].class), eq("image/jpeg"), eq(2))).thenReturn("第二页内容"); + when(jdbc.queryForList(contains("select ocr_text"), eq(String.class), + eq("000000"), eq(101L), eq(9L))).thenReturn(List.of("第一页内容", "第二页内容")); + when(jdbc.queryForList(contains("select page_number"), eq(Integer.class), + eq("000000"), eq(101L), eq(9L))).thenReturn(List.of()); + when(jdbc.update(contains("set status = 'READY'"), any(), eq("[]"), + eq("000000"), eq(101L), eq(9L))).thenReturn(1); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9, 8, 7}); + + assertEquals(true, service.processNextBatch()); + + verify(jdbc).batchUpdate(contains("insert into aihr_personal_fragment"), + any(BatchPreparedStatementSetter.class)); + verify(jdbc).update(contains("set status = 'READY'"), any(), eq("[]"), + eq("000000"), eq(101L), eq(9L)); + verify(jdbc).update(contains("status = 'COMPLETED'"), eq(2), eq(2), eq(0), any(), eq(5L)); + } + + @Test + void retryJobProcessesOnlyFailedPagesThenPublishes() throws Exception { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class); + PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class); + Map job = Map.ofEntries( + Map.entry("id", 5L), Map.entry("tenant_id", "000000"), Map.entry("owner_user_id", 101L), + Map.entry("item_id", 9L), Map.entry("space_id", 7L), Map.entry("oss_id", 88L), + Map.entry("tags_json", "[]"), Map.entry("total_pages", 30), Map.entry("next_page", 31), + Map.entry("status", "RETRY")); + when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job)); + when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1); + when(jdbc.queryForList(contains("attempt_count < 3"), eq(Integer.class), + eq("000000"), eq(101L), eq(9L), eq(20))).thenReturn(List.of(4, 27)); + when(renderer.render(any(byte[].class), eq(3), eq(1))).thenReturn(List.of( + new PersonalPdfPageRenderer.RenderedPage(3, new byte[] {4}, "image/jpeg"))); + when(renderer.render(any(byte[].class), eq(26), eq(1))).thenReturn(List.of( + new PersonalPdfPageRenderer.RenderedPage(26, new byte[] {27}, "image/jpeg"))); + when(vision.recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class))) + .thenReturn("重试成功"); + when(jdbc.queryForObject(contains("attempt_count < 3"), eq(Integer.class), + eq("000000"), eq(101L), eq(9L))).thenReturn(0); + when(jdbc.queryForList(contains("select ocr_text"), eq(String.class), + eq("000000"), eq(101L), eq(9L))).thenReturn(List.of("已有页面", "重试成功", "重试成功")); + when(jdbc.queryForList(contains("select page_number"), eq(Integer.class), + eq("000000"), eq(101L), eq(9L))).thenReturn(List.of()); + when(jdbc.update(contains("set status = 'READY'"), any(), eq("[]"), + eq("000000"), eq(101L), eq(9L))).thenReturn(1); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9}); + + assertEquals(true, service.processNextBatch()); + + verify(renderer).render(any(byte[].class), eq(3), eq(1)); + verify(renderer).render(any(byte[].class), eq(26), eq(1)); + verify(vision, times(2)).recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class)); + verify(jdbc).update(contains("set status = 'READY'"), any(), eq("[]"), + eq("000000"), eq(101L), eq(9L)); + } + + @Test + void recoversExpiredRunningLease() { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, mock(PersonalPdfPageRenderer.class), mock(PersonalVisionOcrService.class), + immediateTransactions(), ignored -> new byte[] {1}); + + service.recoverExpiredJobs(); + + verify(jdbc).update(contains("lease_until < now()")); + } + + @Test + void thirdBatchInfrastructureFailureStopsRetryLoop() throws Exception { + JdbcTemplate jdbc = mock(JdbcTemplate.class); + Map job = Map.ofEntries( + Map.entry("id", 5L), Map.entry("tenant_id", "000000"), Map.entry("owner_user_id", 101L), + Map.entry("item_id", 9L), Map.entry("space_id", 7L), Map.entry("oss_id", 88L), + Map.entry("tags_json", "[]"), Map.entry("total_pages", 2), Map.entry("next_page", 1), + Map.entry("status", "PENDING"), Map.entry("attempt_count", 2)); + when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job)); + when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1); + PersonalPdfOcrService service = PersonalPdfOcrService.forTest( + jdbc, mock(PersonalPdfPageRenderer.class), mock(PersonalVisionOcrService.class), + immediateTransactions(), ignored -> { throw new IllegalStateException("storage unavailable"); }); + + assertEquals(true, service.processNextBatch()); + + verify(jdbc).update(contains("status = 'FAILED'"), eq(5L)); + verify(jdbc).update(contains("PERSONAL_OCR_BATCH_FAILED"), eq("000000"), eq(101L), eq(9L)); + } + + private TransactionTemplate immediateTransactions() { + TransactionTemplate transactions = mock(TransactionTemplate.class); + when(transactions.execute(any())).thenAnswer(invocation -> { + TransactionCallback callback = invocation.getArgument(0); + return callback.doInTransaction(null); + }); + return transactions; + } +}