Properly implement KnowledgeManagerImpl::retryFailedIngestion

Signed-off-by: Florian Hotze <[email protected]>

src/main/java/com/github/llamara/ai/internal/knowledge/KnowledgeManager.java

@@ -204,27 +204,29 @@ void removePermission(UUID id, User user)
     void setLabel(UUID id, String label) throws KnowledgeNotFoundException;
 
     /**
-     * Retry failed ingestion of a knowledge specified by its id. If the ingestion status isn't
-     * {@link com.github.llamara.ai.internal.ingestion.IngestionStatus#FAILED}, does nothing.
+     * Get the source file of the file-based knowledge specified by its id.
      *
      * @param id persistent unique id of knowledge
+     * @return the knowledge source file as {@link InputStream}
      * @throws KnowledgeNotFoundException if no knowledge with the given id was found
      * @throws UnexpectedFileStorageFailureException if a {@link FileStorage} operation failed
      *     unexpectedly
      */
-    void retryFailedIngestion(UUID id)
+    NamedFileContainer getFile(UUID id)
             throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException;
 
     /**
-     * Get the source file of the file-based knowledge specified by its id.
+     * Retry the failed ingestion of a knowledge specified by its id.
+     *
+     * <p>Implementations should do nothing if the current ingestion status is not {@link
+     * IngestionStatus#FAILED}.
      *
      * @param id persistent unique id of knowledge
-     * @return the knowledge source file as {@link InputStream}
      * @throws KnowledgeNotFoundException if no knowledge with the given id was found
      * @throws UnexpectedFileStorageFailureException if a {@link FileStorage} operation failed
      *     unexpectedly
      */
-    NamedFileContainer getFile(UUID id)
+    void retryFailedIngestion(UUID id)
             throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException;
 
     /**

src/main/java/com/github/llamara/ai/internal/knowledge/KnowledgeManagerImpl.java

@@ -32,9 +32,9 @@
 import com.github.llamara.ai.internal.security.PermissionMetadataMapper;
 import com.github.llamara.ai.internal.security.user.User;
 
+import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
-import java.io.InputStream;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -303,40 +303,6 @@ public void setLabel(UUID id, String label) throws KnowledgeNotFoundException {
         repository.persist(knowledge);
     }
 
-    @Override
-    public void retryFailedIngestion(UUID id)
-            throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
-        Knowledge knowledge = getKnowledge(id);
-        if (knowledge.getIngestionStatus() != IngestionStatus.FAILED) {
-            return;
-        }
-
-        QuarkusTransaction.begin();
-        knowledge = getKnowledge(id);
-        repository.persist(knowledge);
-        NamedFileContainer file = getFile(id);
-        Map<String, String> metadata = createEmbeddingMetadata(knowledge);
-        Path tempFile = null;
-        try (InputStream fileContent = file.content()) {
-            tempFile = Files.createTempDirectory("").resolve(UUID.randomUUID().toString() + ".tmp");
-            Files.copy(fileContent, tempFile, StandardCopyOption.REPLACE_EXISTING);
-            knowledge.setIngestionStatus(IngestionStatus.PENDING);
-            ingestToStore(tempFile, metadata);
-        } catch (IOException e) {
-            System.err.println("Error creating or copying file: " + e.getMessage());
-        } finally {
-            if (tempFile != null) {
-                try {
-                    Files.deleteIfExists(tempFile);
-                } catch (IOException e) {
-                    System.err.println("Failed to delete temporary file: " + e.getMessage());
-                }
-            }
-        }
-
-        QuarkusTransaction.commit();
-    }
-
     @Override
     public NamedFileContainer getFile(UUID id)
             throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
@@ -351,6 +317,38 @@ public NamedFileContainer getFile(UUID id)
         }
     }
 
+    @Override
+    public void retryFailedIngestion(UUID id)
+            throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
+        Knowledge knowledge = getKnowledge(id);
+
+        // Do nothing if ingestion status is not FAILED
+        if (knowledge.getIngestionStatus() != IngestionStatus.FAILED) {
+            return;
+        }
+
+        // Get the file from the file storage and save it as temporary file
+        NamedFileContainer fc = getFile(id);
+        File tempFile;
+        try {
+            tempFile = File.createTempFile("knowledge_" + id, null);
+            tempFile.deleteOnExit(); // the file will be deleted when the JVM exits
+            Files.copy(fc.content(), tempFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
+        } catch (IOException e) {
+            throw new UnexpectedFileStorageFailureException(
+                    "Failed to temporary save retrieved file to disk", e);
+        }
+
+        // Create metadata
+        Map<String, String> metadata = createEmbeddingMetadata(knowledge);
+
+        // Reset ingestion status
+        setKnowledgeIngestionStatus(id, IngestionStatus.PENDING);
+
+        // Retry ingestion
+        ingestToStore(tempFile.toPath(), metadata);
+    }
+
     private Map<String, String> createFileMetadata(String checksum, String contentType) {
         return Map.of(MetadataKeys.CHECKSUM, checksum, MetadataKeys.CONTENT_TYPE, contentType);
     }

src/main/java/com/github/llamara/ai/internal/knowledge/storage/UnexpectedFileStorageFailureException.java

@@ -25,7 +25,7 @@
  * @author Florian Hotze - Initial contribution
  */
 public class UnexpectedFileStorageFailureException extends Exception {
-    UnexpectedFileStorageFailureException(String message, Throwable cause) {
+    public UnexpectedFileStorageFailureException(String message, Throwable cause) {
         super(message, cause);
     }
 }

src/main/java/com/github/llamara/ai/internal/rest/KnowledgeResource.java

@@ -425,15 +425,14 @@ public void removeKnowledgeTag(
     @ResponseStatus(200)
     @Operation(
             operationId = "retryFailedIngestion",
-            summary = "Retry failed ingestion of a single knowledge identified by it ID.",
-            description = "If the Ingestion status isn't failed, nothing happens.")
+            summary = "Retry the failed ingestion of a single knowledge identified by it ID.")
     @APIResponse(responseCode = "200", description = "OK.")
     @APIResponse(responseCode = "404", description = "No knowledge with the given id found.")
     public void retryFailedIngestion(
             @PathParam("id")
                     @Parameter(
                             name = "id",
-                            description = "UID of the knowledge to retry",
+                            description = "UID of the knowledge to retry the ingestion for",
                             required = true)
                     UUID id)
             throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {

src/main/java/com/github/llamara/ai/internal/security/knowledge/UserKnowledgeManagerImpl.java

@@ -229,19 +229,18 @@ public void setLabel(UUID id, String label) throws KnowledgeNotFoundException {
         delegate.setLabel(id, label);
     }
 
-    @Override
-    public void retryFailedIngestion(UUID id)
-            throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
-        enforceKnowledgeEditable(id);
-
-        delegate.retryFailedIngestion(id);
-    }
-
     @Override
     public NamedFileContainer getFile(UUID id)
             throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
         userManager.enforceRegistered();
         Knowledge knowledge = getKnowledge(id);
         return delegate.getFile(knowledge.getId());
     }
+
+    @Override
+    public void retryFailedIngestion(UUID id)
+            throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
+        enforceKnowledgeEditable(id);
+        delegate.retryFailedIngestion(id);
+    }
 }

src/test/java/com/github/llamara/ai/internal/knowledge/KnowledgeManagerImplTest.java

@@ -616,24 +616,49 @@ void removeTagDoesNothingIfTagNotExists() throws KnowledgeNotFoundException {
         @Test
         void retryFailedIngestionDoesNothingIfKnowledgeIngestionStatusIsSucceeded()
                 throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
+            // setup
             knowledgeManager.setKnowledgeIngestionStatus(knowledgeId, IngestionStatus.SUCCEEDED);
+
+            // test
             knowledgeManager.retryFailedIngestion(knowledgeId);
+
             verify(documentIngestor, times(0)).ingestDocument(any());
         }
 
         @Test
         void retryFailedIngestionDoesNothingIfKnowledgeIngestionStatusIsPending()
                 throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
+            // setup
             knowledgeManager.setKnowledgeIngestionStatus(knowledgeId, IngestionStatus.PENDING);
+
+            // test
             knowledgeManager.retryFailedIngestion(knowledgeId);
+
             verify(documentIngestor, times(0)).ingestDocument(any());
         }
 
         @Test
-        void retryFailedIngestionDoesDispatchIngestionIfKnowledgeIngestionStatusIsFailed()
+        void retryFailedIngestionSetsIngestionStatusToPendingIfIngestionStatusIsFailed()
+                throws UnexpectedFileStorageFailureException, KnowledgeNotFoundException {
+            // setup
+            knowledgeManager.setKnowledgeIngestionStatus(knowledgeId, IngestionStatus.FAILED);
+
+            // test
+            knowledgeManager.retryFailedIngestion(knowledgeId);
+
+            verify(knowledgeRepository, times(1))
+                    .setStatusFor(knowledgeId, IngestionStatus.PENDING);
+        }
+
+        @Test
+        void retryFailedIngestionDispatchesIngestionIfKnowledgeIngestionStatusIsFailed()
                 throws KnowledgeNotFoundException, UnexpectedFileStorageFailureException {
+            // setup
             knowledgeManager.setKnowledgeIngestionStatus(knowledgeId, IngestionStatus.FAILED);
+
+            // test
             knowledgeManager.retryFailedIngestion(knowledgeId);
+
             verify(documentIngestor, times(1)).ingestDocument(any());
         }
     }