Merge pull request #58 from MuleSoft-AI-Chain-Project/feature/enhance…

…d-module-configuration

Feature/enhanced module configuration

pom.xml

@@ -6,7 +6,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.mulesoft.connectors</groupId>
     <artifactId>mule4-vectors-connector</artifactId>
-    <version>0.2.6-SNAPSHOT</version>
+    <version>0.2.19-SNAPSHOT</version>
     <packaging>mule-extension</packaging>
     <name>MuleSoft Vectors Connector - Mule 4</name>
     <description>MuleSoft Vectors Connector provides access to a broad number of external Vector Stores.</description>

src/main/java/org/mule/extension/vectors/internal/config/Configuration.java

@@ -1,16 +1,19 @@
 package org.mule.extension.vectors.internal.config;
 
-import org.mule.extension.vectors.internal.helper.provider.EmbeddingModelServiceProvider;
-import org.mule.extension.vectors.internal.helper.provider.VectorStoreProvider;
+import org.mule.extension.vectors.internal.model.BaseModelConfiguration;
 import org.mule.extension.vectors.internal.operation.DocumentOperations;
 import org.mule.extension.vectors.internal.operation.EmbeddingOperations;
+import org.mule.extension.vectors.internal.storage.BaseStorageConfiguration;
+import org.mule.extension.vectors.internal.store.BaseStoreConfiguration;
+import org.mule.runtime.api.meta.ExpressionSupport;
 import org.mule.runtime.extension.api.annotation.Alias;
+import org.mule.runtime.extension.api.annotation.Expression;
 import org.mule.runtime.extension.api.annotation.Operations;
+import org.mule.runtime.extension.api.annotation.param.Optional;
 import org.mule.runtime.extension.api.annotation.param.Parameter;
 import org.mule.runtime.extension.api.annotation.param.display.DisplayName;
 import org.mule.runtime.extension.api.annotation.param.display.Placement;
 import org.mule.runtime.extension.api.annotation.param.display.Summary;
-import org.mule.runtime.extension.api.annotation.values.OfValues;
 
 /**
  * This class represents an extension configuration, values set in this class are commonly used across multiple
@@ -25,34 +28,35 @@ public class Configuration {
   @DisplayName("Embedding Model Service")
   @Summary("The embedding model service.")
   @Placement(order = 1, tab = Placement.DEFAULT_TAB)
-  @OfValues(EmbeddingModelServiceProvider.class)
-  private String embeddingModelService;
+  private BaseModelConfiguration modelConfiguration;
 
   @Parameter
   @Alias("vectorStore")
   @DisplayName("Vector Store")
   @Summary("The vector store.")
   @Placement(order = 2, tab = Placement.DEFAULT_TAB)
-  @OfValues(VectorStoreProvider.class)
-  private String vectorStore;
+  @Expression(ExpressionSupport.NOT_SUPPORTED)
+  private BaseStoreConfiguration storeConfiguration;
 
   @Parameter
-  @Alias("configFilePath")
-  @DisplayName("Configuration File Path")
-  @Summary("The configuration file path.")
+  @Alias("storageProvider")
+  @DisplayName("Storage provider")
+  @Summary("The storage provider.")
+  @Optional
   @Placement(order = 3, tab = Placement.DEFAULT_TAB)
-  private String configFilePath;
+  private BaseStorageConfiguration storageConfiguration;
 
-  public String getEmbeddingModelService() {
-    return embeddingModelService;
+  public BaseModelConfiguration getModelConfiguration() {
+    return modelConfiguration;
   }
 
-  public String getVectorStore() {
-    return vectorStore;
+  public BaseStoreConfiguration getStoreConfiguration() {
+    return storeConfiguration;
   }
 
-  public String getConfigFilePath() {
-    return configFilePath;
+  public BaseStorageConfiguration getStorageConfiguration() {
+    return storageConfiguration;
   }
 
+
 }

src/main/java/org/mule/extension/vectors/internal/constant/Constants.java

@@ -13,8 +13,10 @@ private Constants() {}
   public static final String FILE_TYPE_URL = "url";
 
   public static final String STORAGE_TYPE_LOCAL = "Local";
-  public static final String STORAGE_TYPE_S3 = "S3";
-  public static final String STORAGE_TYPE_AZURE_BLOB = "AZURE_BLOB";
+  public static final String STORAGE_TYPE_CLOUD = "Cloud";
+
+  public static final String STORAGE_PROVIDER_AWS_S3 = "S3";
+  public static final String STORAGE_PROVIDER_AZURE_BLOB = "AZURE_BLOB";
 
   public static final String EMBEDDING_MODEL_SERVICE_OPENAI = "OPENAI";
   public static final String EMBEDDING_MODEL_SERVICE_AZURE_OPENAI = "AZURE_OPENAI";

src/main/java/org/mule/extension/vectors/internal/error/MuleVectorsErrorType.java

@@ -7,5 +7,5 @@
  */
 public enum MuleVectorsErrorType implements ErrorTypeDefinition<MuleVectorsErrorType> {
 
-  DOCUMENT_OPERATIONS_FAILURE, EMBEDDING_OPERATIONS_FAILURE, AI_SERVICES_FAILURE, STORE_SERVICES_FAILURE
+  DOCUMENT_OPERATIONS_FAILURE, EMBEDDING_OPERATIONS_FAILURE, AI_SERVICES_FAILURE, STORE_SERVICES_FAILURE, STORAGE_SERVICES_FAILURE
 }

src/main/java/org/mule/extension/vectors/internal/error/provider/DocumentErrorTypeProvider.java

@@ -12,12 +12,13 @@
 import static java.util.Arrays.asList;
 import static java.util.Collections.unmodifiableSet;
 import static org.mule.extension.vectors.internal.error.MuleVectorsErrorType.DOCUMENT_OPERATIONS_FAILURE;
+import static org.mule.extension.vectors.internal.error.MuleVectorsErrorType.STORAGE_SERVICES_FAILURE;
 
 public class DocumentErrorTypeProvider implements ErrorTypeProvider {
 
   @SuppressWarnings("rawtypes")
   @Override
   public Set<ErrorTypeDefinition> getErrorTypes() {
-    return unmodifiableSet(new HashSet<>(asList(DOCUMENT_OPERATIONS_FAILURE)));
+    return unmodifiableSet(new HashSet<>(asList(DOCUMENT_OPERATIONS_FAILURE, STORAGE_SERVICES_FAILURE)));
   }
 }

src/main/java/org/mule/extension/vectors/internal/error/provider/EmbeddingErrorTypeProvider.java

@@ -11,8 +11,7 @@
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.unmodifiableSet;
-import static org.mule.extension.vectors.internal.error.MuleVectorsErrorType.AI_SERVICES_FAILURE;
-import static org.mule.extension.vectors.internal.error.MuleVectorsErrorType.EMBEDDING_OPERATIONS_FAILURE;
+import static org.mule.extension.vectors.internal.error.MuleVectorsErrorType.*;
 
 public class EmbeddingErrorTypeProvider implements ErrorTypeProvider {
 
@@ -21,6 +20,8 @@ public class EmbeddingErrorTypeProvider implements ErrorTypeProvider {
   public Set<ErrorTypeDefinition> getErrorTypes() {
     return unmodifiableSet(new HashSet<>(asList(
         EMBEDDING_OPERATIONS_FAILURE,
-        AI_SERVICES_FAILURE)));
+        AI_SERVICES_FAILURE,
+        STORE_SERVICES_FAILURE,
+        STORAGE_SERVICES_FAILURE)));
   }
 }

src/main/java/org/mule/extension/vectors/internal/extension/Connector.java

@@ -2,9 +2,29 @@
 
 import org.mule.extension.vectors.internal.config.Configuration;
 import org.mule.extension.vectors.internal.error.MuleVectorsErrorType;
+import org.mule.extension.vectors.internal.model.BaseModelConfiguration;
+import org.mule.extension.vectors.internal.model.azureopenai.AzureOpenAIModelConfiguration;
+import org.mule.extension.vectors.internal.model.einstein.EinsteinModelConfiguration;
+import org.mule.extension.vectors.internal.model.huggingface.HuggingFaceModelConfiguration;
+import org.mule.extension.vectors.internal.model.mistralai.MistralAIModelConfiguration;
+import org.mule.extension.vectors.internal.model.nomic.NomicModelConfiguration;
+import org.mule.extension.vectors.internal.model.openai.OpenAIModelConfiguration;
+import org.mule.extension.vectors.internal.storage.BaseStorageConfiguration;
+import org.mule.extension.vectors.internal.storage.azureblob.AzureBlobStorageConfiguration;
+import org.mule.extension.vectors.internal.storage.s3.AWSS3StorageConfiguration;
+import org.mule.extension.vectors.internal.store.BaseStoreConfiguration;
+import org.mule.extension.vectors.internal.store.aisearch.AISearchStoreConfiguration;
+import org.mule.extension.vectors.internal.store.chroma.ChromaStoreConfiguration;
+import org.mule.extension.vectors.internal.store.elasticsearch.ElasticsearchStoreConfiguration;
+import org.mule.extension.vectors.internal.store.milvus.MilvusStoreConfiguration;
+import org.mule.extension.vectors.internal.store.opensearch.OpenSearchStoreConfiguration;
+import org.mule.extension.vectors.internal.store.pgvector.PGVectorStoreConfiguration;
+import org.mule.extension.vectors.internal.store.pinecone.PineconeStoreConfiguration;
+import org.mule.extension.vectors.internal.store.qdrant.QdrantStoreConfiguration;
 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.SubTypeMapping;
 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;
@@ -20,6 +40,28 @@
 @Xml(prefix = "ms-vectors")
 @Extension(name = "MuleSoft Vectors Connector", category = Category.SELECT)
 @Configurations(Configuration.class)
+@SubTypeMapping(baseType = BaseStoreConfiguration.class,
+    subTypes = {
+        AISearchStoreConfiguration.class,
+        ChromaStoreConfiguration.class,
+        ElasticsearchStoreConfiguration.class,
+        MilvusStoreConfiguration.class,
+        OpenSearchStoreConfiguration.class,
+        PGVectorStoreConfiguration.class,
+        PineconeStoreConfiguration.class,
+        QdrantStoreConfiguration.class})
+@SubTypeMapping(baseType = BaseModelConfiguration.class,
+    subTypes = {
+        AzureOpenAIModelConfiguration.class,
+        EinsteinModelConfiguration.class,
+        HuggingFaceModelConfiguration.class,
+        MistralAIModelConfiguration.class,
+        NomicModelConfiguration.class,
+        OpenAIModelConfiguration.class})
+@SubTypeMapping(baseType = BaseStorageConfiguration.class,
+    subTypes = {
+        AWSS3StorageConfiguration.class,
+        AzureBlobStorageConfiguration.class})
 @RequiresEnterpriseLicense(allowEvaluationLicense = true)
 @ErrorTypes(MuleVectorsErrorType.class)
 @JavaVersionSupport({JAVA_8, JAVA_11, JAVA_17})

src/main/java/org/mule/extension/vectors/internal/helper/provider/EmbeddingModelNameProvider.java

@@ -50,7 +50,7 @@ public class EmbeddingModelNameProvider implements ValueProvider {
   @Override
   public Set<Value> resolve() throws ValueResolvingException {
 
-    String embeddingModelService = configuration.getEmbeddingModelService();
+    String embeddingModelService = configuration.getModelConfiguration().getEmbeddingModelService();
     switch (embeddingModelService) {
       case Constants.EMBEDDING_MODEL_SERVICE_OPENAI:
         return VALUES_FOR_OPENAI;

src/main/java/org/mule/extension/vectors/internal/helper/provider/StorageTypeProvider.java

@@ -16,8 +16,7 @@ public Set<Value> resolve() throws ValueResolvingException {
 
     return ValueBuilder.getValuesFor(
         Constants.STORAGE_TYPE_LOCAL,
-        Constants.STORAGE_TYPE_S3,
-        Constants.STORAGE_TYPE_AZURE_BLOB);
+        Constants.STORAGE_TYPE_CLOUD);
   }
 
 }

src/main/java/org/mule/extension/vectors/internal/model/BaseModel.java

@@ -61,8 +61,8 @@ public BaseModel build() {
 
       BaseModel baseModel;
 
-      LOGGER.debug("Embedding Model Service: " + configuration.getEmbeddingModelService());
-      switch (configuration.getEmbeddingModelService()) {
+      LOGGER.debug("Embedding Model Service: " + configuration.getModelConfiguration().getEmbeddingModelService());
+      switch (configuration.getModelConfiguration().getEmbeddingModelService()) {
 
         case Constants.EMBEDDING_MODEL_SERVICE_AZURE_OPENAI:
           baseModel = new AzureOpenAIModel(configuration, embeddingModelParameters);
@@ -90,7 +90,7 @@ public BaseModel build() {
 
         default:
           throw new ModuleException(
-              String.format("Error while initializing embedding model service. \"%s\" is not supported.", configuration.getEmbeddingModelService()),
+              String.format("Error while initializing embedding model service. \"%s\" is not supported.", configuration.getModelConfiguration().getEmbeddingModelService()),
               MuleVectorsErrorType.AI_SERVICES_FAILURE);
       }
       return baseModel;

src/main/java/org/mule/extension/vectors/internal/model/BaseModelConfiguration.java

@@ -0,0 +1,6 @@
+package org.mule.extension.vectors.internal.model;
+
+public interface BaseModelConfiguration {
+
+  String getEmbeddingModelService();
+}

src/main/java/org/mule/extension/vectors/internal/model/azureopenai/AzureOpenAIModel.java

@@ -12,17 +12,16 @@
 
 public class AzureOpenAIModel extends BaseModel {
 
-  private final String apiKey;
   private final String endpoint;
+  private final String apiKey;
 
   public AzureOpenAIModel(Configuration configuration, EmbeddingModelParameters embeddingModelParameters) {
 
     super(configuration,embeddingModelParameters);
-    JSONObject config = readConfigFile(configuration.getConfigFilePath());
-    assert config != null;
-    JSONObject modelConfig = config.getJSONObject(Constants.EMBEDDING_MODEL_SERVICE_AZURE_OPENAI);
-    this.apiKey = modelConfig.getString("AZURE_OPENAI_KEY");
-    this.endpoint = modelConfig.getString("AZURE_OPENAI_ENDPOINT");
+
+    AzureOpenAIModelConfiguration azureOpenAIModelConfiguration = (AzureOpenAIModelConfiguration) configuration.getModelConfiguration();
+    this.endpoint = azureOpenAIModelConfiguration.getEndpoint();
+    this.apiKey = azureOpenAIModelConfiguration.getApiKey();
   }
 
   public EmbeddingModel buildEmbeddingModel() {

src/main/java/org/mule/extension/vectors/internal/model/azureopenai/AzureOpenAIModelConfiguration.java

@@ -0,0 +1,41 @@
+package org.mule.extension.vectors.internal.model.azureopenai;
+
+import org.mule.extension.vectors.internal.constant.Constants;
+import org.mule.extension.vectors.internal.model.BaseModelConfiguration;
+import org.mule.extension.vectors.internal.store.BaseStoreConfiguration;
+import org.mule.runtime.api.meta.ExpressionSupport;
+import org.mule.runtime.extension.api.annotation.Alias;
+import org.mule.runtime.extension.api.annotation.Expression;
+import org.mule.runtime.extension.api.annotation.param.Parameter;
+import org.mule.runtime.extension.api.annotation.param.display.DisplayName;
+import org.mule.runtime.extension.api.annotation.param.display.Password;
+import org.mule.runtime.extension.api.annotation.param.display.Placement;
+
+@Alias("azureOpenAI")
+@DisplayName("Azure OpenAI")
+public class AzureOpenAIModelConfiguration implements BaseModelConfiguration {
+
+  @Parameter
+  @Expression(ExpressionSupport.SUPPORTED)
+  @Placement(order = 1)
+  private String endpoint;
+
+  @Parameter
+  @Password
+  @Expression(ExpressionSupport.SUPPORTED)
+  @Placement(order = 2)
+  private String apiKey;
+
+  @Override
+  public String getEmbeddingModelService() {
+    return Constants.EMBEDDING_MODEL_SERVICE_AZURE_OPENAI;
+  }
+
+  public String getEndpoint() {
+    return endpoint;
+  }
+
+  public String getApiKey() {
+    return apiKey;
+  }
+}