Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Elasticsearch spring boot starter #27

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 55 additions & 0 deletions langchain4j-spring-boot-tests/pom.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-spring</artifactId>
<version>0.32.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

<artifactId>langchain4j-spring-boot-tests</artifactId>
<name>LangChain4j Spring Boot Tests</name>
<packaging>jar</packaging>

<dependencies>

<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-core</artifactId>
</dependency>

<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-embeddings-all-minilm-l6-v2-q</artifactId>
<scope>test</scope>
</dependency>

</dependencies>

<build>
<plugins>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-jar-plugin</artifactId>
<version>3.3.0</version>
<executions>
<execution>
<goals>
<goal>test-jar</goal>
</goals>
</execution>
</executions>
</plugin>
</plugins>
</build>

</project>
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
package dev.langchain4j.store.embedding.spring;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.AllMiniLmL6V2QuantizedEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.junit.jupiter.api.Test;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;

import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.data.Percentage.withPercentage;

public abstract class EmbeddingStoreAutoConfigurationIT {

protected abstract Class<?> autoConfigurationClass();

protected abstract Class<? extends EmbeddingStore<TextSegment>> embeddingStoreClass();

protected abstract String[] properties();

/**
* Property Key to configure {@link EmbeddingStore} dimension (if needed)
*
* @return {@link EmbeddingStore} dimension property key
*/
protected abstract String dimensionPropertyKey();

ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(autoConfigurationClass()));

@Test
void should_provide_embedding_store_without_embedding_model() {
// copy dimension property
EmbeddingModel embeddingModel = new AllMiniLmL6V2QuantizedEmbeddingModel();
String[] properties = new String[properties().length + 1];
System.arraycopy(properties(), 0, properties, 0, properties().length);
properties[properties.length - 1] = dimensionPropertyKey() + "=" + embeddingModel.dimension();

contextRunner
.withPropertyValues(properties)
.run(context -> {
TextSegment segment = TextSegment.from("hello");
Embedding embedding = embeddingModel.embed(segment.text()).content();

assertThat(context.getBean(embeddingStoreClass())).isExactlyInstanceOf(embeddingStoreClass());
EmbeddingStore<TextSegment> embeddingStore = context.getBean(embeddingStoreClass());

String id = embeddingStore.add(embedding, segment);
assertThat(id).isNotBlank();

awaitUntilPersisted();

List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(embedding, 10);
assertThat(relevant).hasSize(1);

EmbeddingMatch<TextSegment> match = relevant.get(0);
assertThat(match.score()).isCloseTo(1, withPercentage(1));
assertThat(match.embeddingId()).isEqualTo(id);
assertThat(match.embedding()).isEqualTo(embedding);
assertThat(match.embedded()).isEqualTo(segment);
});
}

@Test
void should_provide_embedding_store_with_embedding_model() {
contextRunner
.withBean(AllMiniLmL6V2QuantizedEmbeddingModel.class)
.withPropertyValues(properties())
.run(context -> {
TextSegment segment = TextSegment.from("hello");
EmbeddingModel embeddingModel = context.getBean(AllMiniLmL6V2QuantizedEmbeddingModel.class);
Embedding embedding = embeddingModel.embed(segment.text()).content();

assertThat(context.getBean(embeddingStoreClass())).isExactlyInstanceOf(embeddingStoreClass());
EmbeddingStore<TextSegment> embeddingStore = context.getBean(embeddingStoreClass());

String id = embeddingStore.add(embedding, segment);
assertThat(id).isNotBlank();

awaitUntilPersisted();

List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.findRelevant(embedding, 10);
assertThat(relevant).hasSize(1);

EmbeddingMatch<TextSegment> match = relevant.get(0);
assertThat(match.score()).isCloseTo(1, withPercentage(1));
assertThat(match.embeddingId()).isEqualTo(id);
assertThat(match.embedding()).isEqualTo(embedding);
assertThat(match.embedded()).isEqualTo(segment);
});
}

protected void awaitUntilPersisted() {

}
}
89 changes: 89 additions & 0 deletions langchian4j-elasticsearch-spring-boot-starter/pom.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-spring</artifactId>
<version>0.32.0-SNAPSHOT</version>
<relativePath>../pom.xml</relativePath>
</parent>

<artifactId>langchian4j-elasticsearch-spring-boot-starter</artifactId>
<name>LangChain4j Spring Boot starter for Elasticsearch</name>
<packaging>jar</packaging>

<dependencies>

<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-elasticsearch</artifactId>
</dependency>

<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>

<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-autoconfigure-processor</artifactId>
<optional>true</optional>
</dependency>

<!-- should be listed before spring-boot-configuration-processor -->
<dependency>
<groupId>org.projectlombok</groupId>
<artifactId>lombok</artifactId>
<scope>provided</scope>
</dependency>

<!-- needed to generate automatic metadata about available config properties -->
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-configuration-processor</artifactId>
<optional>true</optional>
</dependency>

<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-embeddings-all-minilm-l6-v2-q</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-spring-boot-tests</artifactId>
<version>${project.version}</version>
<classifier>tests</classifier>
<type>test-jar</type>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>elasticsearch</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.tinylog</groupId>
<artifactId>tinylog-impl</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.tinylog</groupId>
<artifactId>slf4j-tinylog</artifactId>
<scope>test</scope>
</dependency>
</dependencies>

</project>
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
package dev.langchain4j.store.embedding.elasticsearch.spring;

import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.lang.Nullable;

import java.util.Optional;

import static dev.langchain4j.store.embedding.elasticsearch.spring.ElasticsearchEmbeddingStoreProperties.*;

@AutoConfiguration
@EnableConfigurationProperties(ElasticsearchEmbeddingStoreProperties.class)
@ConditionalOnProperty(prefix = PREFIX, name = "enabled", havingValue = "true", matchIfMissing = true)
langchain4j marked this conversation as resolved.
Show resolved Hide resolved
public class ElasticsearchEmbeddingStoreAutoConfiguration {

@Bean
@ConditionalOnMissingBean
public ElasticsearchEmbeddingStore elasticsearchEmbeddingStore(ElasticsearchEmbeddingStoreProperties properties,
@Nullable EmbeddingModel embeddingModel) {
String serverUrl = Optional.ofNullable(properties.getServerUrl()).orElse(DEFAULT_SERVER_URL);
String indexName = Optional.ofNullable(properties.getIndexName()).orElse(DEFAULT_INDEX_NAME);
Integer dimension = Optional.ofNullable(properties.getDimension()).orElseGet(() -> embeddingModel == null ? null : embeddingModel.dimension());

return ElasticsearchEmbeddingStore.builder()
.serverUrl(serverUrl)
.apiKey(properties.getApiKey())
.userName(properties.getUserName())
.password(properties.getPassword())
.indexName(indexName)
.dimension(dimension)
.build();
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
package dev.langchain4j.store.embedding.elasticsearch.spring;

import lombok.Getter;
import lombok.Setter;
import org.springframework.boot.context.properties.ConfigurationProperties;

@Getter
@Setter
@ConfigurationProperties(prefix = ElasticsearchEmbeddingStoreProperties.PREFIX)
public class ElasticsearchEmbeddingStoreProperties {

static final String PREFIX = "langchain4j.elasticsearch";
static final String DEFAULT_SERVER_URL = "http://localhost:9200";
static final String DEFAULT_INDEX_NAME = "langchain4j-index";

private String serverUrl;
private String apiKey;
private String userName;
private String password;
private String indexName;
private Integer dimension;
}
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
dev.langchain4j.store.embedding.elasticsearch.spring.ElasticsearchEmbeddingStoreAutoConfiguration
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
package dev.langchain4j.store.embedding.elasticsearch.spring;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import dev.langchain4j.store.embedding.spring.EmbeddingStoreAutoConfigurationIT;
import lombok.SneakyThrows;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.elasticsearch.ElasticsearchContainer;

import static dev.langchain4j.internal.Utils.randomUUID;

class ElasticsearchEmbeddingStoreAutoConfigurationIT extends EmbeddingStoreAutoConfigurationIT {

static ElasticsearchContainer elasticsearch = new ElasticsearchContainer("elasticsearch:8.9.0")
.withEnv("xpack.security.enabled", "false")
.waitingFor(Wait.defaultWaitStrategy());

@BeforeAll
static void beforeAll() {
elasticsearch.start();
}

@AfterAll
static void afterAll() {
elasticsearch.stop();
}

@Override
protected Class<?> autoConfigurationClass() {
return ElasticsearchEmbeddingStoreAutoConfiguration.class;
}

@Override
protected Class<? extends EmbeddingStore<TextSegment>> embeddingStoreClass() {
return ElasticsearchEmbeddingStore.class;
}

@Override
protected String[] properties() {
return new String[]{
"langchain4j.elasticsearch.serverUrl=" + elasticsearch.getHttpHostAddress(),
"langchain4j.elasticsearch.indexName=" + randomUUID()
};
}

@Override
protected String dimensionPropertyKey() {
return "langchain4j.elasticsearch.dimension";
}

@Override
@SneakyThrows
protected void awaitUntilPersisted() {
Thread.sleep(1000);
}
}
Loading
Loading