From 7207c41c88d50e856de71ec5590de2594adf684c Mon Sep 17 00:00:00 2001 From: Tommaso Bolis Date: Tue, 19 Nov 2024 19:21:36 +0100 Subject: [PATCH] Adding json output schema --- .../operation/EmbeddingOperations.java | 37 ++++++++----- .../EmbeddingListSourcesResponse.json | 54 +++++++++++++++++++ .../EmbeddingRemoveFromStoreResponse.json | 10 ++++ 3 files changed, 89 insertions(+), 12 deletions(-) create mode 100644 src/main/resources/api/response/EmbeddingListSourcesResponse.json create mode 100644 src/main/resources/api/response/EmbeddingRemoveFromStoreResponse.json diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/EmbeddingOperations.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/EmbeddingOperations.java index d637302..c12e07e 100644 --- a/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/EmbeddingOperations.java +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/EmbeddingOperations.java @@ -397,9 +397,11 @@ public class EmbeddingOperations { jsonObject.put(Constants.JSON_KEY_SOURCES, sources); - - - return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + return createEmbeddingResponse( + jsonObject.toString(), + new HashMap() {{ + put("storeName", storeName); + }}); } catch (ModuleException me) { throw me; @@ -490,7 +492,6 @@ public class EmbeddingOperations { JSONArray sources = new JSONArray(); JSONObject contentObject; - String fullPath; for (EmbeddingMatch match : embeddingMatches) { Metadata matchMetadata = match.embedded().metadata(); @@ -508,7 +509,13 @@ public class EmbeddingOperations { jsonObject.put(Constants.JSON_KEY_SOURCES, sources); - return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + + return createEmbeddingResponse( + jsonObject.toString(), + new HashMap() {{ + put("storeName", storeName); + put("filter", searchFilterParams.getFilterJSONObject()); + }}); } catch (ModuleException me) { throw me; @@ -522,7 +529,6 @@ public class EmbeddingOperations { } } - /** * Retrieves and lists sources from the specified embedding store. * @@ -539,6 +545,7 @@ public class EmbeddingOperations { @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-list-sources") @Throws(EmbeddingErrorTypeProvider.class) + @OutputJsonType(schema = "api/response/EmbeddingListSourcesResponse.json") public org.mule.runtime.extension.api.runtime.operation.Result listSourcesFromStore( String storeName, @Config Configuration configuration, @@ -560,7 +567,11 @@ public class EmbeddingOperations { JSONObject jsonObject = baseStore.listSources(); - return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + return createEmbeddingResponse( + jsonObject.toString(), + new HashMap() {{ + put("storeName", storeName); + }}); } catch (ModuleException me) { throw me; @@ -581,7 +592,7 @@ public class EmbeddingOperations { @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-remove-from-store-by-filter") @Throws(EmbeddingErrorTypeProvider.class) - + @OutputJsonType(schema = "api/response/EmbeddingRemoveFromStoreResponse.json") public org.mule.runtime.extension.api.runtime.operation.Result removeEmbeddingsByFilter( String storeName, @Config Configuration configuration, @@ -614,12 +625,14 @@ public class EmbeddingOperations { embeddingStore.removeAll(filter); JSONObject jsonObject = new JSONObject(); - jsonObject.put(Constants.JSON_KEY_STORE_NAME, storeName); jsonObject.put(Constants.JSON_KEY_STATUS, Constants.OPERATION_STATUS_DELETED); - //jsonObject.put("filter", removeFilterParams.getFilterJSONObject()); - - return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + return createEmbeddingResponse( + jsonObject.toString(), + new HashMap() {{ + put("storeName", storeName); + put("filter", removeFilterParams.getFilterJSONObject()); + }}); } catch (ModuleException me) { throw me; diff --git a/src/main/resources/api/response/EmbeddingListSourcesResponse.json b/src/main/resources/api/response/EmbeddingListSourcesResponse.json new file mode 100644 index 0000000..b55c666 --- /dev/null +++ b/src/main/resources/api/response/EmbeddingListSourcesResponse.json @@ -0,0 +1,54 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "sourceCount": { + "type": "integer" + }, + "sources": { + "type": "array", + "items": { + "type": "object", + "properties": { + "ingestion_timestamp": { + "type": "integer", + "format": "int64" + }, + "source_id": { + "type": "string", + "format": "uuid" + }, + "source": { + "type": "string" + }, + "url": { + "type": "string" + }, + "title": { + "type": "string" + }, + "segmentCount": { + "type": "integer" + }, + "ingestion_datetime": { + "type": "string", + "format": "date-time" + }, + "absolute_directory_path": { + "type": "string", + "optional": true + }, + "file_name": { + "type": "string", + "optional": true + }, + "file_type": { + "type": "string" + } + } + }, + "minItems": 0 + } + }, + "required": ["sourceCount", "sources"] +} diff --git a/src/main/resources/api/response/EmbeddingRemoveFromStoreResponse.json b/src/main/resources/api/response/EmbeddingRemoveFromStoreResponse.json new file mode 100644 index 0000000..6f9dbf5 --- /dev/null +++ b/src/main/resources/api/response/EmbeddingRemoveFromStoreResponse.json @@ -0,0 +1,10 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "status": { + "type": "string" + } + }, + "required": ["status"] +}