Skip to content

Commit

Permalink
Merge pull request #21 from cloudera/mob/main
Browse files Browse the repository at this point in the history
Reconciler exception handling and updated python prefixes
  • Loading branch information
jkwatson authored Nov 19, 2024
2 parents 62b09db + bd3d2fc commit a595674
Show file tree
Hide file tree
Showing 26 changed files with 358 additions and 168 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -106,8 +106,6 @@ public HttpClient httpClient(OpenTelemetry openTelemetry) {
}

public static String getRagIndexUrl() {
String ragBackendBaseUrl =
Optional.ofNullable(System.getenv("LLM_SERVICE_URL")).orElse("http://rag-backend:8000");
return ragBackendBaseUrl + "/index";
return Optional.ofNullable(System.getenv("LLM_SERVICE_URL")).orElse("http://rag-backend:8000");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,11 @@ public void indexFile(
Types.RagDocument ragDocument, String bucketName, IndexConfiguration configuration) {
try {
client.post(
indexUrl + "/download-and-index",
new IndexRequest(
bucketName, ragDocument.s3Path(), ragDocument.dataSourceId(), configuration));
indexUrl
+ "/data_sources/"
+ ragDocument.dataSourceId()
+ "/documents/download-and-index",
new IndexRequest(bucketName, ragDocument.s3Path(), configuration));
} catch (IOException e) {
throw new RuntimeException(e);
}
Expand Down Expand Up @@ -97,7 +99,6 @@ public void deleteSession(Long sessionId) {
record IndexRequest(
@JsonProperty("s3_bucket_name") String s3BucketName,
@JsonProperty("s3_document_key") String s3DocumentKey,
@JsonProperty("data_source_id") long dataSourceId,
IndexConfiguration configuration) {}

public record SummaryRequest(
Expand All @@ -114,8 +115,12 @@ public static RagBackendClient createNull() {
return createNull(new Tracker<>());
}

public static RagBackendClient createNull(Tracker<TrackedRequest<?>> tracker) {
public static RagBackendClient createNull(
Tracker<TrackedRequest<?>> tracker, RuntimeException... t) {
return new RagBackendClient(SimpleHttpClient.createNull()) {
private final RuntimeException[] exceptions = t;
private int exceptionIndex = 0;

@Override
public void indexFile(
Types.RagDocument ragDocument, String bucketName, IndexConfiguration configuration) {
Expand All @@ -124,32 +129,43 @@ public void indexFile(
new TrackedRequest<>(
new TrackedIndexRequest(
bucketName, ragDocument.s3Path(), ragDocument.dataSourceId(), configuration)));
checkForException();
}

private void checkForException() {
if (exceptionIndex < exceptions.length) {
throw exceptions[exceptionIndex++];
}
}

@Override
public void deleteDataSource(Long dataSourceId) {
super.deleteDataSource(dataSourceId);
tracker.track(new TrackedRequest<>(new TrackedDeleteDataSourceRequest(dataSourceId)));
checkForException();
}

@Override
public String createSummary(Types.RagDocument ragDocument, String bucketName) {
String result = super.createSummary(ragDocument, bucketName);
tracker.track(new TrackedRequest<>(new SummaryRequest(bucketName, ragDocument.s3Path())));
checkForException();
return result;
}

@Override
public void deleteSession(Long sessionId) {
super.deleteSession(sessionId);
tracker.track(new TrackedRequest<>(new TrackedDeleteSessionRequest(sessionId)));
checkForException();
}

@Override
public void deleteDocument(long dataSourceId, String documentId) {
super.deleteDocument(dataSourceId, documentId);
tracker.track(
new TrackedRequest<>(new TrackedDeleteDocumentRequest(dataSourceId, documentId)));
checkForException();
}
};
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.cloudera.cai.rag.Types;
import com.cloudera.cai.rag.configuration.JdbiConfiguration;
import com.cloudera.cai.rag.external.RagBackendClient;
import com.cloudera.cai.util.exceptions.NotFound;
import com.cloudera.cai.util.reconcilers.BaseReconciler;
import com.cloudera.cai.util.reconcilers.ReconcileResult;
import com.cloudera.cai.util.reconcilers.ReconcilerConfig;
Expand Down Expand Up @@ -47,8 +48,12 @@ public void resync() throws Exception {
public ReconcileResult reconcile(Set<Types.RagDocument> documents) throws Exception {
for (Types.RagDocument document : documents) {
log.info("starting deletion of document: {}", document);
ragBackendClient.deleteDocument(document.dataSourceId(), document.documentId());
log.info("finished requesting deletion of document {}", document);
try {
ragBackendClient.deleteDocument(document.dataSourceId(), document.documentId());
log.info("finished requesting deletion of document {}", document);
} catch (NotFound e) {
log.info("got a not found exception from the rag backend: {}", e.getMessage());
}
jdbi.useHandle(
handle ->
handle.execute("DELETE from rag_data_source_document WHERE id = ?", document.id()));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
import com.cloudera.cai.rag.configuration.JdbiConfiguration;
import com.cloudera.cai.rag.datasources.RagDataSourceRepository;
import com.cloudera.cai.rag.external.RagBackendClient;
import com.cloudera.cai.util.exceptions.NotFound;
import com.cloudera.cai.util.reconcilers.*;
import io.opentelemetry.api.OpenTelemetry;
import java.time.Instant;
Expand Down Expand Up @@ -106,8 +107,7 @@ public ReconcileResult reconcile(Set<RagDocument> documents) {
for (RagDocument document : documents) {
log.info("starting indexing document: {}", document);
IndexConfiguration indexConfiguration = fetchIndexConfiguration(document.dataSourceId());
ragBackendClient.indexFile(document, bucketName, indexConfiguration);
log.info("finished requesting indexing of file {}", document);
Instant updateTimestamp = indexFile(document, indexConfiguration);
String updateSql =
"""
UPDATE rag_data_source_document
Expand All @@ -117,13 +117,26 @@ public ReconcileResult reconcile(Set<RagDocument> documents) {
jdbi.useHandle(
handle -> {
try (Update update = handle.createUpdate(updateSql)) {
update.bind("id", document.id()).bind("now", Instant.now()).execute();
update.bind("id", document.id()).bind("now", updateTimestamp).execute();
}
});
}
return new ReconcileResult();
}

private Instant indexFile(RagDocument document, IndexConfiguration indexConfiguration) {
Instant updateTimestamp;
try {
ragBackendClient.indexFile(document, bucketName, indexConfiguration);
updateTimestamp = Instant.now();
log.info("finished requesting indexing of file {}", document);
} catch (NotFound e) {
updateTimestamp = Instant.EPOCH;
log.info("not found exception from RAG Backend: {}", e.getMessage());
}
return updateTimestamp;
}

private IndexConfiguration fetchIndexConfiguration(Long dataSourceId) {
RagDataSource dataSource = ragDataSourceRepository.getRagDataSourceById(dataSourceId);
return new IndexConfiguration(dataSource.chunkSize(), dataSource.chunkOverlapPercent());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@

import com.cloudera.cai.rag.configuration.JdbiConfiguration;
import com.cloudera.cai.rag.external.RagBackendClient;
import com.cloudera.cai.util.exceptions.NotFound;
import com.cloudera.cai.util.reconcilers.*;
import io.opentelemetry.api.OpenTelemetry;
import java.time.Instant;
Expand Down Expand Up @@ -108,8 +109,7 @@ public ReconcileResult reconcile(Set<RagDocument> documents) {
log.info("Document already summarized: {}", document.filename());
continue;
}
String summary = ragBackendClient.createSummary(document, bucketName);
log.info("Doc summary: {}", summary);
var updateTimestamp = generateSummary(document);
log.info("finished requesting summarization of file {}", document);
String updateSql =
"""
Expand All @@ -120,13 +120,25 @@ public ReconcileResult reconcile(Set<RagDocument> documents) {
jdbi.useHandle(
handle -> {
try (Update update = handle.createUpdate(updateSql)) {
update.bind("id", document.id()).bind("now", Instant.now()).execute();
update.bind("id", document.id()).bind("now", updateTimestamp).execute();
}
});
}
return new ReconcileResult();
}

private Instant generateSummary(RagDocument document) {
var updateTimestamp = Instant.now();
try {
String summary = ragBackendClient.createSummary(document, bucketName);
log.info("Doc summary: {}", summary);
} catch (NotFound e) {
updateTimestamp = Instant.EPOCH;
log.info("got a not found exception from the rag backend: {}", e.getMessage());
}
return updateTimestamp;
}

// nullables stuff below here...
public static RagFileSummaryReconciler createNull() {
return new RagFileSummaryReconciler(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@

package com.cloudera.cai.util;

import com.cloudera.cai.util.exceptions.NotFound;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
Expand Down Expand Up @@ -73,6 +74,10 @@ public <T> String post(String url, T bodyObject) throws IOException {
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());
int statusCode = response.statusCode();
if (statusCode == 404) {
throw new NotFound("Failed to post to " + url + " code: " + statusCode);
}

if (statusCode >= 400) {
throw new RuntimeException(
"Failed to post to " + url + " code: " + statusCode + ", body : " + response.body());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,14 @@
package com.cloudera.cai.rag.external;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import com.cloudera.cai.rag.Types.RagDocument;
import com.cloudera.cai.rag.external.RagBackendClient.IndexConfiguration;
import com.cloudera.cai.util.SimpleHttpClient;
import com.cloudera.cai.util.SimpleHttpClient.TrackedHttpRequest;
import com.cloudera.cai.util.Tracker;
import com.cloudera.cai.util.exceptions.NotFound;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.springframework.http.HttpMethod;
Expand All @@ -65,9 +67,8 @@ void indexFile() {
.contains(
new TrackedHttpRequest<>(
HttpMethod.POST,
"http://rag-backend:8000/index/download-and-index",
new RagBackendClient.IndexRequest(
"bucketName", "s3Path", 1234L, indexConfiguration)));
"http://rag-backend:8000/data_sources/" + 1234L + "/documents/download-and-index",
new RagBackendClient.IndexRequest("bucketName", "s3Path", indexConfiguration)));
}

@Test
Expand All @@ -84,7 +85,7 @@ void createSummary() {
.contains(
new TrackedHttpRequest<>(
HttpMethod.POST,
"http://rag-backend:8000/index/data_sources/1234/summarize-document",
"http://rag-backend:8000/data_sources/1234/summarize-document",
new RagBackendClient.SummaryRequest("bucketName", "s3Path")));
}

Expand All @@ -98,7 +99,7 @@ void deleteDataSource() {
.hasSize(1)
.contains(
new TrackedHttpRequest<>(
HttpMethod.DELETE, "http://rag-backend:8000/index/data_sources/1234", null));
HttpMethod.DELETE, "http://rag-backend:8000/data_sources/1234", null));
}

@Test
Expand All @@ -112,7 +113,7 @@ void deleteDocument() {
.contains(
new TrackedHttpRequest<>(
HttpMethod.DELETE,
"http://rag-backend:8000/index/data_sources/1234/documents/documentId",
"http://rag-backend:8000/data_sources/1234/documents/documentId",
null));
}

Expand All @@ -126,7 +127,16 @@ void deleteSession() {
.hasSize(1)
.contains(
new TrackedHttpRequest<>(
HttpMethod.DELETE, "http://rag-backend:8000/index/sessions/1234", null));
HttpMethod.DELETE, "http://rag-backend:8000/sessions/1234", null));
}

@Test
void null_handlesThrowable() {
RagBackendClient client =
RagBackendClient.createNull(new Tracker<>(), new NotFound("not found"));
RagDocument document = indexRequest("s3Path", 1234L);
assertThatThrownBy(() -> client.indexFile(document, "fakeit", null))
.isInstanceOf(NotFound.class);
}

private static RagDocument indexRequest(String s3Path, Long dataSourceId) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,40 @@ void reconciler() throws Exception {
});
}

private static RagFileDeleteReconciler createTestReconciler(Tracker<TrackedRequest<?>> tracker) {
@Test
void reconcile_notFound() throws Exception {
Tracker<TrackedRequest<?>> tracker = new Tracker<>();
RagFileDeleteReconciler reconciler =
createTestReconciler(tracker, new NotFound("data source not found"));

RagFileRepository ragFileRepository = RagFileRepository.createNull();
var dataSourceId = TestData.createTestDataSource(RagDataSourceRepository.createNull());
String documentId = UUID.randomUUID().toString();
var id = TestData.createTestDocument(dataSourceId, documentId, ragFileRepository);

ragFileRepository.deleteById(id);
reconciler.resync();

await()
.untilAsserted(
() -> {
assertThat(tracker.getValues())
.contains(
new TrackedRequest<>(
new TrackedDeleteDocumentRequest(dataSourceId, documentId)));
assertThatThrownBy(() -> ragFileRepository.getRagDocumentById(id))
.isInstanceOf(NotFound.class);
});
}

private static RagFileDeleteReconciler createTestReconciler(
Tracker<TrackedRequest<?>> tracker, RuntimeException... exceptions) {
RagFileDeleteReconciler reconciler =
new RagFileDeleteReconciler(
ReconcilerConfig.createTestConfig(),
OpenTelemetry.noop(),
JdbiConfiguration.createNull(),
RagBackendClient.createNull(tracker));
RagBackendClient.createNull(tracker, exceptions));
reconciler.init();
return reconciler;
}
Expand Down
Loading

0 comments on commit a595674

Please sign in to comment.