diff --git a/pom.xml b/pom.xml index 7c8abb7..77e7fe6 100644 --- a/pom.xml +++ b/pom.xml @@ -4,19 +4,21 @@ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> 4.0.0 - com.mule.mulechain - mulechain-vectors - 0.1.112-SNAPSHOT + com.mulesoft.connectors + mule4-vectors-connector + 0.1.113-SNAPSHOT mule-extension - Mule Vectors Connector - MAC Vectors provides access to a broad number of external Vector Stores and Databases. + MuleSoft Vectors Connector - Mule 4 + MuleSoft Vectors Connector provides access to a broad number of external Vector Stores. https://mac-project.ai/docs/mulechain-vectors/connector-overview + MIT License https://www.opensource.org/licenses/mit-license.php + Amir Khan @@ -26,7 +28,7 @@ Tommaso Bolis - tbolis@salesforce.com + tbolis-at-salesforce-dot-com Salesforce https://www.salesforce.com @@ -61,6 +63,7 @@ validate 0.35.0 2.0.7 + 0.9.0-rc1 @@ -79,38 +82,11 @@ - - org.mule mule-javaee-runtime-bom @@ -122,27 +98,30 @@ + dev.langchain4j langchain4j-azure-ai-search ${langchain4jVersion} + dev.langchain4j langchain4j-document-loader-azure-storage-blob ${langchain4jVersion} + dev.langchain4j langchain4j-document-loader-amazon-s3 ${langchain4jVersion} - - org.mule.sdk - mule-sdk-api - 0.9.0-rc1 - + + org.mule.sdk + mule-sdk-api + ${mule.sdk.api.version} + dev.langchain4j @@ -155,11 +134,13 @@ + org.jsoup jsoup 1.14.3 + org.slf4j slf4j-api @@ -170,6 +151,7 @@ json 20240303 + dev.langchain4j langchain4j-document-parser-apache-tika @@ -279,6 +261,7 @@ + mule-releases @@ -298,4 +281,5 @@ + diff --git a/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/DocumentResponseAttributes.java b/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/DocumentResponseAttributes.java new file mode 100644 index 0000000..3e60821 --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/DocumentResponseAttributes.java @@ -0,0 +1,14 @@ +package org.mule.extension.mulechain.vectors.api.metadata; + +import java.io.Serializable; +import java.util.HashMap; + +public class DocumentResponseAttributes implements Serializable { + + private final HashMap documentAttributes; + + public DocumentResponseAttributes(HashMap documentAttributes) { + + this.documentAttributes = documentAttributes; + } +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/EmbeddingResponseAttributes.java b/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/EmbeddingResponseAttributes.java new file mode 100644 index 0000000..8b54397 --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/api/metadata/EmbeddingResponseAttributes.java @@ -0,0 +1,14 @@ +package org.mule.extension.mulechain.vectors.api.metadata; + +import java.io.Serializable; +import java.util.HashMap; + +public class EmbeddingResponseAttributes implements Serializable { + + private final HashMap embeddingAttributes; + + public EmbeddingResponseAttributes(HashMap embeddingAttributes) { + + this.embeddingAttributes = embeddingAttributes; + } +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/error/MuleVectorsErrorType.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/MuleVectorsErrorType.java new file mode 100644 index 0000000..e762988 --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/MuleVectorsErrorType.java @@ -0,0 +1,11 @@ +package org.mule.extension.mulechain.vectors.internal.error; + +import org.mule.runtime.extension.api.error.ErrorTypeDefinition; + +/** + * (c) 2003-2024 MuleSoft, Inc. The software in this package is published under the terms of the Commercial Free Software license V.1 a copy of which has been included with this distribution in the LICENSE.md file. + */ +public enum MuleVectorsErrorType implements ErrorTypeDefinition { + + DOCUMENT_OPERATIONS_FAILURE, EMBEDDING_OPERATIONS_FAILURE +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/DocumentErrorTypeProvider.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/DocumentErrorTypeProvider.java new file mode 100644 index 0000000..e250d5e --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/DocumentErrorTypeProvider.java @@ -0,0 +1,23 @@ +/** + * (c) 2003-2024 MuleSoft, Inc. The software in this package is published under the terms of the Commercial Free Software license V.1 a copy of which has been included with this distribution in the LICENSE.md file. + */ +package org.mule.extension.mulechain.vectors.internal.error.provider; + +import org.mule.runtime.extension.api.annotation.error.ErrorTypeProvider; +import org.mule.runtime.extension.api.error.ErrorTypeDefinition; + +import java.util.HashSet; +import java.util.Set; + +import static java.util.Arrays.asList; +import static java.util.Collections.unmodifiableSet; +import static org.mule.extension.mulechain.vectors.internal.error.MuleVectorsErrorType.DOCUMENT_OPERATIONS_FAILURE; + +public class DocumentErrorTypeProvider implements ErrorTypeProvider { + + @SuppressWarnings("rawtypes") + @Override + public Set getErrorTypes() { + return unmodifiableSet(new HashSet<>(asList(DOCUMENT_OPERATIONS_FAILURE))); + } +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/EmbeddingErrorTypeProvider.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/EmbeddingErrorTypeProvider.java new file mode 100644 index 0000000..7c4979c --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/error/provider/EmbeddingErrorTypeProvider.java @@ -0,0 +1,23 @@ +/** + * (c) 2003-2024 MuleSoft, Inc. The software in this package is published under the terms of the Commercial Free Software license V.1 a copy of which has been included with this distribution in the LICENSE.md file. + */ +package org.mule.extension.mulechain.vectors.internal.error.provider; + +import org.mule.runtime.extension.api.annotation.error.ErrorTypeProvider; +import org.mule.runtime.extension.api.error.ErrorTypeDefinition; + +import java.util.HashSet; +import java.util.Set; + +import static java.util.Arrays.asList; +import static java.util.Collections.unmodifiableSet; +import static org.mule.extension.mulechain.vectors.internal.error.MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE; + +public class EmbeddingErrorTypeProvider implements ErrorTypeProvider { + + @SuppressWarnings("rawtypes") + @Override + public Set getErrorTypes() { + return unmodifiableSet(new HashSet<>(asList(EMBEDDING_OPERATIONS_FAILURE))); + } +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/extension/Connector.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/extension/Connector.java index f8e65e5..d368975 100644 --- a/src/main/java/org/mule/extension/mulechain/vectors/internal/extension/Connector.java +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/extension/Connector.java @@ -1,9 +1,13 @@ package org.mule.extension.mulechain.vectors.internal.extension; import org.mule.extension.mulechain.vectors.internal.config.Configuration; +import org.mule.extension.mulechain.vectors.internal.error.MuleVectorsErrorType; +import org.mule.runtime.api.meta.Category; import org.mule.runtime.extension.api.annotation.Extension; import org.mule.runtime.extension.api.annotation.Configurations; import org.mule.runtime.extension.api.annotation.dsl.xml.Xml; +import org.mule.runtime.extension.api.annotation.error.ErrorTypes; +import org.mule.runtime.extension.api.annotation.license.RequiresEnterpriseLicense; import org.mule.sdk.api.annotation.JavaVersionSupport; import static org.mule.sdk.api.meta.JavaVersion.JAVA_11; import static org.mule.sdk.api.meta.JavaVersion.JAVA_17; @@ -13,9 +17,11 @@ * This is the main class of an extension, is the entry point from which configurations, connection providers, operations * and sources are going to be declared. */ -@Xml(prefix = "vectors") -@Extension(name = "MAC Vectors") +@Xml(prefix = "ms-vectors") +@Extension(name = "MuleSoft Vectors Connector", category = Category.SELECT) @Configurations(Configuration.class) +@RequiresEnterpriseLicense(allowEvaluationLicense = true) +@ErrorTypes(MuleVectorsErrorType.class) @JavaVersionSupport({JAVA_8, JAVA_11, JAVA_17}) public class Connector { diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/helper/ResponseHelper.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/helper/ResponseHelper.java new file mode 100644 index 0000000..38a4119 --- /dev/null +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/helper/ResponseHelper.java @@ -0,0 +1,42 @@ +package org.mule.extension.mulechain.vectors.internal.helper; + +import org.mule.extension.mulechain.vectors.api.metadata.DocumentResponseAttributes; +import org.mule.extension.mulechain.vectors.api.metadata.EmbeddingResponseAttributes; +import org.mule.runtime.extension.api.runtime.operation.Result; + +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.Map; + +import static org.apache.commons.io.IOUtils.toInputStream; + +public final class ResponseHelper { + + private ResponseHelper() { + } + + public static Result createEmbeddingResponse( + String response, + Map embeddingAttributes) { + + return Result.builder() + .attributes(new EmbeddingResponseAttributes((HashMap) embeddingAttributes)) + .attributesMediaType(org.mule.runtime.api.metadata.MediaType.APPLICATION_JAVA) + .output(toInputStream(response, StandardCharsets.UTF_8)) + .mediaType(org.mule.runtime.api.metadata.MediaType.APPLICATION_JSON) + .build(); + } + + public static Result createDocumentResponse( + String response, + Map documentAttributes) { + + return Result.builder() + .attributes(new DocumentResponseAttributes((HashMap) documentAttributes)) + .attributesMediaType(org.mule.runtime.api.metadata.MediaType.APPLICATION_JAVA) + .output(toInputStream(response, StandardCharsets.UTF_8)) + .mediaType(org.mule.runtime.api.metadata.MediaType.APPLICATION_JSON) + .build(); + } +} diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/DocumentOperations.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/DocumentOperations.java index 4658a1a..c57c036 100644 --- a/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/DocumentOperations.java +++ b/src/main/java/org/mule/extension/mulechain/vectors/internal/operation/DocumentOperations.java @@ -9,14 +9,19 @@ import dev.langchain4j.data.document.transformer.jsoup.HtmlToTextDocumentTransformer; import dev.langchain4j.data.segment.TextSegment; import org.json.JSONObject; +import org.mule.extension.mulechain.vectors.api.metadata.DocumentResponseAttributes; import org.mule.extension.mulechain.vectors.internal.constant.Constants; +import org.mule.extension.mulechain.vectors.internal.error.MuleVectorsErrorType; +import org.mule.extension.mulechain.vectors.internal.error.provider.DocumentErrorTypeProvider; import org.mule.extension.mulechain.vectors.internal.helper.parameter.FileTypeParameters; import org.mule.extension.mulechain.vectors.internal.helper.parameter.SegmentationParameters; import org.mule.runtime.extension.api.annotation.Alias; +import org.mule.runtime.extension.api.annotation.error.Throws; import org.mule.runtime.extension.api.annotation.param.MediaType; import org.mule.runtime.extension.api.annotation.param.ParameterGroup; import org.mule.runtime.extension.api.annotation.param.display.DisplayName; import org.mule.runtime.extension.api.annotation.param.display.Summary; +import org.mule.runtime.extension.api.exception.ModuleException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -25,107 +30,131 @@ import java.net.URL; import java.nio.charset.StandardCharsets; import java.util.Arrays; +import java.util.HashMap; import java.util.List; import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument; -import static org.apache.commons.io.IOUtils.toInputStream; +import static org.mule.extension.mulechain.vectors.internal.helper.ResponseHelper.createDocumentResponse; import static org.mule.runtime.extension.api.annotation.param.MediaType.APPLICATION_JSON; public class DocumentOperations { - private static final Logger LOGGER = LoggerFactory.getLogger(DocumentOperations.class); - - /** - * Splits a document provided by full path in to a defined set of chucks and overlaps - */ - @MediaType(value = APPLICATION_JSON, strict = false) - @Alias("Document-split-into-chunks") - public InputStream documentSplitter(@Alias("contextPath") @DisplayName("Context Path") @Summary("The context path.") String contextPath, - @ParameterGroup(name = "Context") FileTypeParameters fileType, - @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters){ - - List segments; - DocumentSplitter splitter; - Document document = null; - switch (fileType.getFileType()) { - case Constants.FILE_TYPE_TEXT: - document = loadDocument(contextPath, new TextDocumentParser()); - splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); - segments = splitter.split(document); - break; - case Constants.FILE_TYPE_ANY: - document = loadDocument(contextPath, new ApacheTikaDocumentParser()); - splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); - segments = splitter.split(document); - break; - case Constants.FILE_TYPE_URL: - URL url = null; - try { - url = new URL(contextPath); - } catch (MalformedURLException e) { - LOGGER.error(e.getMessage() + " " + Arrays.toString(e.getStackTrace())); - } - - Document htmlDocument = UrlDocumentLoader.load(url, new TextDocumentParser()); - HtmlToTextDocumentTransformer transformer = new HtmlToTextDocumentTransformer(null, null, true); - document = transformer.transform(htmlDocument); - document.metadata().put(Constants.METADATA_KEY_URL, contextPath); - splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); - segments = splitter.split(document); - break; - default: - throw new IllegalArgumentException("Unsupported File Type: " + fileType.getFileType()); - } - - JSONObject jsonObject = new JSONObject(); - jsonObject.put("contextPath", contextPath); - jsonObject.put("fileType", fileType.getFileType()); - jsonObject.put("segments", segments.toString()); - - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + private static final Logger LOGGER = LoggerFactory.getLogger(DocumentOperations.class); + + /** + * Splits a document provided by full path in to a defined set of chucks and overlaps + */ + @MediaType(value = APPLICATION_JSON, strict = false) + @Alias("Document-split-into-chunks") + @Throws(DocumentErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + documentSplitter( @Alias("contextPath") @DisplayName("Context Path") @Summary("The context path.") String contextPath, + @ParameterGroup(name = "Context") FileTypeParameters fileType, + @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters){ + + try { + + List segments; + DocumentSplitter splitter; + Document document = null; + switch (fileType.getFileType()) { + case Constants.FILE_TYPE_TEXT: + document = loadDocument(contextPath, new TextDocumentParser()); + splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); + segments = splitter.split(document); + break; + case Constants.FILE_TYPE_ANY: + document = loadDocument(contextPath, new ApacheTikaDocumentParser()); + splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); + segments = splitter.split(document); + break; + case Constants.FILE_TYPE_URL: + URL url = null; + try { + url = new URL(contextPath); + } catch (MalformedURLException e) { + LOGGER.error(e.getMessage() + " " + Arrays.toString(e.getStackTrace())); + } + + Document htmlDocument = UrlDocumentLoader.load(url, new TextDocumentParser()); + HtmlToTextDocumentTransformer transformer = new HtmlToTextDocumentTransformer(null, null, true); + document = transformer.transform(htmlDocument); + document.metadata().put(Constants.METADATA_KEY_URL, contextPath); + splitter = DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars()); + segments = splitter.split(document); + break; + default: + throw new IllegalArgumentException("Unsupported File Type: " + fileType.getFileType()); + } + + JSONObject jsonObject = new JSONObject(); + jsonObject.put("contextPath", contextPath); + jsonObject.put("fileType", fileType.getFileType()); + jsonObject.put("segments", segments.toString()); + + return createDocumentResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while splitting document %s.", contextPath), + MuleVectorsErrorType.DOCUMENT_OPERATIONS_FAILURE, + e); } - - /** - * Parses a document by filepath and returns the text - */ - @MediaType(value = APPLICATION_JSON, strict = false) - @Alias("Document-parser") - public InputStream documentParser(@Alias("contextPath") @DisplayName("Context Path") @Summary("The context path.") String contextPath, - @ParameterGroup(name = "Context") FileTypeParameters fileType){ - - Document document = null; - switch (fileType.getFileType()) { - case Constants.FILE_TYPE_TEXT: - document = loadDocument(contextPath, new TextDocumentParser()); - break; - case Constants.FILE_TYPE_ANY: - document = loadDocument(contextPath, new ApacheTikaDocumentParser()); - break; - case Constants.FILE_TYPE_URL: - URL url = null; - try { - url = new URL(contextPath); - } catch (MalformedURLException e) { - LOGGER.error(e.getMessage() + " " + Arrays.toString(e.getStackTrace())); - } - - Document htmlDocument = UrlDocumentLoader.load(url, new TextDocumentParser()); - HtmlToTextDocumentTransformer transformer = new HtmlToTextDocumentTransformer(null, null, true); - document = transformer.transform(htmlDocument); - document.metadata().put(Constants.METADATA_KEY_URL, contextPath); - - break; - default: - throw new IllegalArgumentException("Unsupported File Type: " + fileType.getFileType()); - } - - JSONObject jsonObject = new JSONObject(); - jsonObject.put("contextPath", contextPath); - jsonObject.put("fileType", fileType.getFileType()); - jsonObject.put("documentText",document.text()); - jsonObject.put("metadata",document.metadata()); - - - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + } + + /** + * Parses a document by filepath and returns the text + */ + @MediaType(value = APPLICATION_JSON, strict = false) + @Alias("Document-parser") + @Throws(DocumentErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + documentParser( @Alias("contextPath") @DisplayName("Context Path") @Summary("The context path.") String contextPath, + @ParameterGroup(name = "Context") FileTypeParameters fileType){ + + try { + + Document document = null; + switch (fileType.getFileType()) { + case Constants.FILE_TYPE_TEXT: + document = loadDocument(contextPath, new TextDocumentParser()); + break; + case Constants.FILE_TYPE_ANY: + document = loadDocument(contextPath, new ApacheTikaDocumentParser()); + break; + case Constants.FILE_TYPE_URL: + URL url = null; + try { + url = new URL(contextPath); + } catch (MalformedURLException e) { + LOGGER.error(e.getMessage() + " " + Arrays.toString(e.getStackTrace())); + } + + Document htmlDocument = UrlDocumentLoader.load(url, new TextDocumentParser()); + HtmlToTextDocumentTransformer transformer = new HtmlToTextDocumentTransformer(null, null, true); + document = transformer.transform(htmlDocument); + document.metadata().put(Constants.METADATA_KEY_URL, contextPath); + + break; + default: + throw new IllegalArgumentException("Unsupported File Type: " + fileType.getFileType()); + } + + JSONObject jsonObject = new JSONObject(); + jsonObject.put("contextPath", contextPath); + jsonObject.put("fileType", fileType.getFileType()); + jsonObject.put("documentText",document.text()); + jsonObject.put("metadata",document.metadata()); + + return createDocumentResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while splitting document %s.", contextPath), + MuleVectorsErrorType.DOCUMENT_OPERATIONS_FAILURE, + e); } + } } 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 816d1b0..80b0758 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 @@ -1,14 +1,17 @@ package org.mule.extension.mulechain.vectors.internal.operation; -import static org.apache.commons.io.IOUtils.toInputStream; +import static org.mule.extension.mulechain.vectors.internal.helper.ResponseHelper.createEmbeddingResponse; import static org.mule.runtime.extension.api.annotation.param.MediaType.APPLICATION_JSON; import java.io.InputStream; +import java.util.HashMap; import java.util.List; -import java.nio.charset.StandardCharsets; import dev.langchain4j.data.document.Document; +import org.mule.extension.mulechain.vectors.api.metadata.EmbeddingResponseAttributes; import org.mule.extension.mulechain.vectors.internal.constant.Constants; +import org.mule.extension.mulechain.vectors.internal.error.MuleVectorsErrorType; +import org.mule.extension.mulechain.vectors.internal.error.provider.EmbeddingErrorTypeProvider; import org.mule.extension.mulechain.vectors.internal.helper.EmbeddingOperationValidator; import org.mule.extension.mulechain.vectors.internal.helper.parameter.*; import org.mule.extension.mulechain.vectors.internal.config.Configuration; @@ -21,6 +24,7 @@ import org.mule.extension.mulechain.vectors.internal.store.BaseStore; import org.mule.extension.mulechain.vectors.internal.util.JsonUtils; import org.mule.runtime.extension.api.annotation.Alias; +import org.mule.runtime.extension.api.annotation.error.Throws; import org.mule.runtime.extension.api.annotation.param.*; import static java.util.stream.Collectors.joining; @@ -31,6 +35,7 @@ import dev.langchain4j.model.embedding.EmbeddingModel; import org.mule.runtime.extension.api.annotation.param.display.DisplayName; +import org.mule.runtime.extension.api.exception.ModuleException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -46,37 +51,50 @@ public class EmbeddingOperations { */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("Embedding-add-text-to-store") - public InputStream addTextToStore( @Config Configuration configuration, - @Alias("text") @DisplayName("Text") String text, - @Alias("storeName") @DisplayName("Store Name") String storeName, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ - - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); - - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); - - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - - TextSegment textSegment = TextSegment.from(text); - Embedding textEmbedding = embeddingModel.embed(textSegment).content(); - embeddingStore.add(textEmbedding, textSegment); - - JSONObject jsonObject = new JSONObject(); - jsonObject.put("status", Constants.OPERATION_STATUS_ADDED); - jsonObject.put("textSegment", textSegment.toString()); - jsonObject.put("textEmbedding", textEmbedding.toString()); - jsonObject.put("storeName", storeName); - - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + @Throws(EmbeddingErrorTypeProvider.class) + + public org.mule.runtime.extension.api.runtime.operation.Result + addTextToStore( @Config Configuration configuration, + @Alias("text") @DisplayName("Text") String text, + @Alias("storeName") @DisplayName("Store Name") String storeName, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ + + try { + + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); + + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); + + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + + TextSegment textSegment = TextSegment.from(text); + Embedding textEmbedding = embeddingModel.embed(textSegment).content(); + embeddingStore.add(textEmbedding, textSegment); + + JSONObject jsonObject = new JSONObject(); + jsonObject.put("status", Constants.OPERATION_STATUS_ADDED); + jsonObject.put("textSegment", textSegment.toString()); + jsonObject.put("textEmbedding", textEmbedding.toString()); + jsonObject.put("storeName", storeName); + + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while adding text %s into the store %s", text, storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } /** @@ -84,27 +102,38 @@ public InputStream addTextToStore( @Config Configuration configuration, */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("Embedding-generate-from-text") - public InputStream generateEmbedding( @Config Configuration configuration, - @Alias("text") @DisplayName("Text") String text, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ + @Throws(EmbeddingErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + generateEmbedding(@Config Configuration configuration, + @Alias("text") @DisplayName("Text") String text, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ + + try { + + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + TextSegment textSegment = TextSegment.from(text); + Embedding textEmbedding = embeddingModel.embed(textSegment).content(); - TextSegment textSegment = TextSegment.from(text); - Embedding textEmbedding = embeddingModel.embed(textSegment).content(); + JSONObject jsonObject = new JSONObject(); + jsonObject.put("Segment", textSegment.toString()); + jsonObject.put("Embedding", textEmbedding.toString()); + jsonObject.put("Dimension", textEmbedding.dimension()); - JSONObject jsonObject = new JSONObject(); - jsonObject.put("Segment", textSegment.toString()); - jsonObject.put("Embedding", textEmbedding.toString()); - jsonObject.put("Dimension", textEmbedding.dimension()); + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + } catch (Exception e) { - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + throw new ModuleException( + String.format("Error while generating embedding from text %s", text), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } /** @@ -112,53 +141,66 @@ public InputStream generateEmbedding( @Config Configuration configuration, */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("Embedding-add-folder-to-store") - public InputStream addFolderToStore(@Config Configuration configuration, - @Alias("storeName") @DisplayName("Store Name") String storeName, - @ParameterGroup(name = "Documents") DocumentParameters documentParameters, - @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ - - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_STORE_METADATA,configuration.getVectorStore()); - - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); - - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); - - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - - EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder() - .documentSplitter(DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars())) - .embeddingModel(embeddingModel) - .embeddingStore(embeddingStore) - .build(); - - BaseStorage baseStorage = BaseStorage.builder() - .configuration(configuration) - .storageType(documentParameters.getStorageType()) - .contextPath(documentParameters.getContextPath()) - .fileType(documentParameters.getFileType()) - .build(); - - long documentNumber = 0; - while(baseStorage.hasNext()) { - - Document document = baseStorage.next(); - embeddingStoreIngestor.ingest(document); - documentNumber ++; - } - JSONObject jsonObject = JsonUtils.createFolderIngestionStatusObject(storeName, documentNumber, documentParameters.getFileType()); + @Throws(EmbeddingErrorTypeProvider.class) + + public org.mule.runtime.extension.api.runtime.operation.Result + addFolderToStore( @Config Configuration configuration, + @Alias("storeName") @DisplayName("Store Name") String storeName, + @ParameterGroup(name = "Documents") DocumentParameters documentParameters, + @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters){ - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + try { + + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_STORE_METADATA,configuration.getVectorStore()); + + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); + + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); + + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + + EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder() + .documentSplitter(DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars())) + .embeddingModel(embeddingModel) + .embeddingStore(embeddingStore) + .build(); + + BaseStorage baseStorage = BaseStorage.builder() + .configuration(configuration) + .storageType(documentParameters.getStorageType()) + .contextPath(documentParameters.getContextPath()) + .fileType(documentParameters.getFileType()) + .build(); + + long documentNumber = 0; + while(baseStorage.hasNext()) { + + Document document = baseStorage.next(); + embeddingStoreIngestor.ingest(document); + documentNumber ++; + } + JSONObject jsonObject = JsonUtils.createFolderIngestionStatusObject(storeName, documentNumber, documentParameters.getFileType()); + + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while adding folder %s into the store %s", documentParameters.getContextPath(), storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } /** @@ -167,50 +209,62 @@ public InputStream addFolderToStore(@Config Configuration configuration, */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-add-document-to-store") - public InputStream addFileEmbedding( - @Config Configuration configuration, - @Alias("storeName") @DisplayName("Store Name") String storeName, - @ParameterGroup(name = "Document") DocumentParameters documentParameters, - @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { - - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_STORE_METADATA,configuration.getVectorStore()); - - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); - - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); - - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - - EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder() - .documentSplitter(DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars())) - .embeddingModel(embeddingModel) - .embeddingStore(embeddingStore) - .build(); - - BaseStorage baseStorage = BaseStorage.builder() - .configuration(configuration) - .storageType(documentParameters.getStorageType()) - .contextPath(documentParameters.getContextPath()) - .fileType(documentParameters.getFileType()) - .build(); - Document document = baseStorage.getSingleDocument(); - - embeddingStoreIngestor.ingest(document); - - JSONObject jsonObject = JsonUtils.createFileIngestionStatusObject(storeName, documentParameters.getFileType(), documentParameters.getContextPath()); - - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + @Throws(EmbeddingErrorTypeProvider.class) + + public org.mule.runtime.extension.api.runtime.operation.Result + addFileEmbedding( @Config Configuration configuration, + @Alias("storeName") @DisplayName("Store Name") String storeName, + @ParameterGroup(name = "Document") DocumentParameters documentParameters, + @ParameterGroup(name = "Segmentation") SegmentationParameters segmentationParameters, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { + + try { + + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_STORE_METADATA,configuration.getVectorStore()); + + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); + + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); + + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + + EmbeddingStoreIngestor embeddingStoreIngestor = EmbeddingStoreIngestor.builder() + .documentSplitter(DocumentSplitters.recursive(segmentationParameters.getMaxSegmentSizeInChar(), segmentationParameters.getMaxOverlapSizeInChars())) + .embeddingModel(embeddingModel) + .embeddingStore(embeddingStore) + .build(); + + BaseStorage baseStorage = BaseStorage.builder() + .configuration(configuration) + .storageType(documentParameters.getStorageType()) + .contextPath(documentParameters.getContextPath()) + .fileType(documentParameters.getFileType()) + .build(); + Document document = baseStorage.getSingleDocument(); + + embeddingStoreIngestor.ingest(document); + + JSONObject jsonObject = JsonUtils.createFileIngestionStatusObject(storeName, documentParameters.getFileType(), documentParameters.getContextPath()); + + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while adding document %s into the store %s", documentParameters.getContextPath(), storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } /** @@ -218,78 +272,89 @@ public InputStream addFileEmbedding( */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-query-from-store") - public InputStream queryFromEmbedding( - @Config Configuration configuration, - @Alias("storeName") @DisplayName("Store Name") String storeName, - String question, - Number maxResults, - Double minScore, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { - int maximumResults = (int) maxResults; - if (minScore == null) { //|| minScore == 0) { - minScore = Constants.EMBEDDING_SEARCH_REQUEST_DEFAULT_MIN_SCORE; - } + @Throws(EmbeddingErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + queryFromEmbedding( @Config Configuration configuration, + @Alias("storeName") @DisplayName("Store Name") String storeName, + String question, + Number maxResults, + Double minScore, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { + try { - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); + int maximumResults = (int) maxResults; + if (minScore == null) { //|| minScore == 0) { + minScore = Constants.EMBEDDING_SEARCH_REQUEST_DEFAULT_MIN_SCORE; + } - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); - Embedding questionEmbedding = embeddingModel.embed(question).content(); + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder() - .queryEmbedding(questionEmbedding) - .maxResults(maximumResults) - .minScore(minScore) - .build(); + Embedding questionEmbedding = embeddingModel.embed(question).content(); - EmbeddingSearchResult searchResult = embeddingStore.search(searchRequest); - List> embeddingMatches = searchResult.matches(); + EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder() + .queryEmbedding(questionEmbedding) + .maxResults(maximumResults) + .minScore(minScore) + .build(); - String information = embeddingMatches.stream() - .map(match -> match.embedded().text()) - .collect(joining("\n\n")); + EmbeddingSearchResult searchResult = embeddingStore.search(searchRequest); + List> embeddingMatches = searchResult.matches(); - JSONObject jsonObject = new JSONObject(); - jsonObject.put("response", information); - jsonObject.put("storeName", storeName); - jsonObject.put("question", question); - JSONArray sources = new JSONArray(); + String information = embeddingMatches.stream() + .map(match -> match.embedded().text()) + .collect(joining("\n\n")); - JSONObject contentObject; - for (EmbeddingMatch match : embeddingMatches) { - Metadata matchMetadata = match.embedded().metadata(); + JSONObject jsonObject = new JSONObject(); + jsonObject.put("response", information); + jsonObject.put("storeName", storeName); + jsonObject.put("question", question); + JSONArray sources = new JSONArray(); - contentObject = new JSONObject(); - contentObject.put("embeddingId", match.embeddingId()); - contentObject.put("text", match.embedded().text()); - contentObject.put("score", match.score()); + JSONObject contentObject; + for (EmbeddingMatch match : embeddingMatches) { + Metadata matchMetadata = match.embedded().metadata(); - JSONObject metadataObject = new JSONObject(matchMetadata.toMap()); - contentObject.put("metadata", metadataObject); + contentObject = new JSONObject(); + contentObject.put("embeddingId", match.embeddingId()); + contentObject.put("text", match.embedded().text()); + contentObject.put("score", match.score()); - sources.put(contentObject); - } + JSONObject metadataObject = new JSONObject(matchMetadata.toMap()); + contentObject.put("metadata", metadataObject); + + sources.put(contentObject); + } + + jsonObject.put("sources", sources); + + jsonObject.put("maxResults", maxResults); + jsonObject.put("minScore", minScore); + jsonObject.put("question", question); + jsonObject.put("storeName", storeName); - jsonObject.put("sources", sources); - jsonObject.put("maxResults", maxResults); - jsonObject.put("minScore", minScore); - jsonObject.put("question", question); - jsonObject.put("storeName", storeName); - + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + } catch (Exception e) { + + throw new ModuleException( + String.format("Error while querying embeddings with question %s from the store %s", question, storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } @@ -299,94 +364,105 @@ public InputStream queryFromEmbedding( */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-query-from-store-with-filter") - public InputStream queryByFilterFromEmbedding( String storeName, - String question, - Number maxResults, - Double minScore, - @Config Configuration configuration, - @ParameterGroup(name = "Filter") MetadataFilterParameters.SearchFilterParameters searchFilterParams, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { - - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); - - int maximumResults = (int) maxResults; - if (minScore == null) { //|| minScore == 0) { - minScore = Constants.EMBEDDING_SEARCH_REQUEST_DEFAULT_MIN_SCORE; - } + @Throws(EmbeddingErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + queryByFilterFromEmbedding( String storeName, + String question, + Number maxResults, + Double minScore, + @Config Configuration configuration, + @ParameterGroup(name = "Filter") MetadataFilterParameters.SearchFilterParameters searchFilterParams, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); + try { - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); + int maximumResults = (int) maxResults; + if (minScore == null) { //|| minScore == 0) { + minScore = Constants.EMBEDDING_SEARCH_REQUEST_DEFAULT_MIN_SCORE; + } - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); - Embedding questionEmbedding = embeddingModel.embed(question).content(); + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - EmbeddingSearchRequest.EmbeddingSearchRequestBuilder searchRequestBuilder = EmbeddingSearchRequest.builder() - .queryEmbedding(questionEmbedding) - .maxResults(maximumResults) - .minScore(minScore); + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); - JSONObject jsonObject = new JSONObject(); + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - if(searchFilterParams.areFilterParamsSet()) { + Embedding questionEmbedding = embeddingModel.embed(question).content(); - Filter filter = searchFilterParams.buildMetadataFilter(); - searchRequestBuilder.filter(filter); - jsonObject.put("filter", searchFilterParams.getFilterJSONObject()); - } + EmbeddingSearchRequest.EmbeddingSearchRequestBuilder searchRequestBuilder = EmbeddingSearchRequest.builder() + .queryEmbedding(questionEmbedding) + .maxResults(maximumResults) + .minScore(minScore); - EmbeddingSearchRequest searchRequest = searchRequestBuilder.build(); + JSONObject jsonObject = new JSONObject(); - EmbeddingSearchResult searchResult = embeddingStore.search(searchRequest); - List> embeddingMatches = searchResult.matches(); + if(searchFilterParams.areFilterParamsSet()) { - String information = embeddingMatches.stream() - .map(match -> match.embedded().text()) - .collect(joining("\n\n")); + Filter filter = searchFilterParams.buildMetadataFilter(); + searchRequestBuilder.filter(filter); + jsonObject.put("filter", searchFilterParams.getFilterJSONObject()); + } - jsonObject.put("response", information); - jsonObject.put("storeName", storeName); - jsonObject.put("question", question); + EmbeddingSearchRequest searchRequest = searchRequestBuilder.build(); - JSONArray sources = new JSONArray(); + EmbeddingSearchResult searchResult = embeddingStore.search(searchRequest); + List> embeddingMatches = searchResult.matches(); - JSONObject contentObject; - String fullPath; - for (EmbeddingMatch match : embeddingMatches) { - Metadata matchMetadata = match.embedded().metadata(); + String information = embeddingMatches.stream() + .map(match -> match.embedded().text()) + .collect(joining("\n\n")); + jsonObject.put("response", information); + jsonObject.put("storeName", storeName); + jsonObject.put("question", question); - contentObject = new JSONObject(); - contentObject.put("embeddingId", match.embeddingId()); - contentObject.put("text", match.embedded().text()); - contentObject.put("score", match.score()); + JSONArray sources = new JSONArray(); - JSONObject metadataObject = new JSONObject(matchMetadata.toMap()); - contentObject.put("metadata", metadataObject); + JSONObject contentObject; + String fullPath; + for (EmbeddingMatch match : embeddingMatches) { + Metadata matchMetadata = match.embedded().metadata(); - sources.put(contentObject); - } - jsonObject.put("sources", sources); + contentObject = new JSONObject(); + contentObject.put("embeddingId", match.embeddingId()); + contentObject.put("text", match.embedded().text()); + contentObject.put("score", match.score()); + + JSONObject metadataObject = new JSONObject(matchMetadata.toMap()); + contentObject.put("metadata", metadataObject); + + sources.put(contentObject); + } + + jsonObject.put("sources", sources); + + jsonObject.put("maxResults", maxResults); + jsonObject.put("minScore", minScore); + jsonObject.put("question", question); + jsonObject.put("storeName", storeName); - jsonObject.put("maxResults", maxResults); - jsonObject.put("minScore", minScore); - jsonObject.put("question", question); - jsonObject.put("storeName", storeName); + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + } catch (Exception e) { - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + throw new ModuleException( + String.format("Error while querying embeddings with question %s from the store %s", question, storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } @@ -405,25 +481,37 @@ public InputStream queryByFilterFromEmbedding( String storeName, */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-list-sources") - public InputStream listSourcesFromStore( String storeName, - @Config Configuration configuration, - @ParameterGroup(name = "Querying Strategy") QueryParameters queryParams + @Throws(EmbeddingErrorTypeProvider.class) + public org.mule.runtime.extension.api.runtime.operation.Result + listSourcesFromStore( String storeName, + @Config Configuration configuration, + @ParameterGroup(name = "Querying Strategy") QueryParameters queryParams ) { - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_QUERY_ALL,configuration.getVectorStore()); - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); + try { + + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_QUERY_ALL,configuration.getVectorStore()); + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); + + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .queryParams(queryParams) + .build(); + + JSONObject jsonObject = baseStore.listSources(); - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .queryParams(queryParams) - .build(); + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); - JSONObject jsonObject = baseStore.listSources(); + } catch (Exception e) { - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + throw new ModuleException( + String.format("Error while listing sources from the store %s", storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } @@ -432,40 +520,51 @@ public InputStream listSourcesFromStore( String storeName, */ @MediaType(value = APPLICATION_JSON, strict = false) @Alias("EMBEDDING-remove-from-store-by-filter") - public InputStream removeEmbeddingsByFilter( String storeName, - @Config Configuration configuration, - @ParameterGroup(name = "Filter") MetadataFilterParameters.RemoveFilterParameters removeFilterParams, - @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { + @Throws(EmbeddingErrorTypeProvider.class) + + public org.mule.runtime.extension.api.runtime.operation.Result + removeEmbeddingsByFilter( String storeName, + @Config Configuration configuration, + @ParameterGroup(name = "Filter") MetadataFilterParameters.RemoveFilterParameters removeFilterParams, + @ParameterGroup(name = "Embedding Model") EmbeddingModelParameters embeddingModelParameters) { - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_REMOVE_EMBEDDINGS,configuration.getVectorStore()); - EmbeddingOperationValidator.validateOperationType( - Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); + try { + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_REMOVE_EMBEDDINGS,configuration.getVectorStore()); + EmbeddingOperationValidator.validateOperationType( + Constants.EMBEDDING_OPERATION_TYPE_FILTER_BY_METADATA,configuration.getVectorStore()); - BaseModel baseModel = BaseModel.builder() - .configuration(configuration) - .embeddingModelParameters(embeddingModelParameters) - .build(); + BaseModel baseModel = BaseModel.builder() + .configuration(configuration) + .embeddingModelParameters(embeddingModelParameters) + .build(); - EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); + EmbeddingModel embeddingModel = baseModel.buildEmbeddingModel(); - BaseStore baseStore = BaseStore.builder() - .storeName(storeName) - .configuration(configuration) - .dimension(embeddingModel.dimension()) - .build(); + BaseStore baseStore = BaseStore.builder() + .storeName(storeName) + .configuration(configuration) + .dimension(embeddingModel.dimension()) + .build(); - EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); + EmbeddingStore embeddingStore = baseStore.buildEmbeddingStore(); - Filter filter = removeFilterParams.buildMetadataFilter(); + Filter filter = removeFilterParams.buildMetadataFilter(); - embeddingStore.removeAll(filter); + embeddingStore.removeAll(filter); - JSONObject jsonObject = new JSONObject(); - jsonObject.put("storeName", storeName); - jsonObject.put("filter", removeFilterParams.getFilterJSONObject()); - jsonObject.put("status", Constants.OPERATION_STATUS_DELETED); + JSONObject jsonObject = new JSONObject(); + jsonObject.put("storeName", storeName); + jsonObject.put("filter", removeFilterParams.getFilterJSONObject()); + jsonObject.put("status", Constants.OPERATION_STATUS_DELETED); - return toInputStream(jsonObject.toString(), StandardCharsets.UTF_8); + return createEmbeddingResponse(jsonObject.toString(), new HashMap<>()); + + } catch (Exception e) { + throw new ModuleException( + String.format("Error while removing embeddings from the store %s", storeName), + MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE, + e); + } } }