Skip to content

Commit ec2576c

Browse files
jiafu1115jiafu1115
committed
Support chat memory based on Elasticsearch
Signed-off-by: stroller <[email protected]>
1 parent 940bcf3 commit ec2576c

File tree

11 files changed

+1143
-0
lines changed

11 files changed

+1143
-0
lines changed

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/pom.xml‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!--
3+
~ Copyright 2023-2025 the original author or authors.
4+
~
5+
~ Licensed under the Apache License, Version 2.0 (the "License");
6+
~ you may not use this file except in compliance with the License.
7+
~ You may obtain a copy of the License at
8+
~
9+
~ https://www.apache.org/licenses/LICENSE-2.0
10+
~
11+
~ Unless required by applicable law or agreed to in writing, software
12+
~ distributed under the License is distributed on an "AS IS" BASIS,
13+
~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
~ See the License for the specific language governing permissions and
15+
~ limitations under the License.
16+
-->
17+
18+
<project xmlns="http://maven.apache.org/POM/4.0.0"
19+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
20+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
21+
<modelVersion>4.0.0</modelVersion>
22+
<parent>
23+
<groupId>org.springframework.ai</groupId>
24+
<artifactId>spring-ai-parent</artifactId>
25+
<version>1.1.0-SNAPSHOT</version>
26+
<relativePath>../../../../../../pom.xml</relativePath>
27+
</parent>
28+
29+
<artifactId>spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch</artifactId>
30+
<name>Spring AI Auto Configuration - Chat Memory Repository - Elasticsearch</name>
31+
<description>Spring AI Auto Configuration for Elasticsearch Chat Memory Repository</description>
32+
33+
<url>https://github.com/spring-projects/spring-ai</url>
34+
35+
<scm>
36+
<url>https://github.com/spring-projects/spring-ai</url>
37+
<connection>git://github.com/spring-projects/spring-ai.git</connection>
38+
<developerConnection>[email protected]:spring-projects/spring-ai.git</developerConnection>
39+
</scm>
40+
41+
<dependencies>
42+
<dependency>
43+
<groupId>org.springframework.ai</groupId>
44+
<artifactId>spring-ai-model-chat-memory-repository-elasticsearch</artifactId>
45+
<version>${project.version}</version>
46+
</dependency>
47+
48+
<dependency>
49+
<groupId>org.springframework.ai</groupId>
50+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
51+
<version>${project.parent.version}</version>
52+
</dependency>
53+
54+
<dependency>
55+
<groupId>org.springframework.boot</groupId>
56+
<artifactId>spring-boot-autoconfigure</artifactId>
57+
</dependency>
58+
59+
<dependency>
60+
<groupId>org.springframework.boot</groupId>
61+
<artifactId>spring-boot-configuration-processor</artifactId>
62+
<optional>true</optional>
63+
</dependency>
64+
65+
<dependency>
66+
<groupId>org.springframework.boot</groupId>
67+
<artifactId>spring-boot-autoconfigure-processor</artifactId>
68+
<optional>true</optional>
69+
</dependency>
70+
71+
<dependency>
72+
<groupId>co.elastic.clients</groupId>
73+
<artifactId>elasticsearch-java</artifactId>
74+
<optional>true</optional>
75+
</dependency>
76+
77+
<!-- Test dependencies -->
78+
<dependency>
79+
<groupId>org.springframework.boot</groupId>
80+
<artifactId>spring-boot-starter-test</artifactId>
81+
<scope>test</scope>
82+
</dependency>
83+
84+
<dependency>
85+
<groupId>org.springframework.ai</groupId>
86+
<artifactId>spring-ai-test</artifactId>
87+
<version>${project.version}</version>
88+
<scope>test</scope>
89+
</dependency>
90+
91+
<dependency>
92+
<groupId>org.springframework.boot</groupId>
93+
<artifactId>spring-boot-testcontainers</artifactId>
94+
<scope>test</scope>
95+
</dependency>
96+
97+
<dependency>
98+
<groupId>org.testcontainers</groupId>
99+
<artifactId>elasticsearch</artifactId>
100+
<scope>test</scope>
101+
</dependency>
102+
103+
<dependency>
104+
<groupId>org.testcontainers</groupId>
105+
<artifactId>junit-jupiter</artifactId>
106+
<scope>test</scope>
107+
</dependency>
108+
</dependencies>
109+
</project>
110+

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/model/chat/memory/repository/elasticsearch/autoconfigure/ElasticSearchChatMemoryRepositoryAutoConfiguration.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure;
18+
19+
import co.elastic.clients.elasticsearch.ElasticsearchClient;
20+
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
21+
import co.elastic.clients.transport.rest_client.RestClientTransport;
22+
import com.fasterxml.jackson.databind.DeserializationFeature;
23+
import com.fasterxml.jackson.databind.ObjectMapper;
24+
import org.elasticsearch.client.RestClient;
25+
26+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepository;
27+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig;
28+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
29+
import org.springframework.boot.autoconfigure.AutoConfiguration;
30+
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
31+
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
32+
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
33+
import org.springframework.boot.context.properties.EnableConfigurationProperties;
34+
import org.springframework.context.annotation.Bean;
35+
36+
/**
37+
* {@link AutoConfiguration Auto-configuration} for
38+
* {@link ElasticSearchChatMemoryRepository}.
39+
*
40+
* @author Fu Jian
41+
* @since 1.1.0
42+
*/
43+
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
44+
@ConditionalOnClass({ ElasticSearchChatMemoryRepository.class, RestClient.class })
45+
@EnableConfigurationProperties(ElasticSearchChatMemoryRepositoryProperties.class)
46+
public class ElasticSearchChatMemoryRepositoryAutoConfiguration {
47+
48+
@Bean
49+
@ConditionalOnMissingBean
50+
public ElasticSearchChatMemoryRepositoryConfig elasticSearchChatMemoryRepositoryConfig(
51+
ElasticSearchChatMemoryRepositoryProperties properties, RestClient restClient) {
52+
ElasticsearchClient elasticsearchClient = new ElasticsearchClient(new RestClientTransport(restClient,
53+
new JacksonJsonpMapper(
54+
new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false))))
55+
.withTransportOptions(t -> t.addHeader("user-agent", "spring-ai-chat-memory elastic-java"));
56+
return ElasticSearchChatMemoryRepositoryConfig.builder()
57+
.withClient(elasticsearchClient)
58+
.withIndexName(properties.getIndexName())
59+
.build();
60+
}
61+
62+
@Bean
63+
@ConditionalOnMissingBean
64+
public ElasticSearchChatMemoryRepository elasticSearchChatMemoryRepository(
65+
ElasticSearchChatMemoryRepositoryConfig config) {
66+
return ElasticSearchChatMemoryRepository.create(config);
67+
}
68+
69+
}

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/model/chat/memory/repository/elasticsearch/autoconfigure/ElasticSearchChatMemoryRepositoryProperties.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure;
18+
19+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig;
20+
import org.springframework.boot.context.properties.ConfigurationProperties;
21+
22+
/**
23+
* Configuration properties for Elasticsearch chat memory.
24+
*
25+
* @author Fu Jian
26+
* @since 1.1.0
27+
*/
28+
@ConfigurationProperties(ElasticSearchChatMemoryRepositoryProperties.CONFIG_PREFIX)
29+
public class ElasticSearchChatMemoryRepositoryProperties {
30+
31+
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.elasticsearch";
32+
33+
private String indexName = ElasticSearchChatMemoryRepositoryConfig.DEFAULT_INDEX_NAME;
34+
35+
public String getIndexName() {
36+
return this.indexName;
37+
}
38+
39+
public void setIndexName(String indexName) {
40+
this.indexName = indexName;
41+
}
42+
43+
}

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure.ElasticSearchChatMemoryRepositoryAutoConfiguration
2+

‎auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/model/chat/memory/repository/elasticsearch/autoconfigure/ElasticSearchChatMemoryRepositoryAutoConfigurationIT.java‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure;
18+
19+
import java.util.List;
20+
import java.util.UUID;
21+
import java.util.concurrent.TimeUnit;
22+
23+
import org.junit.jupiter.api.Test;
24+
import org.springframework.ai.chat.messages.Message;
25+
import org.testcontainers.elasticsearch.ElasticsearchContainer;
26+
import org.testcontainers.junit.jupiter.Container;
27+
import org.testcontainers.junit.jupiter.Testcontainers;
28+
29+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepository;
30+
import org.springframework.ai.chat.messages.AssistantMessage;
31+
import org.springframework.ai.chat.messages.MessageType;
32+
import org.springframework.ai.chat.messages.UserMessage;
33+
import org.springframework.boot.autoconfigure.AutoConfigurations;
34+
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
35+
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
36+
37+
import static org.assertj.core.api.Assertions.assertThat;
38+
39+
/**
40+
* Integration tests for {@link ElasticSearchChatMemoryRepositoryAutoConfiguration}.
41+
*
42+
* @author Fu Jian
43+
* @since 1.1.0
44+
*/
45+
@Testcontainers
46+
class ElasticSearchChatMemoryRepositoryAutoConfigurationIT {
47+
48+
@Container
49+
static ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer(
50+
"docker.elastic.co/elasticsearch/elasticsearch:8.10.2")
51+
.withEnv("xpack.security.enabled", "false")
52+
.withEnv("xpack.security.http.ssl.enabled", "false");
53+
54+
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
55+
.withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class,
56+
ElasticSearchChatMemoryRepositoryAutoConfiguration.class))
57+
.withPropertyValues("spring.elasticsearch.uris=http://" + elasticsearchContainer.getHost() + ":"
58+
+ elasticsearchContainer.getMappedPort(9200))
59+
.withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=autoconfig-test-chat-memory");
60+
61+
@Test
62+
void addAndGet() {
63+
this.contextRunner.run(context -> {
64+
ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class);
65+
66+
String conversationId = UUID.randomUUID().toString();
67+
assertThat(memory.findByConversationId(conversationId)).isEmpty();
68+
69+
memory.saveAll(conversationId, List.of(new UserMessage("test question")));
70+
71+
sleepForSearchable();
72+
73+
List<Message> conversation = memory.findByConversationId(conversationId);
74+
assertThat(conversation).hasSize(1);
75+
assertThat(conversation.get(0).getMessageType()).isEqualTo(MessageType.USER);
76+
assertThat(conversation.get(0).getText()).isEqualTo("test question");
77+
78+
memory.deleteByConversationId(conversationId);
79+
sleepForSearchable();
80+
assertThat(memory.findByConversationId(conversationId)).isEmpty();
81+
82+
memory.saveAll(conversationId,
83+
List.of(new UserMessage("test question"), new AssistantMessage("test answer")));
84+
sleepForSearchable();
85+
86+
conversation = memory.findByConversationId(conversationId);
87+
assertThat(conversation).hasSize(2);
88+
assertThat(conversation.get(0).getMessageType()).isEqualTo(MessageType.USER);
89+
assertThat(conversation.get(0).getText()).isEqualTo("test question");
90+
assertThat(conversation.get(1).getMessageType()).isEqualTo(MessageType.ASSISTANT);
91+
assertThat(conversation.get(1).getText()).isEqualTo("test answer");
92+
});
93+
}
94+
95+
private static void sleepForSearchable() throws InterruptedException {
96+
TimeUnit.SECONDS.sleep(2);
97+
}
98+
99+
@Test
100+
void propertiesConfiguration() {
101+
this.contextRunner
102+
.withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=custom-testindex")
103+
.run(context -> {
104+
ElasticSearchChatMemoryRepositoryProperties properties = context
105+
.getBean(ElasticSearchChatMemoryRepositoryProperties.class);
106+
assertThat(properties.getIndexName()).isEqualTo("custom-testindex");
107+
});
108+
}
109+
110+
@Test
111+
void findConversationIds() {
112+
this.contextRunner.run(context -> {
113+
ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class);
114+
115+
String conversationId1 = UUID.randomUUID().toString();
116+
String conversationId2 = UUID.randomUUID().toString();
117+
118+
memory.saveAll(conversationId1, List.of(new UserMessage("test question 1")));
119+
memory.saveAll(conversationId2, List.of(new UserMessage("test question 2")));
120+
121+
sleepForSearchable();
122+
123+
List<String> conversationIds = memory.findConversationIds();
124+
assertThat(conversationIds).contains(conversationId1, conversationId2);
125+
});
126+
}
127+
128+
}

0 commit comments

Comments
 (0)