Skip to content

Commit 3953ace

Browse files
ianardeeianardee
committed
✨ Allow CRUD operations on RAG documents
1 parent 2c11160 commit 3953ace

35 files changed

Lines changed: 1799 additions & 120 deletions

‎pom.xml‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,40 @@
5959
<artifactId>maven-install-plugin</artifactId>
6060
<version>${org.apache.maven.install.version}</version>
6161
</plugin>
62-
<!-- site lifecycle, see https://maven.apache.org/ref/current/maven-core/lifecycles.html#site_Lifecycle -->
62+
<plugin>
63+
<groupId>org.projectlombok</groupId>
64+
<artifactId>lombok-maven-plugin</artifactId>
65+
<version>1.18.20.0</version>
66+
<executions>
67+
<execution>
68+
<phase>generate-sources</phase>
69+
<goals>
70+
<goal>delombok</goal>
71+
</goals>
72+
</execution>
73+
</executions>
74+
<configuration>
75+
<sourceDirectory>${project.basedir}/src/main/java</sourceDirectory>
76+
<outputDirectory>${project.build.directory}/delombok</outputDirectory>
77+
<addOutputDirectory>false</addOutputDirectory>
78+
</configuration>
79+
</plugin>
80+
<plugin>
81+
<groupId>org.apache.maven.plugins</groupId>
82+
<artifactId>maven-javadoc-plugin</artifactId>
83+
<version>${org.apache.maven.javadoc.version}</version>
84+
<configuration>
85+
<sourcepath>${project.build.directory}/delombok</sourcepath>
86+
</configuration>
87+
<executions>
88+
<execution>
89+
<id>attach-javadocs</id>
90+
<goals>
91+
<goal>jar</goal>
92+
</goals>
93+
</execution>
94+
</executions>
95+
</plugin>
6396
<plugin>
6497
<artifactId>maven-site-plugin</artifactId>
6598
<version>${org.apache.maven.site.version}</version>

‎src/main/java/com/mindee/exceptions/MindeeInputException.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package com.mindee.exceptions;
2+
3+
import com.mindee.MindeeException;
4+
5+
/**
6+
* Represent an invalid or malformed input.
7+
*/
8+
public class MindeeInputException extends MindeeException {
9+
10+
/**
11+
* {@link Exception}
12+
*/
13+
public MindeeInputException(String message) {
14+
super(message);
15+
}
16+
17+
/**
18+
* {@link Exception}
19+
*/
20+
public MindeeInputException(String message, Exception innerException) {
21+
super(message, innerException);
22+
}
23+
}

‎src/main/java/com/mindee/v2/MindeeClient.java‎

Lines changed: 238 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@
33
import com.mindee.MindeeException;
44
import com.mindee.input.LocalInputSource;
55
import com.mindee.input.URLInputSource;
6+
import com.mindee.v2.clientoptions.BaseAnnotationParameters;
67
import com.mindee.v2.clientoptions.BaseProductParameters;
8+
import com.mindee.v2.clientoptions.BaseRagDocumentUploadParameters;
79
import com.mindee.v2.clientoptions.BaseSearchParameters;
810
import com.mindee.v2.clientoptions.PollingOptions;
911
import com.mindee.v2.http.MindeeApiV2;
1012
import com.mindee.v2.http.MindeeHttpApiV2;
1113
import com.mindee.v2.http.MindeeHttpExceptionV2;
14+
import com.mindee.v2.parsing.BaseRagAnnotationResponse;
1215
import com.mindee.v2.parsing.BaseResponse;
1316
import com.mindee.v2.parsing.JobResponse;
1417
import com.mindee.v2.parsing.error.ErrorResponse;
@@ -78,10 +81,10 @@ public JobResponse enqueue(
7881
* Can be used for polling.
7982
*/
8083
public JobResponse getJobFromUrl(String pollingUrl) {
81-
logger.log(System.Logger.Level.INFO, "Getting Job at: {0}", pollingUrl);
8284
if (pollingUrl == null || pollingUrl.isBlank()) {
8385
throw new IllegalArgumentException("Job URL cannot be null or blank.");
8486
}
87+
logger.log(System.Logger.Level.INFO, "Getting Job at: {0}", pollingUrl);
8588
return mindeeApi.reqGetJobByUrl(pollingUrl);
8689
}
8790

@@ -90,10 +93,10 @@ public JobResponse getJobFromUrl(String pollingUrl) {
9093
* Can be used for polling.
9194
*/
9295
public JobResponse getJob(String jobId) {
93-
logger.log(System.Logger.Level.INFO, "Getting job ID: {0}", jobId);
9496
if (jobId == null || jobId.isBlank()) {
9597
throw new IllegalArgumentException("jobId must not be null or blank.");
9698
}
99+
logger.log(System.Logger.Level.INFO, "Getting job ID: {0}", jobId);
97100
return mindeeApi.reqGetJobById(jobId);
98101
}
99102

@@ -105,10 +108,10 @@ public <TResponse extends BaseResponse> TResponse getResult(
105108
Class<TResponse> responseClass,
106109
String inferenceId
107110
) {
108-
logger.log(System.Logger.Level.INFO, "Getting result with ID: {0}", inferenceId);
109-
if (inferenceId == null || inferenceId.trim().isEmpty()) {
111+
if (inferenceId == null || inferenceId.isBlank()) {
110112
throw new IllegalArgumentException("inferenceId must not be null or blank.");
111113
}
114+
logger.log(System.Logger.Level.INFO, "Getting result with ID: {0}", inferenceId);
112115
return mindeeApi.reqGetResultById(responseClass, inferenceId);
113116
}
114117

@@ -120,10 +123,10 @@ public <TResponse extends BaseResponse> TResponse getResultFromUrl(
120123
Class<TResponse> responseClass,
121124
String inferenceUrl
122125
) {
123-
logger.log(System.Logger.Level.INFO, "Getting result at: {0}", inferenceUrl);
124-
if (inferenceUrl == null || inferenceUrl.trim().isEmpty()) {
126+
if (inferenceUrl == null || inferenceUrl.isBlank()) {
125127
throw new IllegalArgumentException("inferenceUrl must not be null or blank.");
126128
}
129+
logger.log(System.Logger.Level.INFO, "Getting result at: {0}", inferenceUrl);
127130
return mindeeApi.reqGetResultByUrl(responseClass, inferenceUrl);
128131
}
129132

@@ -277,6 +280,235 @@ public SearchResponse searchModels(String modelName, String modelType) {
277280
.reqGetSearch(ModelSearchParameters.builder().name(modelName).modelType(modelType).build());
278281
}
279282

283+
/**
284+
* Not recommended for general use, prefer {@link #uploadAndGetRagDocument}.
285+
* You will need to poll until the document is ready for use.
286+
* Add a document to the RAG database.
287+
*
288+
* @param inputSource The file to upload.
289+
* @param parameters RAG document upload parameters.
290+
* @return an instance of {@link BaseRagAnnotationResponse}.
291+
*/
292+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse uploadRagDocument(
293+
LocalInputSource inputSource,
294+
BaseRagDocumentUploadParameters<TAnnotationResponse> parameters
295+
) throws IOException {
296+
logger.log(System.Logger.Level.INFO, "Adding a document to the RAG database");
297+
return mindeeApi.reqPostRagDocument(parameters, inputSource);
298+
}
299+
300+
/**
301+
* Add a document to the RAG database, poll, and return the initial annotation.
302+
*
303+
* @param inputSource The file to upload.
304+
* @param parameters RAG document upload parameters.
305+
* @return an instance of {@link BaseRagAnnotationResponse}.
306+
*/
307+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse uploadAndGetRagDocument(
308+
LocalInputSource inputSource,
309+
BaseRagDocumentUploadParameters<TAnnotationResponse> parameters
310+
) throws IOException, InterruptedException {
311+
return uploadAndGetRagDocument(inputSource, parameters, null);
312+
}
313+
314+
/**
315+
* Add a document to the RAG database, poll, and return the initial annotation.
316+
*
317+
* @param inputSource The file to upload.
318+
* @param parameters RAG document upload parameters.
319+
* @param pollingOptions Polling options (if null, default options are used).
320+
* @return an instance of {@link BaseRagAnnotationResponse}.
321+
*/
322+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse uploadAndGetRagDocument(
323+
LocalInputSource inputSource,
324+
BaseRagDocumentUploadParameters<TAnnotationResponse> parameters,
325+
PollingOptions pollingOptions
326+
) throws IOException, InterruptedException {
327+
if (pollingOptions == null) {
328+
pollingOptions = PollingOptions.builder().build();
329+
}
330+
TAnnotationResponse initialResponse = uploadRagDocument(inputSource, parameters);
331+
if (!"Processing".equals(initialResponse.getStatus())) {
332+
return initialResponse;
333+
}
334+
return pollForRagDocument(parameters.getResponseClass(), initialResponse, pollingOptions);
335+
}
336+
337+
/**
338+
* Not recommended for general use, prefer {@link #getReadyRagDocument}.
339+
* You will need to poll until the document is ready for use.
340+
* Get a document's info and annotations from the RAG database.
341+
*
342+
* @param responseClass The class of the response.
343+
* @param documentId The document's ID.
344+
* @return an instance of {@link BaseRagAnnotationResponse}.
345+
*/
346+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse getRagDocument(
347+
Class<TAnnotationResponse> responseClass,
348+
String documentId
349+
) {
350+
if (documentId == null || documentId.isBlank()) {
351+
throw new IllegalArgumentException("documentId must not be null or blank.");
352+
}
353+
logger.log(System.Logger.Level.INFO, "Getting RAG document ID: {0}", documentId);
354+
return mindeeApi.reqGetRagAnnotation(responseClass, documentId);
355+
}
356+
357+
/**
358+
* Get a document's info and annotations from the RAG database, polling if it is still processing.
359+
*
360+
* @param responseClass The class of the response.
361+
* @param documentId The document's ID.
362+
* @return an instance of {@link BaseRagAnnotationResponse}.
363+
*/
364+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse getReadyRagDocument(
365+
Class<TAnnotationResponse> responseClass,
366+
String documentId
367+
) throws InterruptedException {
368+
return getReadyRagDocument(responseClass, documentId, null);
369+
}
370+
371+
/**
372+
* Get a document's info and annotations from the RAG database, polling if it is still processing.
373+
*
374+
* @param responseClass The class of the response.
375+
* @param documentId The document's ID.
376+
* @param pollingOptions Polling options (if null, default options are used).
377+
* @return an instance of {@link BaseRagAnnotationResponse}.
378+
*/
379+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse getReadyRagDocument(
380+
Class<TAnnotationResponse> responseClass,
381+
String documentId,
382+
PollingOptions pollingOptions
383+
) throws InterruptedException {
384+
TAnnotationResponse initialResponse = getRagDocument(responseClass, documentId);
385+
if (!"Processing".equals(initialResponse.getStatus())) {
386+
return initialResponse;
387+
}
388+
389+
if (pollingOptions == null) {
390+
pollingOptions = PollingOptions.builder().build();
391+
}
392+
return pollForRagDocument(responseClass, initialResponse, pollingOptions);
393+
}
394+
395+
/**
396+
* Not recommended for general use, prefer {@link #updateAndGetRagAnnotation}.
397+
* You will need to poll until the document is ready for use.
398+
* Update a document's annotations in the RAG database.
399+
*
400+
* @param parameters Annotation parameters.
401+
* @return an instance of {@link BaseRagAnnotationResponse}.
402+
*/
403+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse updateRagAnnotation(
404+
BaseAnnotationParameters<TAnnotationResponse> parameters
405+
) {
406+
logger
407+
.log(System.Logger.Level.INFO, "Updating RAG document ID: {0}", parameters.getDocumentId());
408+
return mindeeApi.reqPatchRagAnnotation(parameters);
409+
}
410+
411+
/**
412+
* Update a document's annotations in the RAG database and poll until complete.
413+
*
414+
* @param parameters Annotation parameters.
415+
* @return an instance of {@link BaseRagAnnotationResponse}.
416+
*/
417+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse updateAndGetRagAnnotation(
418+
BaseAnnotationParameters<TAnnotationResponse> parameters
419+
) throws InterruptedException {
420+
return updateAndGetRagAnnotation(parameters, null);
421+
}
422+
423+
/**
424+
* Update a document's annotations in the RAG database and poll until complete.
425+
*
426+
* @param parameters Annotation parameters.
427+
* @param pollingOptions Polling options (if null, default options are used).
428+
* @return an instance of {@link BaseRagAnnotationResponse}.
429+
*/
430+
public <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse updateAndGetRagAnnotation(
431+
BaseAnnotationParameters<TAnnotationResponse> parameters,
432+
PollingOptions pollingOptions
433+
) throws InterruptedException {
434+
TAnnotationResponse initialResponse = updateRagAnnotation(parameters);
435+
if (!"Processing".equals(initialResponse.getStatus())) {
436+
return initialResponse;
437+
}
438+
439+
if (pollingOptions == null) {
440+
pollingOptions = PollingOptions.builder().build();
441+
}
442+
return pollForRagDocument(parameters.getResponseClass(), initialResponse, pollingOptions);
443+
}
444+
445+
/**
446+
* Delete a document from the RAG database.
447+
* For extraction models only.
448+
*
449+
* @param documentId The document's ID.
450+
* @return true if successful.
451+
*/
452+
public boolean deleteExtractionRagDocument(String documentId) {
453+
if (documentId == null || documentId.isBlank()) {
454+
throw new IllegalArgumentException("documentId must not be null or blank.");
455+
}
456+
logger.log(System.Logger.Level.INFO, "Deleting RAG document ID: {0}", documentId);
457+
return mindeeApi.reqDeleteExtractionRagDocument(documentId);
458+
}
459+
460+
/**
461+
* Poll until the RAG document is finished processing or the max number of attempts is reached.
462+
*
463+
* @param responseClass The class of the response.
464+
* @param initialResponse The initial annotation response.
465+
* @param pollingOptions Polling options.
466+
* @return an instance of {@link BaseRagAnnotationResponse}.
467+
* @throws InterruptedException Throws if the thread is interrupted.
468+
*/
469+
private <TAnnotationResponse extends BaseRagAnnotationResponse> TAnnotationResponse pollForRagDocument(
470+
Class<TAnnotationResponse> responseClass,
471+
BaseRagAnnotationResponse initialResponse,
472+
PollingOptions pollingOptions
473+
) throws InterruptedException {
474+
logger
475+
.log(System.Logger.Level.INFO, "Polling for RAG document ID: {0}", initialResponse.getId());
476+
int maxRetries = pollingOptions.getMaxRetries() + 1;
477+
478+
logger
479+
.log(
480+
System.Logger.Level.DEBUG,
481+
"Waiting {0} seconds before attempting to retrieve the result...",
482+
pollingOptions.getInitialDelaySec()
483+
);
484+
485+
interruptibleSleep(
486+
(long) (pollingOptions.getInitialDelaySec() * 1000),
487+
pollingOptions.getCancelToken()
488+
);
489+
490+
String documentId = initialResponse.getId();
491+
long intervalMillis = (long) (pollingOptions.getIntervalSec() * 1000);
492+
int retryCount = 1;
493+
494+
while (retryCount < maxRetries) {
495+
logger.log(System.Logger.Level.DEBUG, "Poll attempt {0} of {1}", retryCount, maxRetries);
496+
497+
TAnnotationResponse response = getRagDocument(responseClass, documentId);
498+
retryCount++;
499+
500+
String status = response.getStatus();
501+
if ("Processing".equals(status)) {
502+
interruptibleSleep(intervalMillis, pollingOptions.getCancelToken());
503+
} else if ("Failed".equals(status)) {
504+
throw new MindeeException("RAG failed without an error payload.");
505+
} else {
506+
return response;
507+
}
508+
}
509+
throw new MindeeException("RAG polling not complete after " + retryCount + " attempts.");
510+
}
511+
280512
/**
281513
* Common logic for polling an asynchronous job for local & url files.
282514
*

0 commit comments

Comments
 (0)