Merge pull request #26 from MuleSoft-AI-Chain-Project/opensearch

opensearch support added
MuleSoft-AI-Chain-Project · Nov 11, 2024 · 24b97c5 · 24b97c5
2 parents b09f4f7 + 50aa4ba
commit 24b97c5
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 3 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/pom.xml b/pom.xml
@@ -173,6 +173,12 @@
             <version>${langchain4jVersion}</version>
         </dependency>
 
+		<dependency>
+			<groupId>dev.langchain4j</groupId>
+			<artifactId>langchain4j-opensearch</artifactId>
+			<version>${langchain4jVersion}</version>
+		</dependency>
+
         <dependency>
             <groupId>dev.langchain4j</groupId>
             <artifactId>langchain4j-pgvector</artifactId>

diff --git a/src/.DS_Store b/src/.DS_Store
diff --git a/src/main/java/org/mule/extension/mulechain/vectors/internal/constant/Constants.java b/src/main/java/org/mule/extension/mulechain/vectors/internal/constant/Constants.java
@@ -32,6 +32,8 @@ private Constants() {}
 
   public static final String VECTOR_STORE_PGVECTOR = "PGVECTOR";
   public static final String VECTOR_STORE_ELASTICSEARCH = "ELASTICSEARCH";
+
+  public static final String VECTOR_STORE_OPENSEARCH = "OPENSEARCH";
   public static final String VECTOR_STORE_MILVUS = "MILVUS";
   public static final String VECTOR_STORE_CHROMA = "CHROMA";
   public static final String VECTOR_STORE_PINECONE = "PINECONE";

diff --git a/...ava/org/mule/extension/mulechain/vectors/internal/helper/EmbeddingOperationValidator.java b/...ava/org/mule/extension/mulechain/vectors/internal/helper/EmbeddingOperationValidator.java
@@ -43,6 +43,7 @@ public class EmbeddingOperationValidator {
             new HashSet<>(Arrays.asList(
               Constants.VECTOR_STORE_PGVECTOR,
               Constants.VECTOR_STORE_ELASTICSEARCH,
+                    Constants.VECTOR_STORE_OPENSEARCH,
               Constants.VECTOR_STORE_MILVUS,
               Constants.VECTOR_STORE_CHROMA,
               Constants.VECTOR_STORE_PINECONE,
@@ -55,6 +56,7 @@ public class EmbeddingOperationValidator {
             new HashSet<>(Arrays.asList(
               Constants.VECTOR_STORE_PGVECTOR,
               Constants.VECTOR_STORE_ELASTICSEARCH,
+                    Constants.VECTOR_STORE_OPENSEARCH,
               Constants.VECTOR_STORE_MILVUS,
               Constants.VECTOR_STORE_CHROMA,
               Constants.VECTOR_STORE_PINECONE,
@@ -65,6 +67,7 @@ public class EmbeddingOperationValidator {
             new HashSet<>(Arrays.asList(
               Constants.VECTOR_STORE_PGVECTOR,
               Constants.VECTOR_STORE_ELASTICSEARCH,
+                    Constants.VECTOR_STORE_OPENSEARCH,
               Constants.VECTOR_STORE_MILVUS,
               Constants.VECTOR_STORE_CHROMA,
               Constants.VECTOR_STORE_PINECONE,

diff --git a/...a/org/mule/extension/mulechain/vectors/internal/helper/factory/EmbeddingStoreFactory.java b/...a/org/mule/extension/mulechain/vectors/internal/helper/factory/EmbeddingStoreFactory.java
@@ -5,6 +5,7 @@
 import dev.langchain4j.store.embedding.azure.search.AzureAiSearchEmbeddingStore;
 import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
 import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
+import dev.langchain4j.store.embedding.opensearch.OpenSearchEmbeddingStore;
 import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
 import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
 import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
@@ -65,6 +66,15 @@ public static EmbeddingStore<TextSegment> createStore(Configuration configuratio
                 password = vectorType.getString("ELASTICSEARCH_PASSWORD");
                 store = createElasticStore(vectorUrl, userName, password, indexName, dimension);
                 break;
+
+            case Constants.VECTOR_STORE_OPENSEARCH:
+                vectorType = config.getJSONObject(Constants.VECTOR_STORE_OPENSEARCH);
+                vectorUrl = vectorType.getString("OPENSEARCH_URL");
+                userName = vectorType.getString("OPENSEARCH_USER");
+                password = vectorType.getString("OPENSEARCH_PASSWORD");
+                store = createOpenSrchStore(vectorUrl, userName, password, indexName, dimension);
+                break;
+
             case Constants.VECTOR_STORE_PGVECTOR:
                 vectorType = config.getJSONObject(Constants.VECTOR_STORE_PGVECTOR);
                 vectorHost = vectorType.getString("POSTGRES_HOST");
@@ -125,6 +135,15 @@ private static EmbeddingStore<TextSegment> createElasticStore(String baseUrl, St
                 .build();
     }
 
+    private static EmbeddingStore<TextSegment> createOpenSrchStore(String baseUrl, String userName, String password, String collectionName, Integer dimension) {
+        return OpenSearchEmbeddingStore.builder()
+                .serverUrl(baseUrl)
+                .userName(userName)
+                .password(password)
+                .indexName(collectionName)
+                .build();
+    }
+
     private static EmbeddingStore<TextSegment> createPineconeStore(String apiKey, String cloudProvider, String cloudRegion, String collectionName, Integer dimension) {
         return PineconeEmbeddingStore.builder()
                 .apiKey(apiKey)

diff --git a/...va/org/mule/extension/mulechain/vectors/internal/helper/provider/VectorStoreProvider.java b/...va/org/mule/extension/mulechain/vectors/internal/helper/provider/VectorStoreProvider.java
@@ -1,13 +1,13 @@
 package org.mule.extension.mulechain.vectors.internal.helper.provider;
 
-import java.util.Set;
-
 import org.mule.extension.mulechain.vectors.internal.constant.Constants;
 import org.mule.runtime.api.value.Value;
 import org.mule.runtime.extension.api.values.ValueBuilder;
 import org.mule.runtime.extension.api.values.ValueProvider;
 import org.mule.runtime.extension.api.values.ValueResolvingException;
 
+import java.util.Set;
+
 public class VectorStoreProvider implements ValueProvider {
 
   @Override
@@ -16,11 +16,13 @@ public Set<Value> resolve() throws ValueResolvingException {
      return ValueBuilder.getValuesFor(
             Constants.VECTOR_STORE_PGVECTOR,
             Constants.VECTOR_STORE_ELASTICSEARCH,
+             Constants.VECTOR_STORE_OPENSEARCH,
             Constants.VECTOR_STORE_MILVUS,
             Constants.VECTOR_STORE_CHROMA,
             Constants.VECTOR_STORE_PINECONE,
             Constants.VECTOR_STORE_WEAVIATE,
-            Constants.VECTOR_STORE_AI_SEARCH
+            Constants.VECTOR_STORE_AI_SEARCH,
+             Constants.VECTOR_STORE_OPENSEARCH
     ); // MuleChainVectorsConstants.VECTOR_STORE_NEO4J
   }