Add wait flag for extractions. (#5101)

DARPA-ASKEM · Oct 9, 2024 · 794e558 · 794e558
1 parent f1fc1aa
commit 794e558
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/.../java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java b/.../java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java
@@ -21,6 +21,7 @@
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeoutException;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;
@@ -797,14 +798,23 @@ public ResponseEntity<Void> variableExtractions(
 	)
 	public ResponseEntity<Void> pdfExtractions(
 		@RequestParam("document-id") final UUID documentId,
-		@RequestParam(name = "domain", defaultValue = "epi") final String domain,
-		@RequestParam(name = "project-id", required = false) final UUID projectId
+		@RequestParam(name = "project-id", required = false) final UUID projectId,
+		@RequestParam(name = "wait", required = false, defaultValue = "false") final Boolean wait
 	) {
 		final Schema.Permission permission = projectService.checkPermissionCanWrite(
 			currentUserService.get().getId(),
 			projectId
 		);
-		extractionService.extractPDFAndApplyToDocument(documentId, projectId, permission);
+
+		final Future<DocumentAsset> f = extractionService.extractPDFAndApplyToDocument(documentId, projectId, permission);
+		if (wait) {
+			try {
+				f.get();
+			} catch (InterruptedException | ExecutionException e) {
+				log.error("Error extracting PDF", e);
+				throw new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR, messages.get("document.extracton.failed"));
+			}
+		}
 		return ResponseEntity.accepted().build();
 	}
 

diff --git a/packages/server/src/main/resources/application.properties b/packages/server/src/main/resources/application.properties
@@ -87,6 +87,7 @@ terarium.serviceRequestPatterns[11]=/simulations/**
 terarium.serviceRequestPatterns[12]=/code-asset/**
 terarium.serviceRequestPatterns[13]=/workflows/**
 terarium.serviceRequestPatterns[14]=/interventions/**
+terarium.serviceRequestPatterns[15]=/knowledge/**
 terarium.extractionService.poolSize=10
 terarium.cacheHeadersMaxAge=86400
 management.endpoints.web.exposure.include=health,prometheus

diff --git a/packages/server/src/main/resources/messages.properties b/packages/server/src/main/resources/messages.properties
@@ -32,6 +32,7 @@ document.not-found = We couldn't find the requested document. Please verify that
 document.unable-to-delete = We couldn't delete the document. Please try again later.
 document.unable-to-update = We couldn't update the document. Please try again later.
 document.text-length-exceeded = This document exceeds our 600,000-word limit. If you still need to analyze this document, contact support.
+document.extracton.failed = Document extraction failed. Please try again later.
 
 mira.concept.bad-curies = Our domain knowledge service (MIRA) couldn't find any valid concepts for the selected IDs. Please review the IDs that you selected and try again.
 mira.concept.bad-query = Our domain knowledge service (MIRA) couldn't find any valid concepts for the entered query. Please review your search criteria and try again.