From a684b2f6be9cb60d6b4f4e5ffb20b2e8b51bbda6 Mon Sep 17 00:00:00 2001
From: Anthony Kugel <32197641+aokugel@users.noreply.github.com>
Date: Fri, 22 Nov 2024 10:14:09 -0600
Subject: [PATCH] added elasticsearch vector db (#7)

* added elasticsearch vector db

* fixed checkstyle errors
---
 pom.xml                                       |   5 +
 .../ContentRetrieverClientFactory.java        |   6 +-
 .../ElasticsearchContentRetrieverClient.java  |  93 ++++++++++++
 .../model/enums/ContentRetrieverType.java     |   3 +-
 .../ElasticsearchConnectionEntity.java        | 137 +++++++++++++++++
 .../retriever/BaseRetrieverRequest.java       |   3 +-
 .../retriever/ElasticsearchRequest.java       | 142 ++++++++++++++++++
 .../composer/util/mappers/MapperUtil.java     |   6 +
 .../mappers/RetrieverConnectionMapper.java    |  22 ++-
 src/main/resources/application.properties     |   8 +
 10 files changed, 420 insertions(+), 5 deletions(-)
 create mode 100644 src/main/java/com/redhat/composer/config/retriever/contentretriever/ElasticsearchContentRetrieverClient.java
 create mode 100644 src/main/java/com/redhat/composer/model/mongo/contentretrieverentites/ElasticsearchConnectionEntity.java
 create mode 100644 src/main/java/com/redhat/composer/model/request/retriever/ElasticsearchRequest.java

diff --git a/pom.xml b/pom.xml
index ef2db2a..21962a2 100644
--- a/pom.xml
+++ b/pom.xml
@@ -67,6 +67,11 @@
             <artifactId>langchain4j-weaviate</artifactId>
             <version>${langchain4j.version}</version>
         </dependency>
+        <dependency>
+            <groupId>dev.langchain4j</groupId>
+            <artifactId>langchain4j-elasticsearch</artifactId>
+            <version>${langchain4j.version}</version>
+        </dependency>
         <!-- <dependency>
             <groupId>io.quarkiverse.langchain4j</groupId>
             <artifactId>quarkus-langchain4j-openai</artifactId>
diff --git a/src/main/java/com/redhat/composer/config/retriever/contentretriever/ContentRetrieverClientFactory.java b/src/main/java/com/redhat/composer/config/retriever/contentretriever/ContentRetrieverClientFactory.java
index b324411..36659db 100644
--- a/src/main/java/com/redhat/composer/config/retriever/contentretriever/ContentRetrieverClientFactory.java
+++ b/src/main/java/com/redhat/composer/config/retriever/contentretriever/ContentRetrieverClientFactory.java
@@ -16,8 +16,10 @@ public class ContentRetrieverClientFactory {
   
   @Inject
   Neo4jContentRetrieverClient neo4jContentRetrieverClient;
-  
 
+  @Inject
+  ElasticsearchContentRetrieverClient elasticsearchContentRetrieverClient;
+  
   static final ContentRetrieverType DEFAULT_CONTENT_RETRIEVER = ContentRetrieverType.WEAVIATE;
 
   /**
@@ -35,6 +37,8 @@ public BaseContentRetrieverClient getContentRetrieverClient(ContentRetrieverType
         return weaviateEmbeddingStoreClient;
       case ContentRetrieverType.NEO4J:
         return neo4jContentRetrieverClient;
+      case ContentRetrieverType.ELASTICSEARCH:
+        return elasticsearchContentRetrieverClient;
       default:
         throw new RuntimeException("Content Retriever type not found: " + contentRetrieverType);
     }
diff --git a/src/main/java/com/redhat/composer/config/retriever/contentretriever/ElasticsearchContentRetrieverClient.java b/src/main/java/com/redhat/composer/config/retriever/contentretriever/ElasticsearchContentRetrieverClient.java
new file mode 100644
index 0000000..d76f98a
--- /dev/null
+++ b/src/main/java/com/redhat/composer/config/retriever/contentretriever/ElasticsearchContentRetrieverClient.java
@@ -0,0 +1,93 @@
+package com.redhat.composer.config.retriever.contentretriever;
+
+import org.eclipse.microprofile.config.inject.ConfigProperty;
+import org.jboss.logging.Logger;
+
+//import com.redhat.composer.config.retriever.contentretriever.custom.WeaviateEmbeddingStoreCustom;
+import com.redhat.composer.model.request.RetrieverRequest;
+//import com.redhat.composer.model.request.retriever.WeaviateRequest;
+import com.redhat.composer.model.request.retriever.ElasticsearchRequest;
+
+import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.model.embedding.EmbeddingModel;
+import dev.langchain4j.rag.content.retriever.ContentRetriever;
+import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
+import dev.langchain4j.store.embedding.EmbeddingStore;
+import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
+import org.apache.http.impl.client.BasicCredentialsProvider;
+import org.apache.http.message.BasicHeader;
+import org.apache.http.auth.UsernamePasswordCredentials;
+import org.apache.http.client.CredentialsProvider;
+import org.apache.http.auth.AuthScope;
+import org.apache.http.Header;
+import org.apache.http.HttpHost;
+import org.elasticsearch.client.RestClient;
+import jakarta.inject.Singleton;
+
+/**
+ * Elasticsearch Content Retriever Client.
+ */
+@Singleton
+public class ElasticsearchContentRetrieverClient extends BaseContentRetrieverClient {
+
+  Logger log = Logger.getLogger(ElasticsearchContentRetrieverClient.class);
+
+  @ConfigProperty(name = "elasticsearch.default.host")
+  private String elasticHost;
+
+  @ConfigProperty(name = "elasticsearch.default.user")
+  private String elasticUser;
+
+  @ConfigProperty(name = "elasticsearch.default.password")
+  private String elasticPassword;
+
+  @ConfigProperty(name = "elasticsearch.default.index")
+  private String elasticIndex;
+
+  /**
+   * Get the Content Retriever.
+   * @param request the RetrieverRequest
+   * @return the Content Retriever
+   */
+  public ContentRetriever getContentRetriever(RetrieverRequest request) {
+    ElasticsearchRequest elasticsearchRequest = (ElasticsearchRequest) request.getBaseRetrieverRequest();
+    if (elasticsearchRequest == null) {
+      elasticsearchRequest = new ElasticsearchRequest();
+    }
+    String host = elasticsearchRequest.getHost() != null ? elasticsearchRequest.getHost() : elasticHost;
+    String user = elasticsearchRequest.getUser() != null ? elasticsearchRequest.getUser() : elasticUser;
+    String pass = elasticsearchRequest.getPassword() != null ? elasticsearchRequest.getPassword() : elasticPassword;
+    String index = elasticsearchRequest.getIndex() != null ? elasticsearchRequest.getIndex() : elasticIndex;
+
+    // TODO: Make this configurable for different authentication types
+    final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
+    credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(user, pass));
+
+    RestClient restClient = RestClient
+        .builder(HttpHost.create(host))
+        .setHttpClientConfigCallback(httpClientBuilder -> {
+          httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
+          return httpClientBuilder;
+        })
+        .build();
+  
+    log.debug("Attempting to connect to Elasticsearch at " + host + " with index " + index);
+
+    EmbeddingStore<TextSegment> store = ElasticsearchEmbeddingStore.builder()
+                    .indexName(index)
+                    .restClient(restClient)
+                    .build();
+      
+    // Retrieve the embedding model
+    EmbeddingModel embeddingModel = getEmbeddingModel(request.getEmbeddingType());
+
+    // Create the content retriever
+    ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
+          .embeddingStore(store)
+          .embeddingModel(embeddingModel)
+        .build();
+      
+    return contentRetriever;
+  }
+
+}
diff --git a/src/main/java/com/redhat/composer/model/enums/ContentRetrieverType.java b/src/main/java/com/redhat/composer/model/enums/ContentRetrieverType.java
index a9e0a25..b17317c 100644
--- a/src/main/java/com/redhat/composer/model/enums/ContentRetrieverType.java
+++ b/src/main/java/com/redhat/composer/model/enums/ContentRetrieverType.java
@@ -6,7 +6,8 @@
 public enum ContentRetrieverType {
 
   WEAVIATE("weaviate"),
-  NEO4J("neo4j");
+  NEO4J("neo4j"),
+  ELASTICSEARCH("elasticsearch");
 
   private final String type;
 
diff --git a/src/main/java/com/redhat/composer/model/mongo/contentretrieverentites/ElasticsearchConnectionEntity.java b/src/main/java/com/redhat/composer/model/mongo/contentretrieverentites/ElasticsearchConnectionEntity.java
new file mode 100644
index 0000000..b78f6a3
--- /dev/null
+++ b/src/main/java/com/redhat/composer/model/mongo/contentretrieverentites/ElasticsearchConnectionEntity.java
@@ -0,0 +1,137 @@
+package com.redhat.composer.model.mongo.contentretrieverentites;
+
+import org.apache.commons.lang3.builder.EqualsBuilder;
+import org.bson.codecs.pojo.annotations.BsonDiscriminator;
+
+import com.redhat.composer.model.enums.ContentRetrieverType;
+
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Elasticsearch Connection Entity.
+ */
+@SuppressWarnings("all")
+@BsonDiscriminator("Elasticsearch")
+public class ElasticsearchConnectionEntity extends BaseRetrieverConnectionEntity {
+
+  // Key of the value containing the text used for retrieval and passed into the LLM
+  String textKey;
+
+  // List of metadata fields to be retrieved as part of the content
+  List<String> metadataFields = List.of("source");
+
+  String index;
+
+  String host;
+
+  String user;
+
+  String password;
+
+  public ElasticsearchConnectionEntity() {
+    contentRetrieverType = ContentRetrieverType.ELASTICSEARCH;
+  }
+
+  public String getTextKey() {
+    return this.textKey;
+  }
+
+  public void setTextKey(String textKey) {
+    this.textKey = textKey;
+  }
+
+  public List<String> getMetadataFields() {
+    return this.metadataFields;
+  }
+
+  public void setMetadataFields(List<String> metadataFields) {
+    this.metadataFields = metadataFields;
+  }
+
+  public String getIndex() {
+    return this.index;
+  }
+
+  public void setIndex(String index) {
+    this.index = index;
+  }
+
+  public String getHost() {
+    return this.host;
+  }
+
+  public void setHost(String host) {
+    this.host = host;
+  }
+
+  public String getUser() {
+    return this.user;
+  }
+
+  public String getPassword() {
+    return this.password;
+  }
+
+  public void setUser(String user) {
+    this.user = user;
+  }
+
+  public void setPassword(String password) {
+    this.password = password;
+  }
+
+  public ElasticsearchConnectionEntity textKey(String textKey) {
+    setTextKey(textKey);
+    return this;
+  }
+
+  public ElasticsearchConnectionEntity metadataFields(List<String> metadataFields) {
+    setMetadataFields(metadataFields);
+    return this;
+  }
+
+  public ElasticsearchConnectionEntity index(String index) {
+    setIndex(index);
+    return this;
+  }
+
+  public ElasticsearchConnectionEntity host(String host) {
+    setHost(host);
+    return this;
+  }
+
+  public ElasticsearchConnectionEntity user(String user) {
+    setUser(user);
+    return this;
+  }
+
+  public ElasticsearchConnectionEntity password(String password) {
+    setPassword(password);
+    return this;
+  }
+
+  @Override
+    public boolean equals(Object o) {
+    return EqualsBuilder.reflectionEquals(this, o);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(textKey, metadataFields, index, host, user, password);
+  }
+
+  @Override
+  public String toString() {
+    return "{" +
+      " textKey='" + getTextKey() + "'" +
+      ", metadataFields='" + getMetadataFields() + "'" +
+      ", index='" + getIndex() + "'" +
+      ", host='" + getHost() + "'" +
+      ", apiKey='" + getUser() + "'" +
+      ", apiKey='" + getPassword() + "'" +
+      "}";
+  }
+  
+  
+}
diff --git a/src/main/java/com/redhat/composer/model/request/retriever/BaseRetrieverRequest.java b/src/main/java/com/redhat/composer/model/request/retriever/BaseRetrieverRequest.java
index 891af7a..e7b80fc 100644
--- a/src/main/java/com/redhat/composer/model/request/retriever/BaseRetrieverRequest.java
+++ b/src/main/java/com/redhat/composer/model/request/retriever/BaseRetrieverRequest.java
@@ -9,7 +9,8 @@
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.EXISTING_PROPERTY, property = "contentRetrieverType", visible = true)
 @JsonSubTypes({
     @JsonSubTypes.Type(value = WeaviateRequest.class, name = "weaviate"),
-    @JsonSubTypes.Type(value = Neo4JRequest.class, name = "neo4j")
+    @JsonSubTypes.Type(value = Neo4JRequest.class, name = "neo4j"),
+    @JsonSubTypes.Type(value = ElasticsearchRequest.class, name = "elasticsearch")
 })
 public class BaseRetrieverRequest {
 
diff --git a/src/main/java/com/redhat/composer/model/request/retriever/ElasticsearchRequest.java b/src/main/java/com/redhat/composer/model/request/retriever/ElasticsearchRequest.java
new file mode 100644
index 0000000..cba3704
--- /dev/null
+++ b/src/main/java/com/redhat/composer/model/request/retriever/ElasticsearchRequest.java
@@ -0,0 +1,142 @@
+package com.redhat.composer.model.request.retriever;
+
+import java.util.List;
+import java.util.Objects;
+
+import org.apache.commons.lang3.builder.EqualsBuilder;
+
+import com.redhat.composer.model.enums.ContentRetrieverType;
+
+@SuppressWarnings("all")
+public class ElasticsearchRequest extends BaseRetrieverRequest {
+
+  // Key of the value containing the text used for retrieval and passed into the LLM
+  String textKey;
+
+  // List of metadata fields to be retrieved as part of the content
+  List<String> metadataFields = List.of("source","title");
+
+  String index;
+
+  String host;
+
+  String user;
+
+  String password;
+  
+
+  public ElasticsearchRequest() {
+    contentRetrieverType = ContentRetrieverType.ELASTICSEARCH.getType();
+  }
+
+  public ElasticsearchRequest(String textKey, List<String> metadataFields, String index, String host, String user, String password) {
+    this.textKey = textKey;
+    this.metadataFields = metadataFields;
+    this.index = index;
+    this.host = host;
+    this.user = user;
+    this.password = password;
+  }
+
+  public String getTextKey() {
+    return this.textKey;
+  }
+
+  public void setTextKey(String textKey) {
+    this.textKey = textKey;
+  }
+
+  public List<String> getMetadataFields() {
+    return this.metadataFields;
+  }
+
+  public void setMetadataFields(List<String> metadataFields) {
+    this.metadataFields = metadataFields;
+  }
+
+  public String getIndex() {
+    return this.index;
+  }
+
+  public void setIndex(String index) {
+    this.index = index;
+  }
+
+  public String getHost() {
+    return this.host;
+  }
+
+  public void setHost(String host) {
+    this.host = host;
+  }
+
+  public String getUser() {
+    return this.user;
+  }
+
+  public String getPassword() {
+    return this.password;
+  }
+
+  public void setUser(String user) {
+    this.user = user;
+  }
+
+  public void setPassword(String password) {
+    this.password = password;
+  }
+
+  public ElasticsearchRequest textKey(String textKey) {
+    setTextKey(textKey);
+    return this;
+  }
+
+  public ElasticsearchRequest metadataFields(List<String> metadataFields) {
+    setMetadataFields(metadataFields);
+    return this;
+  }
+
+  public ElasticsearchRequest index(String index) {
+    setIndex(index);
+    return this;
+  }
+
+  public ElasticsearchRequest host(String host) {
+    setHost(host);
+    return this;
+  }
+
+  public ElasticsearchRequest user(String user) {
+    setUser(user);
+    return this;
+  }
+
+  public ElasticsearchRequest password(String password) {
+    password(password);
+    return this;
+  }
+  @Override
+    public boolean equals(Object o) {
+      return EqualsBuilder.reflectionEquals(this, o);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(textKey, metadataFields, index, host, user, password);
+  }
+
+  @Override
+  public String toString() {
+    return "{" +
+      " textKey='" + getTextKey() + "'" +
+      ", metadataFields='" + getMetadataFields() + "'" +
+      ", index='" + getIndex() + "'" +
+      ", host='" + getHost() + "'" +
+      ", user='" + getUser() + "'" +
+      ", password='" + getPassword() + "'" +
+      "}";
+  }
+
+  
+  
+}
diff --git a/src/main/java/com/redhat/composer/util/mappers/MapperUtil.java b/src/main/java/com/redhat/composer/util/mappers/MapperUtil.java
index 85c08cf..ef6cc5b 100644
--- a/src/main/java/com/redhat/composer/util/mappers/MapperUtil.java
+++ b/src/main/java/com/redhat/composer/util/mappers/MapperUtil.java
@@ -9,6 +9,7 @@
 import com.redhat.composer.model.mongo.LlmConnectionEntity;
 import com.redhat.composer.model.mongo.RetrieverConnectionEntity;
 import com.redhat.composer.model.mongo.contentretrieverentites.BaseRetrieverConnectionEntity;
+import com.redhat.composer.model.mongo.contentretrieverentites.ElasticsearchConnectionEntity;
 import com.redhat.composer.model.mongo.contentretrieverentites.Neo4jEntity;
 import com.redhat.composer.model.mongo.contentretrieverentites.WeaviateConnectionEntity;
 import com.redhat.composer.model.request.LLMRequest;
@@ -16,6 +17,7 @@
 import com.redhat.composer.model.request.retriever.BaseRetrieverRequest;
 import com.redhat.composer.model.request.retriever.Neo4JRequest;
 import com.redhat.composer.model.request.retriever.WeaviateRequest;
+import com.redhat.composer.model.request.retriever.ElasticsearchRequest;
 
 import jakarta.enterprise.inject.Default;
 
@@ -77,6 +79,8 @@ default BaseRetrieverConnectionEntity mapToBaseEntity(BaseRetrieverRequest reque
         return retrieverConnectionMapper.toEntity((WeaviateRequest) request);
       case ContentRetrieverType.NEO4J:
         return retrieverConnectionMapper.toEntity((Neo4JRequest) request);
+      case ContentRetrieverType.ELASTICSEARCH:
+        return retrieverConnectionMapper.toEntity((ElasticsearchRequest) request);
       default:
         return null;
     }
@@ -99,6 +103,8 @@ default BaseRetrieverRequest mapToBaseRequest(BaseRetrieverConnectionEntity enti
         return retrieverConnectionMapper.toRequest((WeaviateConnectionEntity) entity);
       case ContentRetrieverType.NEO4J:
         return retrieverConnectionMapper.toRequest((Neo4jEntity) entity);
+      case ContentRetrieverType.ELASTICSEARCH:
+        return retrieverConnectionMapper.toRequest((ElasticsearchConnectionEntity) entity);
       default:
         return null;
     }
diff --git a/src/main/java/com/redhat/composer/util/mappers/RetrieverConnectionMapper.java b/src/main/java/com/redhat/composer/util/mappers/RetrieverConnectionMapper.java
index 66ec816..70f9388 100644
--- a/src/main/java/com/redhat/composer/util/mappers/RetrieverConnectionMapper.java
+++ b/src/main/java/com/redhat/composer/util/mappers/RetrieverConnectionMapper.java
@@ -5,8 +5,10 @@
 import com.redhat.composer.model.enums.ContentRetrieverType;
 import com.redhat.composer.model.mongo.contentretrieverentites.Neo4jEntity;
 import com.redhat.composer.model.mongo.contentretrieverentites.WeaviateConnectionEntity;
+import com.redhat.composer.model.mongo.contentretrieverentites.ElasticsearchConnectionEntity;
 import com.redhat.composer.model.request.retriever.Neo4JRequest;
 import com.redhat.composer.model.request.retriever.WeaviateRequest;
+import com.redhat.composer.model.request.retriever.ElasticsearchRequest;
 
 /**
  * RetrieverConnectionMapper interface.
@@ -31,9 +33,17 @@ public interface RetrieverConnectionMapper {
   Neo4jEntity toEntity(Neo4JRequest request);
 
   /**
-   * Maps a WeaviateConnectionEntity to a WeaviateRequest.
+   * Maps a ElasticsearchConnectionEntity to a ElasticsearchRequest.
+
+   * @param request the ElasticsearchConnectionEntity to map
+   * @return the WeaviateRequest
+   */
+  ElasticsearchConnectionEntity toEntity(ElasticsearchRequest request);
+
+  /**
+   * Maps a WeaviateEntity to a WeaviateRequest.
 
-   * @param entity the WeaviateConnectionEntity to map
+   * @param entity the WeaviateEntity to map
    * @return the WeaviateRequest
    */
   WeaviateRequest toRequest(WeaviateConnectionEntity entity);
@@ -46,6 +56,14 @@ public interface RetrieverConnectionMapper {
    */
   Neo4JRequest toRequest(Neo4jEntity entity);
 
+  /**
+   * Maps a ElasticsearchEntity to a ElasticsearchRequest.
+
+   * @param entity the ElasticsearchEntity to map
+   * @return the ElasticsearchRequest
+   */
+  ElasticsearchRequest toRequest(ElasticsearchConnectionEntity entity);
+
   /**
    * Maps a ContentRetrieverType to a String.
 
diff --git a/src/main/resources/application.properties b/src/main/resources/application.properties
index 9ea219d..bb62b7a 100644
--- a/src/main/resources/application.properties
+++ b/src/main/resources/application.properties
@@ -17,6 +17,14 @@ neo4j.default.url=bolt://localhost:7687
 neo4j.default.username=neo4j
 neo4j.default.password=<REPLACE_ME>
 
+############################################################
+# Default Elasticsearch
+############################################################
+elasticsearch.default.host=https://localhost
+elasticsearch.default.index=test_index
+elasticsearch.default.user=elastic
+elasticsearch.default.password=<REPLACE_ME>
+
 ############################################################
 # Default Mistral Model
 ############################################################