Skip to content

Commit d6e073f

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

File tree

11 files changed

+1131
-0
lines changed

11 files changed

+1131
-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: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
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+
22+
import org.junit.jupiter.api.Test;
23+
24+
import org.testcontainers.elasticsearch.ElasticsearchContainer;
25+
import org.testcontainers.junit.jupiter.Container;
26+
import org.testcontainers.junit.jupiter.Testcontainers;
27+
28+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepository;
29+
import org.springframework.ai.chat.messages.AssistantMessage;
30+
import org.springframework.ai.chat.messages.MessageType;
31+
import org.springframework.ai.chat.messages.UserMessage;
32+
import org.springframework.boot.autoconfigure.AutoConfigurations;
33+
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
34+
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
35+
36+
import static org.assertj.core.api.Assertions.assertThat;
37+
38+
/**
39+
* Integration tests for {@link ElasticSearchChatMemoryRepositoryAutoConfiguration}.
40+
*
41+
* @author Fu Jian
42+
* @since 1.1.0
43+
*/
44+
@Testcontainers
45+
class ElasticSearchChatMemoryRepositoryAutoConfigurationIT {
46+
47+
@Container
48+
static ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer(
49+
"docker.elastic.co/elasticsearch/elasticsearch:8.10.2")
50+
.withEnv("xpack.security.enabled", "false")
51+
.withEnv("xpack.security.http.ssl.enabled", "false");
52+
53+
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
54+
.withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class,
55+
ElasticSearchChatMemoryRepositoryAutoConfiguration.class))
56+
.withPropertyValues("spring.elasticsearch.uris=http://" + elasticsearchContainer.getHost() + ":"
57+
+ elasticsearchContainer.getMappedPort(9200))
58+
.withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=autoconfig-test-chat-memory");
59+
60+
@Test
61+
void addAndGet() {
62+
this.contextRunner.run(context -> {
63+
ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class);
64+
65+
String conversationId = UUID.randomUUID().toString();
66+
assertThat(memory.findByConversationId(conversationId)).isEmpty();
67+
68+
memory.saveAll(conversationId, List.of(new UserMessage("test question")));
69+
70+
assertThat(memory.findByConversationId(conversationId)).hasSize(1);
71+
assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER);
72+
assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question");
73+
74+
memory.deleteByConversationId(conversationId);
75+
assertThat(memory.findByConversationId(conversationId)).isEmpty();
76+
77+
memory.saveAll(conversationId,
78+
List.of(new UserMessage("test question"), new AssistantMessage("test answer")));
79+
80+
assertThat(memory.findByConversationId(conversationId)).hasSize(2);
81+
assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER);
82+
assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question");
83+
assertThat(memory.findByConversationId(conversationId).get(1).getMessageType())
84+
.isEqualTo(MessageType.ASSISTANT);
85+
assertThat(memory.findByConversationId(conversationId).get(1).getText()).isEqualTo("test answer");
86+
});
87+
}
88+
89+
@Test
90+
void propertiesConfiguration() {
91+
this.contextRunner
92+
.withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=custom-testindex")
93+
.run(context -> {
94+
ElasticSearchChatMemoryRepositoryProperties properties = context
95+
.getBean(ElasticSearchChatMemoryRepositoryProperties.class);
96+
assertThat(properties.getIndexName()).isEqualTo("custom-testindex");
97+
});
98+
}
99+
100+
@Test
101+
void findConversationIds() {
102+
this.contextRunner.run(context -> {
103+
ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class);
104+
105+
String conversationId1 = UUID.randomUUID().toString();
106+
String conversationId2 = UUID.randomUUID().toString();
107+
108+
memory.saveAll(conversationId1, List.of(new UserMessage("test question 1")));
109+
memory.saveAll(conversationId2, List.of(new UserMessage("test question 2")));
110+
111+
List<String> conversationIds = memory.findConversationIds();
112+
assertThat(conversationIds).contains(conversationId1, conversationId2);
113+
});
114+
}
115+
116+
}

‎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/ElasticSearchChatMemoryRepositoryPropertiesTest.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
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.junit.jupiter.api.Test;
20+
21+
import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig;
22+
23+
import static org.assertj.core.api.Assertions.assertThat;
24+
25+
/**
26+
* Unit tests for {@link ElasticSearchChatMemoryRepositoryProperties}.
27+
*
28+
* @author Fu Jian
29+
* @since 1.1.0
30+
*/
31+
class ElasticSearchChatMemoryRepositoryPropertiesTest {
32+
33+
@Test
34+
void defaultValues() {
35+
var props = new ElasticSearchChatMemoryRepositoryProperties();
36+
assertThat(props.getIndexName()).isEqualTo(ElasticSearchChatMemoryRepositoryConfig.DEFAULT_INDEX_NAME);
37+
}
38+
39+
@Test
40+
void customValues() {
41+
var props = new ElasticSearchChatMemoryRepositoryProperties();
42+
props.setIndexName("custom_chat_memory");
43+
44+
assertThat(props.getIndexName()).isEqualTo("custom_chat_memory");
45+
}
46+
47+
}

0 commit comments

Comments
 (0)