feat(personal): process scanned PDFs in OCR batches

This commit is contained in:
2026-07-12 21:33:18 +08:00
parent d63532cc90
commit 0e3a68c3c1
4 changed files with 729 additions and 5 deletions
@@ -35,6 +35,7 @@ public class PersonalIngestionWorker {
private final TransactionTemplate transactionTemplate; private final TransactionTemplate transactionTemplate;
private final StoredObjectReader objectReader; private final StoredObjectReader objectReader;
private final TenantRunner tenantRunner; private final TenantRunner tenantRunner;
private final PdfOcrHandler pdfOcrHandler;
private final long maxInputBytes; private final long maxInputBytes;
private final int chunkSize; private final int chunkSize;
private final int chunkOverlap; private final int chunkOverlap;
@@ -44,22 +45,25 @@ public class PersonalIngestionWorker {
@Autowired @Autowired
public PersonalIngestionWorker(JdbcTemplate jdbcTemplate, ISysOssService ossService, public PersonalIngestionWorker(JdbcTemplate jdbcTemplate, ISysOssService ossService,
KnowledgeDocumentParser parser, PersonalKnowledgeProperties properties, KnowledgeDocumentParser parser, PersonalKnowledgeProperties properties,
PlatformTransactionManager transactionManager) { PlatformTransactionManager transactionManager,
PersonalPdfOcrService pdfOcrService) {
this(jdbcTemplate, parser, new TransactionTemplate(transactionManager), this(jdbcTemplate, parser, new TransactionTemplate(transactionManager),
defaultReader(ossService, PersonalIngestionWorker::ossClient), configuredMaxBytes(properties), defaultReader(ossService, PersonalIngestionWorker::ossClient), configuredMaxBytes(properties),
properties.getChunkSize(), properties.getChunkOverlap(), properties.getParsingLeaseMinutes(), properties.getChunkSize(), properties.getChunkOverlap(), properties.getParsingLeaseMinutes(),
properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant); properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant, pdfOcrService::enqueue);
} }
private PersonalIngestionWorker(JdbcTemplate jdbcTemplate, KnowledgeDocumentParser parser, private PersonalIngestionWorker(JdbcTemplate jdbcTemplate, KnowledgeDocumentParser parser,
TransactionTemplate transactionTemplate, StoredObjectReader objectReader, TransactionTemplate transactionTemplate, StoredObjectReader objectReader,
long maxInputBytes, int chunkSize, int chunkOverlap, long maxInputBytes, int chunkSize, int chunkOverlap,
int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner) { int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner,
PdfOcrHandler pdfOcrHandler) {
this.jdbcTemplate = jdbcTemplate; this.jdbcTemplate = jdbcTemplate;
this.parser = parser; this.parser = parser;
this.transactionTemplate = transactionTemplate; this.transactionTemplate = transactionTemplate;
this.objectReader = objectReader; this.objectReader = objectReader;
this.tenantRunner = tenantRunner; this.tenantRunner = tenantRunner;
this.pdfOcrHandler = pdfOcrHandler;
this.maxInputBytes = maxInputBytes; this.maxInputBytes = maxInputBytes;
if (chunkSize <= 0 || chunkOverlap < 0 || chunkOverlap >= chunkSize if (chunkSize <= 0 || chunkOverlap < 0 || chunkOverlap >= chunkSize
|| parsingLeaseMinutes <= 0 || maxParseAttempts <= 0) { || parsingLeaseMinutes <= 0 || maxParseAttempts <= 0) {
@@ -76,7 +80,8 @@ public class PersonalIngestionWorker {
TransactionTemplate transactionTemplate, TransactionTemplate transactionTemplate,
StoredObjectReader objectReader) { StoredObjectReader objectReader) {
return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, 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, public static PersonalIngestionWorker forTest(JdbcTemplate jdbcTemplate, ISysOssService ossService,
@@ -84,7 +89,17 @@ public class PersonalIngestionWorker {
TransactionTemplate transactionTemplate, TransactionTemplate transactionTemplate,
StoredObjectReader objectReader, TenantRunner tenantRunner) { StoredObjectReader objectReader, TenantRunner tenantRunner) {
return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader, 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, public static StoredObjectReader objectReaderForTest(ISysOssService ossService,
@@ -173,6 +188,12 @@ public class PersonalIngestionWorker {
ParsedDocument document = parser.parse(stored.fileName(), item.mimeType(), stored.bytes()); ParsedDocument document = parser.parse(stored.fileName(), item.mimeType(), stored.bytes());
List<String> chunks = document.chunks(chunkSize, chunkOverlap); List<String> chunks = document.chunks(chunkSize, chunkOverlap);
if (chunks.isEmpty()) { 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( throw new KnowledgeDocumentParser.ParseException(
KnowledgeDocumentParser.Failure.EMPTY, "document contains no text"); KnowledgeDocumentParser.Failure.EMPTY, "document contains no text");
} }
@@ -308,6 +329,12 @@ public class PersonalIngestionWorker {
&& "PERSONAL_OSS_NOT_PRIVATE".equals(serviceError.getMessage())) { && "PERSONAL_OSS_NOT_PRIVATE".equals(serviceError.getMessage())) {
return new Failure("PERSONAL_OSS_NOT_PRIVATE", "个人资料存储策略不可用"); 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(); candidate = candidate.getCause();
} }
return new Failure("PERSONAL_PARSE_FAILED", "资料处理失败,请稍后重试"); return new Failure("PERSONAL_PARSE_FAILED", "资料处理失败,请稍后重试");
@@ -351,6 +378,18 @@ public class PersonalIngestionWorker {
return "personal/" + item.tenantId() + "/" + item.ownerUserId() + "/" + item.id() + "/"; 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 @FunctionalInterface
public interface StoredObjectReader { public interface StoredObjectReader {
StoredObject read(long ossId, String expectedPrefix, long ownerUserId, long maxBytes) throws Exception; StoredObject read(long ossId, String expectedPrefix, long ownerUserId, long maxBytes) throws Exception;
@@ -371,6 +410,11 @@ public class PersonalIngestionWorker {
OssClient get(String configKey); OssClient get(String configKey);
} }
@FunctionalInterface
public interface PdfOcrHandler {
void enqueue(PersonalPdfOcrService.OcrItem item, byte[] pdfBytes);
}
public record StoredObject(String fileName, byte[] bytes) { public record StoredObject(String fileName, byte[] bytes) {
public StoredObject { public StoredObject {
bytes = bytes == null ? new byte[0] : bytes.clone(); bytes = bytes == null ? new byte[0] : bytes.clone();
@@ -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<Integer> batchPageNumbers(int firstPage, int totalPages) {
if (firstPage < 1 || totalPages < firstPage) {
return List.of();
}
int end = Math.min(totalPages, firstPage + PersonalPdfPageRenderer.BATCH_SIZE - 1);
List<Integer> 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<Map<String, Object>> 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<Integer> pageNumbers = retryJob(job) ? retryPageNumbers(job)
: batchPageNumbers(job.nextPage(), job.totalPages());
int success = 0;
int failed = 0;
try {
byte[] pdfBytes = sourceReader.read(job);
List<PersonalPdfPageRenderer.RenderedPage> 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<Integer> 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<PersonalPdfPageRenderer.RenderedPage> renderPages(byte[] pdfBytes, List<Integer> pageNumbers,
boolean retry) {
if (!retry) {
if (pageNumbers.isEmpty()) return List.of();
return renderer.render(pdfBytes, pageNumbers.get(0) - 1, pageNumbers.size());
}
List<PersonalPdfPageRenderer.RenderedPage> rendered = new ArrayList<>(pageNumbers.size());
for (int pageNumber : pageNumbers) {
List<PersonalPdfPageRenderer.RenderedPage> 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<String> 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<Integer> 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<Integer> failedPages) {
List<String> 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<Integer> 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<String, Object> 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<String, Object> 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<String, Object> 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) {
}
}
@@ -37,6 +37,28 @@ import static org.mockito.Mockito.when;
@Tag("dev") @Tag("dev")
class PersonalIngestionWorkerTest { 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<org.dromara.aihr.personal.service.PersonalPdfOcrService.OcrItem> 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 @Test
void workerReadsPrivateObjectInsideItemTenantScope() throws Exception { void workerReadsPrivateObjectInsideItemTenantScope() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class); JdbcTemplate jdbc = mock(JdbcTemplate.class);
@@ -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<BatchPreparedStatementSetter> 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<String, Object> 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<PersonalPdfPageRenderer.RenderedPage> 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.<Integer>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<String, Object> 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<String, Object> 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<String, Object> 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;
}
}