pass OpenSearchClient directly to OpenSearchEmbeddingStore (langchain…

…4j#309) see langchain4j#308
hwy110 · Nov 25, 2023 · 11e7364 · 11e7364
1 parent f98c040
commit 11e7364
Showing 1 changed file with 23 additions and 2 deletions.
diff --git a/...ch/src/main/java/dev/langchain4j/store/embedding/opensearch/OpenSearchEmbeddingStore.java b/...ch/src/main/java/dev/langchain4j/store/embedding/opensearch/OpenSearchEmbeddingStore.java
@@ -70,7 +70,7 @@ public class OpenSearchEmbeddingStore implements EmbeddingStore<TextSegment> {
      * @param apiKey    OpenSearch API key (optional)
      * @param userName  OpenSearch username (optional)
      * @param password  OpenSearch password (optional)
-     * @param indexName OpenSearch index name (optional). Default value: "default"
+     * @param indexName OpenSearch index name.
      */
     public OpenSearchEmbeddingStore(String serverUrl,
                                     String apiKey,
@@ -121,7 +121,7 @@ public OpenSearchEmbeddingStore(String serverUrl,
      * @param serviceName The AWS signing service name, one of `es` (Amazon OpenSearch) or `aoss` (Amazon OpenSearch Serverless).
      * @param region      The AWS region for which requests will be signed. This should typically match the region in `serverUrl`.
      * @param options     The options to establish connection with the service. It must include which credentials should be used.
-     * @param indexName   OpenSearch index name (optional). Default value: "default"
+     * @param indexName   OpenSearch index name.
      */
     public OpenSearchEmbeddingStore(String serverUrl,
                                     String serviceName,
@@ -138,6 +138,19 @@ public OpenSearchEmbeddingStore(String serverUrl,
         this.indexName = ensureNotNull(indexName, "indexName");
     }
 
+    /**
+     * Creates an instance of OpenSearchEmbeddingStore using provided OpenSearchClient
+     *
+     * @param openSearchClient OpenSearch client provided
+     * @param indexName        OpenSearch index name.
+     */
+    public OpenSearchEmbeddingStore(OpenSearchClient openSearchClient,
+                                    String indexName) {
+
+        this.client = ensureNotNull(openSearchClient, "openSearchClient");
+        this.indexName = ensureNotNull(indexName, "indexName");
+    }
+
     public static Builder builder() {
         return new Builder();
     }
@@ -152,6 +165,7 @@ public static class Builder {
         private String region;
         private AwsSdk2TransportOptions options;
         private String indexName = "default";
+        private OpenSearchClient openSearchClient;
 
         public Builder serverUrl(String serverUrl) {
             this.serverUrl = serverUrl;
@@ -193,7 +207,14 @@ public Builder indexName(String indexName) {
             return this;
         }
 
+        public Builder openSearchClient(OpenSearchClient openSearchClient) {
+            this.openSearchClient = openSearchClient;
+            return this;
+        }
         public OpenSearchEmbeddingStore build() {
+            if (openSearchClient != null) {
+                return new OpenSearchEmbeddingStore(openSearchClient, indexName);
+            }
             if (!isNullOrBlank(serviceName) && !isNullOrBlank(region) && options != null) {
                 return new OpenSearchEmbeddingStore(serverUrl, serviceName, region, options, indexName);
             }