feat(personal): expose OCR progress and retry

This commit is contained in:
2026-07-12 21:42:56 +08:00
parent 0e3a68c3c1
commit 2e8961c75c
6 changed files with 117 additions and 4 deletions
@@ -6,6 +6,7 @@ import org.dromara.aihr.personal.domain.PersonalAssistantDto.AskResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.DownloadUrlResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.ItemCreatedResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.ItemResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.OcrProgressResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.PageResponse;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.PersonalSearchRequest;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.PersonalSearchResponse;
@@ -17,6 +18,7 @@ import org.dromara.aihr.personal.domain.PersonalAssistantDto.UrlItemRequest;
import org.dromara.aihr.personal.service.PersonalAnswerService;
import org.dromara.aihr.personal.service.PersonalCleanupService;
import org.dromara.aihr.personal.service.PersonalIngestionService;
import org.dromara.aihr.personal.service.PersonalPdfOcrService;
import org.dromara.aihr.personal.service.PersonalRetrievalService;
import org.dromara.aihr.personal.service.PersonalSpaceService;
import org.dromara.aihr.personal.service.PersonalUrlFetchService;
@@ -53,6 +55,7 @@ public class PersonalAssistantController {
private final PersonalRetrievalService retrievalService;
private final PersonalAnswerService answerService;
private final PersonalCleanupService cleanupService;
private final PersonalPdfOcrService pdfOcrService;
@GetMapping("/space")
public R<SpaceResponse> space() {
@@ -95,14 +98,20 @@ public class PersonalAssistantController {
@GetMapping("/items/{id}")
public R<ItemResponse> item(@PathVariable long id) {
return R.ok(spaceService.itemResponse(owner(), id));
PersonalOwner owner = owner();
return R.ok(withOcr(spaceService.itemResponse(owner, id), pdfOcrService.progress(owner, id)));
}
@PostMapping("/items/{id}/retry")
public R<ItemResponse> retry(@PathVariable long id) {
PersonalOwner owner = owner();
ingestionService.retry(owner, id);
return R.ok(spaceService.itemResponse(owner, id));
return R.ok(withOcr(spaceService.itemResponse(owner, id), pdfOcrService.progress(owner, id)));
}
@PostMapping("/items/{id}/ocr/retry-failed")
public R<OcrProgressResponse> retryFailedOcrPages(@PathVariable long id) {
return R.ok(pdfOcrService.retryFailedPages(owner(), id));
}
@DeleteMapping("/items/{id}")
@@ -145,4 +154,10 @@ public class PersonalAssistantController {
private PersonalOwner owner() {
return ownerProvider.current();
}
private static ItemResponse withOcr(ItemResponse item, OcrProgressResponse ocr) {
return new ItemResponse(item.id(), item.sourceType(), item.title(), item.originalUrl(), item.mimeType(),
item.sizeBytes(), item.status(), item.errorCode(), item.errorMessage(), item.summary(), item.tags(),
item.capturedAt(), item.parsedAt(), ocr);
}
}
@@ -132,6 +132,14 @@ public class PersonalCleanupService {
delete from aihr_personal_fragment
where tenant_id = ? and owner_user_id = ? and item_id = ?
""", item.owner().tenantId(), item.owner().userId(), item.itemId());
jdbc.update("""
delete from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ?
""", item.owner().tenantId(), item.owner().userId(), item.itemId());
jdbc.update("""
delete from aihr_personal_ocr_job
where tenant_id = ? and owner_user_id = ? and item_id = ?
""", item.owner().tenantId(), item.owner().userId(), item.itemId());
if (item.ossId() != null && item.ossId() > 0) {
ossCleanup.delete(item.ossId());
}
@@ -1,7 +1,9 @@
package org.dromara.aihr.personal.service;
import org.dromara.aihr.knowledge.parse.ParsedDocument;
import org.dromara.aihr.personal.domain.PersonalAssistantDto.OcrProgressResponse;
import org.dromara.aihr.personal.support.PersonalKnowledgeProperties;
import org.dromara.aihr.personal.support.PersonalOwner;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.oss.core.OssClient;
import org.dromara.common.oss.enums.AccessPolicyType;
@@ -110,6 +112,57 @@ public class PersonalPdfOcrService {
return List.copyOf(pages);
}
public OcrProgressResponse progress(PersonalOwner owner, long itemId) {
List<Map<String, Object>> rows = jdbcTemplate.queryForList("""
select status, total_pages, processed_pages, success_pages, failed_pages
from aihr_personal_ocr_job
where tenant_id = ? and owner_user_id = ? and item_id = ?
limit 1
""", owner.tenantId(), owner.userId(), itemId);
if (rows.isEmpty()) return null;
Map<String, Object> row = rows.get(0);
List<Integer> failed = 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, owner.tenantId(), owner.userId(), itemId);
return new OcrProgressResponse(true, String.valueOf(row.get("status")),
intValue(row, "total_pages", 0), intValue(row, "processed_pages", 0),
intValue(row, "success_pages", 0), intValue(row, "failed_pages", 0), List.copyOf(failed));
}
public OcrProgressResponse retryFailedPages(PersonalOwner owner, long itemId) {
OcrProgressResponse current = progress(owner, itemId);
if (current == null || current.failedPages() < 1) {
throw new ServiceException("PERSONAL_OCR_RETRY_NOT_AVAILABLE");
}
transactions.execute(status -> {
int pages = jdbcTemplate.update("""
update aihr_personal_ocr_page
set status = 'FAILED', attempt_count = 0, last_error_code = null, update_time = now()
where tenant_id = ? and owner_user_id = ? and item_id = ? and status = 'FAILED'
""", owner.tenantId(), owner.userId(), itemId);
if (pages < 1) throw new ServiceException("PERSONAL_OCR_RETRY_NOT_AVAILABLE");
int job = jdbcTemplate.update("""
update aihr_personal_ocr_job
set status = 'RETRY', attempt_count = 0, lease_until = null,
last_error_code = null, update_time = now()
where tenant_id = ? and owner_user_id = ? and item_id = ?
and status in ('COMPLETED','FAILED')
""", owner.tenantId(), owner.userId(), itemId);
int item = jdbcTemplate.update("""
update aihr_personal_item
set status = 'PARSING', error_code = null, error_message = null, update_time = now()
where tenant_id = ? and owner_user_id = ? and id = ?
and status in ('READY','FAILED')
""", owner.tenantId(), owner.userId(), itemId);
if (job != 1 || item != 1) throw new ServiceException("PERSONAL_OCR_RETRY_NOT_AVAILABLE");
return null;
});
return new OcrProgressResponse(true, "RETRY", current.totalPages(), current.processedPages(),
current.successPages(), current.failedPages(), current.failedPageNumbers());
}
@Scheduled(fixedDelayString = "${aihr.personal.ocr-delay-ms:2000}", scheduler = "personalTaskScheduler")
public void poll() {
recoverExpiredJobs();
@@ -20,6 +20,7 @@ import java.time.LocalDateTime;
import java.util.List;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -34,17 +35,23 @@ class PersonalAssistantControllerTest {
PersonalOwner owner = new PersonalOwner("000000", 101L, null);
when(owners.current()).thenReturn(owner);
PersonalSpaceService spaces = mock(PersonalSpaceService.class);
when(spaces.itemResponse(eq(owner), eq(9L))).thenReturn(
new org.dromara.aihr.personal.domain.PersonalAssistantDto.ItemResponse(
9L, "FILE", "资料", null, "application/pdf", 1L, "PARSING",
null, null, null, List.of(), null, null));
PersonalIngestionService ingestion = mock(PersonalIngestionService.class);
PersonalUrlFetchService urls = mock(PersonalUrlFetchService.class);
PersonalRetrievalService retrieval = mock(PersonalRetrievalService.class);
PersonalAnswerService answers = mock(PersonalAnswerService.class);
PersonalCleanupService cleanup = mock(PersonalCleanupService.class);
org.dromara.aihr.personal.service.PersonalPdfOcrService pdfOcr =
mock(org.dromara.aihr.personal.service.PersonalPdfOcrService.class);
when(urls.fetch("https://example.com/a")).thenReturn(new PersonalUrlFetchService.FetchResult(
java.net.URI.create("https://example.com/a"), 200, "text/plain", "a".getBytes(),
java.time.Instant.now(), "hash"));
PersonalAssistantController controller = new PersonalAssistantController(owners, spaces, ingestion, urls,
retrieval, answers, cleanup);
retrieval, answers, cleanup, pdfOcr);
controller.space();
controller.items(1, 20, null, null, null, null, null);
controller.createText(new TextItemRequest("note", "body", null, List.of()));
@@ -52,6 +59,7 @@ class PersonalAssistantControllerTest {
controller.createUrl(new UrlItemRequest("https://example.com/a", null, null));
controller.item(9L);
controller.retry(9L);
controller.retryFailedOcrPages(9L);
controller.deleteItem(9L);
controller.downloadUrl(9L);
controller.search(null);
@@ -60,7 +68,7 @@ class PersonalAssistantControllerTest {
controller.session(3L);
controller.deleteSession(3L);
verify(owners, times(14)).current();
verify(owners, times(15)).current();
assertFalse(PersonalAssistantController.class.isAnnotationPresent(SaIgnore.class));
for (var method : PersonalAssistantController.class.getDeclaredMethods()) {
assertFalse(method.isAnnotationPresent(SaIgnore.class), method.getName());
@@ -134,6 +134,8 @@ class PersonalCleanupServiceTest {
InOrder order = inOrder(vectors, jdbc, oss);
order.verify(vectors).deleteItem(OWNER, 9L);
order.verify(jdbc).update(contains("delete from aihr_personal_fragment"), eq("000000"), eq(101L), eq(9L));
order.verify(jdbc).update(contains("delete from aihr_personal_ocr_page"), eq("000000"), eq(101L), eq(9L));
order.verify(jdbc).update(contains("delete from aihr_personal_ocr_job"), eq("000000"), eq(101L), eq(9L));
order.verify(oss).delete(55L);
order.verify(jdbc).update(contains("set status = 'DELETED'"), eq("000000"), eq(101L), eq(9L));
}
@@ -4,6 +4,7 @@ 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.dromara.aihr.personal.support.PersonalOwner;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
@@ -188,6 +189,32 @@ class PersonalPdfOcrServiceTest {
verify(jdbc).update(contains("PERSONAL_OCR_BATCH_FAILED"), eq("000000"), eq(101L), eq(9L));
}
@Test
void progressAndManualRetryRemainOwnerScoped() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
when(jdbc.queryForList(contains("from aihr_personal_ocr_job"),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of(Map.of(
"status", "COMPLETED", "total_pages", 8, "processed_pages", 8,
"success_pages", 7, "failed_pages", 1)));
when(jdbc.queryForList(contains("from aihr_personal_ocr_page"), eq(Integer.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of(6));
when(jdbc.update(contains("set status = 'FAILED', attempt_count = 0"),
eq("000000"), eq(101L), eq(9L))).thenReturn(1);
when(jdbc.update(contains("set status = 'RETRY'"), eq("000000"), eq(101L), eq(9L))).thenReturn(1);
when(jdbc.update(contains("set status = 'PARSING'"), eq("000000"), eq(101L), eq(9L))).thenReturn(1);
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, mock(PersonalPdfPageRenderer.class), mock(PersonalVisionOcrService.class),
immediateTransactions(), ignored -> new byte[] {1});
PersonalOwner owner = new PersonalOwner("000000", 101L, "13800000000");
assertEquals(List.of(6), service.progress(owner, 9L).failedPageNumbers());
assertEquals("RETRY", service.retryFailedPages(owner, 9L).status());
verify(jdbc).update(contains("set status = 'FAILED', attempt_count = 0"),
eq("000000"), eq(101L), eq(9L));
verify(jdbc).update(contains("set status = 'PARSING'"), eq("000000"), eq(101L), eq(9L));
}
private TransactionTemplate immediateTransactions() {
TransactionTemplate transactions = mock(TransactionTemplate.class);
when(transactions.execute(any())).thenAnswer(invocation -> {