diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/pom.xml b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/pom.xml new file mode 100644 index 00000000000..5905237372b --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch/pom.xml @@ -0,0 +1,110 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 1.1.0-SNAPSHOT + ../../../../../../pom.xml + + + spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch + Spring AI Auto Configuration - Chat Memory Repository - Elasticsearch + Spring AI Auto Configuration for Elasticsearch Chat Memory Repository + + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + org.springframework.ai + spring-ai-model-chat-memory-repository-elasticsearch + ${project.version} + + + + org.springframework.ai + spring-ai-autoconfigure-model-chat-memory + ${project.parent.version} + + + + org.springframework.boot + spring-boot-autoconfigure + + + + org.springframework.boot + spring-boot-configuration-processor + true + + + + org.springframework.boot + spring-boot-autoconfigure-processor + true + + + + co.elastic.clients + elasticsearch-java + true + + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.springframework.ai + spring-ai-test + ${project.version} + test + + + + org.springframework.boot + spring-boot-testcontainers + test + + + + org.testcontainers + elasticsearch + test + + + + org.testcontainers + junit-jupiter + test + + + + diff --git a/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 b/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 new file mode 100644 index 00000000000..528c1a77dd5 --- /dev/null +++ b/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 @@ -0,0 +1,69 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure; + +import co.elastic.clients.elasticsearch.ElasticsearchClient; +import co.elastic.clients.json.jackson.JacksonJsonpMapper; +import co.elastic.clients.transport.rest_client.RestClientTransport; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.elasticsearch.client.RestClient; + +import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepository; +import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig; +import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * {@link AutoConfiguration Auto-configuration} for + * {@link ElasticSearchChatMemoryRepository}. + * + * @author Fu Jian + * @since 1.1.0 + */ +@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class) +@ConditionalOnClass({ ElasticSearchChatMemoryRepository.class, RestClient.class }) +@EnableConfigurationProperties(ElasticSearchChatMemoryRepositoryProperties.class) +public class ElasticSearchChatMemoryRepositoryAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + public ElasticSearchChatMemoryRepositoryConfig elasticSearchChatMemoryRepositoryConfig( + ElasticSearchChatMemoryRepositoryProperties properties, RestClient restClient) { + ElasticsearchClient elasticsearchClient = new ElasticsearchClient(new RestClientTransport(restClient, + new JacksonJsonpMapper( + new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)))) + .withTransportOptions(t -> t.addHeader("user-agent", "spring-ai-chat-memory elastic-java")); + return ElasticSearchChatMemoryRepositoryConfig.builder() + .withClient(elasticsearchClient) + .withIndexName(properties.getIndexName()) + .build(); + } + + @Bean + @ConditionalOnMissingBean + public ElasticSearchChatMemoryRepository elasticSearchChatMemoryRepository( + ElasticSearchChatMemoryRepositoryConfig config) { + return ElasticSearchChatMemoryRepository.create(config); + } + +} diff --git a/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 b/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 new file mode 100644 index 00000000000..1d8d6b58717 --- /dev/null +++ b/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 @@ -0,0 +1,43 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure; + +import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * Configuration properties for Elasticsearch chat memory. + * + * @author Fu Jian + * @since 1.1.0 + */ +@ConfigurationProperties(ElasticSearchChatMemoryRepositoryProperties.CONFIG_PREFIX) +public class ElasticSearchChatMemoryRepositoryProperties { + + public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.elasticsearch"; + + private String indexName = ElasticSearchChatMemoryRepositoryConfig.DEFAULT_INDEX_NAME; + + public String getIndexName() { + return this.indexName; + } + + public void setIndexName(String indexName) { + this.indexName = indexName; + } + +} diff --git a/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 b/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 new file mode 100644 index 00000000000..9921dcd20bf --- /dev/null +++ b/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 @@ -0,0 +1,2 @@ +org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure.ElasticSearchChatMemoryRepositoryAutoConfiguration + diff --git a/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 b/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 new file mode 100644 index 00000000000..0e52b5224bc --- /dev/null +++ b/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 @@ -0,0 +1,115 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure; + +import java.util.List; +import java.util.UUID; + +import org.junit.jupiter.api.Test; +import org.testcontainers.elasticsearch.ElasticsearchContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link ElasticSearchChatMemoryRepositoryAutoConfiguration}. + * + * @author Fu Jian + * @since 1.1.0 + */ +@Testcontainers +class ElasticSearchChatMemoryRepositoryAutoConfigurationIT { + + @Container + static ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer( + "docker.elastic.co/elasticsearch/elasticsearch:8.10.2") + .withEnv("xpack.security.enabled", "false") + .withEnv("xpack.security.http.ssl.enabled", "false"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(ElasticsearchRestClientAutoConfiguration.class, + ElasticSearchChatMemoryRepositoryAutoConfiguration.class)) + .withPropertyValues("spring.elasticsearch.uris=http://" + elasticsearchContainer.getHost() + ":" + + elasticsearchContainer.getMappedPort(9200)) + .withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=autoconfig-test-chat-memory"); + + @Test + void addAndGet() { + this.contextRunner.run(context -> { + ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class); + + String conversationId = UUID.randomUUID().toString(); + assertThat(memory.findByConversationId(conversationId)).isEmpty(); + + memory.saveAll(conversationId, List.of(new UserMessage("test question"))); + + assertThat(memory.findByConversationId(conversationId)).hasSize(1); + assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER); + assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question"); + + memory.deleteByConversationId(conversationId); + assertThat(memory.findByConversationId(conversationId)).isEmpty(); + + memory.saveAll(conversationId, + List.of(new UserMessage("test question"), new AssistantMessage("test answer"))); + + assertThat(memory.findByConversationId(conversationId)).hasSize(2); + assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER); + assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question"); + assertThat(memory.findByConversationId(conversationId).get(1).getMessageType()) + .isEqualTo(MessageType.ASSISTANT); + assertThat(memory.findByConversationId(conversationId).get(1).getText()).isEqualTo("test answer"); + }); + } + + @Test + void propertiesConfiguration() { + this.contextRunner + .withPropertyValues("spring.ai.chat.memory.repository.elasticsearch.index-name=custom-testindex") + .run(context -> { + ElasticSearchChatMemoryRepositoryProperties properties = context + .getBean(ElasticSearchChatMemoryRepositoryProperties.class); + assertThat(properties.getIndexName()).isEqualTo("custom-testindex"); + }); + } + + @Test + void findConversationIds() { + this.contextRunner.run(context -> { + ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class); + + String conversationId1 = UUID.randomUUID().toString(); + String conversationId2 = UUID.randomUUID().toString(); + + memory.saveAll(conversationId1, List.of(new UserMessage("test question 1"))); + memory.saveAll(conversationId2, List.of(new UserMessage("test question 2"))); + + List conversationIds = memory.findConversationIds(); + assertThat(conversationIds).contains(conversationId1, conversationId2); + }); + } + +} diff --git a/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 b/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 new file mode 100644 index 00000000000..8c10cc6a83b --- /dev/null +++ b/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 @@ -0,0 +1,47 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.model.chat.memory.repository.elasticsearch.autoconfigure; + +import org.junit.jupiter.api.Test; + +import org.springframework.ai.chat.memory.repository.elasticsearch.ElasticSearchChatMemoryRepositoryConfig; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link ElasticSearchChatMemoryRepositoryProperties}. + * + * @author Fu Jian + * @since 1.1.0 + */ +class ElasticSearchChatMemoryRepositoryPropertiesTest { + + @Test + void defaultValues() { + var props = new ElasticSearchChatMemoryRepositoryProperties(); + assertThat(props.getIndexName()).isEqualTo(ElasticSearchChatMemoryRepositoryConfig.DEFAULT_INDEX_NAME); + } + + @Test + void customValues() { + var props = new ElasticSearchChatMemoryRepositoryProperties(); + props.setIndexName("custom_chat_memory"); + + assertThat(props.getIndexName()).isEqualTo("custom_chat_memory"); + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/pom.xml b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/pom.xml new file mode 100644 index 00000000000..c9de07e0880 --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/pom.xml @@ -0,0 +1,88 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 1.1.0-SNAPSHOT + ../../../pom.xml + + + spring-ai-model-chat-memory-repository-elasticsearch + Spring AI Elasticsearch Chat Memory Repository + Spring AI Elasticsearch Chat Memory Repository implementation + + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + org.springframework.ai + spring-ai-model + ${project.version} + + + + co.elastic.clients + elasticsearch-java + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.springframework.ai + spring-ai-test + ${project.version} + test + + + + org.springframework.boot + spring-boot-testcontainers + test + + + + org.testcontainers + elasticsearch + test + + + + org.testcontainers + junit-jupiter + test + + + + + + + diff --git a/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepository.java new file mode 100644 index 00000000000..723f1330049 --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepository.java @@ -0,0 +1,230 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.memory.repository.elasticsearch; + +import java.io.IOException; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.stream.Collectors; + +import co.elastic.clients.elasticsearch.ElasticsearchClient; +import co.elastic.clients.elasticsearch._types.SortOrder; +import co.elastic.clients.elasticsearch.core.BulkRequest; +import co.elastic.clients.elasticsearch.core.SearchResponse; +import co.elastic.clients.elasticsearch.core.search.Hit; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.chat.memory.ChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.util.Assert; + +/** + * An implementation of {@link ChatMemoryRepository} for Elasticsearch. + * + * @author Fu Jian + * @since 1.1.0 + */ +public final class ElasticSearchChatMemoryRepository implements ChatMemoryRepository { + + public static final String CONVERSATION_TS = ElasticSearchChatMemoryRepository.class.getSimpleName() + + "_message_timestamp"; + + private static final Logger logger = LoggerFactory.getLogger(ElasticSearchChatMemoryRepository.class); + + private final ElasticsearchClient client; + + private final String indexName; + + private ElasticSearchChatMemoryRepository(ElasticSearchChatMemoryRepositoryConfig config) { + Assert.notNull(config, "config cannot be null"); + this.client = config.getClient(); + this.indexName = config.getIndexName(); + } + + public static ElasticSearchChatMemoryRepository create(ElasticSearchChatMemoryRepositoryConfig config) { + return new ElasticSearchChatMemoryRepository(config); + } + + @Override + public List findConversationIds() { + logger.info("Finding all conversation IDs from Elasticsearch"); + + try { + SearchResponse response = this.client.search( + s -> s.index(this.indexName) + .size(0) + .aggregations("unique_conversations", a -> a.terms(t -> t.field("conversationId").size(10000))), + Map.class); + + return response.aggregations() + .get("unique_conversations") + .sterms() + .buckets() + .array() + .stream() + .map(b -> b.key().stringValue()) + .collect(Collectors.toList()); + } + catch (IOException e) { + throw new RuntimeException("Failed to find conversation IDs", e); + } + } + + @Override + public List findByConversationId(String conversationId) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + logger.info("Finding messages for conversation: {}", conversationId); + + try { + SearchResponse response = this.client.search(s -> s.index(this.indexName) + .query(q -> q.term(t -> t.field("conversationId").value(conversationId))) + .sort(sort -> sort.field(f -> f.field("sequenceNumber").order(SortOrder.Asc))) + .size(10000), Map.class); + + return response.hits() + .hits() + .stream() + .map(Hit::source) + .map(this::mapToMessage) + .collect(Collectors.toList()); + } + catch (IOException e) { + throw new RuntimeException("Failed to find messages for conversation: " + conversationId, e); + } + } + + @Override + public void saveAll(String conversationId, List messages) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + Assert.notNull(messages, "messages cannot be null"); + Assert.noNullElements(messages, "messages cannot contain null elements"); + + logger.info("Saving {} messages for conversation: {}", messages.size(), conversationId); + + // First delete existing messages for this conversation + deleteByConversationId(conversationId); + + // Then save the new messages + Instant timestamp = Instant.now(); + BulkRequest.Builder bulkBuilder = new BulkRequest.Builder(); + + for (int i = 0; i < messages.size(); i++) { + Message message = messages.get(i); + Map doc = createMessageDocument(conversationId, message, timestamp, i); + String docId = (String) doc.get("id"); + + bulkBuilder.operations(op -> op.index(idx -> idx.index(this.indexName).id(docId).document(doc))); + } + + bulkRequest(bulkBuilder.build()); + } + + @Override + public void deleteByConversationId(String conversationId) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + logger.info("Deleting messages for conversation: {}", conversationId); + + try { + this.client.deleteByQuery(d -> d.index(this.indexName) + .query(q -> q.term(t -> t.field("conversationId").value(conversationId)))); + } + catch (IOException e) { + throw new RuntimeException("Failed to delete messages for conversation: " + conversationId, e); + } + } + + private void bulkRequest(BulkRequest bulkRequest) { + try { + co.elastic.clients.elasticsearch.core.BulkResponse bulkResponse = this.client.bulk(bulkRequest); + if (bulkResponse.errors()) { + throw new IllegalStateException("Bulk operation failed"); + } + } + catch (IOException e) { + throw new RuntimeException("Failed to execute bulk request", e); + } + } + + private Map createMessageDocument(String conversationId, Message message, Instant timestamp, + int sequenceNumber) { + Map doc = new HashMap<>(); + doc.put("id", UUID.randomUUID().toString()); + doc.put("conversationId", conversationId); + doc.put("messageType", message.getMessageType().name()); + doc.put("content", message.getText()); + doc.put("sequenceNumber", sequenceNumber); + + // Add timestamp from metadata or use provided timestamp + Instant messageTimestamp = (Instant) message.getMetadata().get(CONVERSATION_TS); + if (messageTimestamp == null) { + messageTimestamp = timestamp; + message.getMetadata().put(CONVERSATION_TS, messageTimestamp); + } + doc.put("messageTimestamp", messageTimestamp.toEpochMilli()); + + // Store any additional metadata + Map filteredMetadata = message.getMetadata() + .entrySet() + .stream() + .filter(entry -> !CONVERSATION_TS.equals(entry.getKey())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + if (!filteredMetadata.isEmpty()) { + doc.put("metadata", filteredMetadata); + } + + return doc; + } + + private Message mapToMessage(Map doc) { + String content = (String) doc.get("content"); + String messageTypeStr = (String) doc.get("messageType"); + MessageType messageType = MessageType.valueOf(messageTypeStr); + + // Reconstruct metadata + Map metadata = new HashMap<>(); + if (doc.containsKey("messageTimestamp")) { + Long timestampMillis = ((Number) doc.get("messageTimestamp")).longValue(); + metadata.put(CONVERSATION_TS, Instant.ofEpochMilli(timestampMillis)); + } + + // Add any additional metadata that was stored + @SuppressWarnings("unchecked") + Map additionalMetadata = (Map) doc.get("metadata"); + if (additionalMetadata != null) { + metadata.putAll(additionalMetadata); + } + + return switch (messageType) { + case ASSISTANT -> new AssistantMessage(content, metadata); + case USER -> UserMessage.builder().text(content).metadata(metadata).build(); + case SYSTEM -> SystemMessage.builder().text(content).metadata(metadata).build(); + case TOOL -> ToolResponseMessage.builder().responses(List.of()).metadata(metadata).build(); + default -> throw new IllegalStateException(String.format("Unknown message type: %s", messageTypeStr)); + }; + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryConfig.java b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryConfig.java new file mode 100644 index 00000000000..85afbe320d4 --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/main/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryConfig.java @@ -0,0 +1,107 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.memory.repository.elasticsearch; + +import java.io.IOException; + +import co.elastic.clients.elasticsearch.ElasticsearchClient; + +import org.springframework.util.Assert; + +/** + * Configuration for the Elasticsearch Chat Memory store. + * + * @author Fu Jian + * @since 1.1.0 + */ +public final class ElasticSearchChatMemoryRepositoryConfig { + + public static final String DEFAULT_INDEX_NAME = "ai_chat_memory"; + + private final ElasticsearchClient client; + + private final String indexName; + + private ElasticSearchChatMemoryRepositoryConfig(Builder builder) { + this.client = builder.client; + this.indexName = builder.indexName; + this.initializeIndex(); + } + + public static Builder builder() { + return new Builder(); + } + + public ElasticsearchClient getClient() { + return this.client; + } + + public String getIndexName() { + return this.indexName; + } + + private void initializeIndex() { + try { + // Check if index exists + boolean exists = this.client.indices().exists(ex -> ex.index(this.indexName)).value(); + + if (!exists) { + // Create index with mapping + this.client.indices() + .create(cr -> cr.index(this.indexName) + .mappings(m -> m.properties("conversationId", p -> p.keyword(k -> k)) + .properties("messageType", p -> p.keyword(k -> k)) + .properties("content", p -> p.text(t -> t)) + .properties("sequenceNumber", p -> p.integer(i -> i)) + .properties("messageTimestamp", p -> p.long_(l -> l)) + .properties("metadata", p -> p.object(o -> o.enabled(true))))); + } + } + catch (IOException e) { + throw new RuntimeException("Failed to initialize Elasticsearch index", e); + } + } + + public static final class Builder { + + private ElasticsearchClient client; + + private String indexName = DEFAULT_INDEX_NAME; + + private Builder() { + } + + public Builder withClient(ElasticsearchClient client) { + this.client = client; + return this; + } + + public Builder withIndexName(String indexName) { + this.indexName = indexName; + return this; + } + + public ElasticSearchChatMemoryRepositoryConfig build() { + Assert.notNull(this.client, "ElasticsearchClient cannot be null"); + Assert.hasText(this.indexName, "indexName cannot be null or empty"); + + return new ElasticSearchChatMemoryRepositoryConfig(this); + } + + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryIT.java b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryIT.java new file mode 100644 index 00000000000..6579c87c152 --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-elasticsearch/src/test/java/org/springframework/ai/chat/memory/repository/elasticsearch/ElasticSearchChatMemoryRepositoryIT.java @@ -0,0 +1,317 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.memory.repository.elasticsearch; + +import java.io.IOException; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.TimeUnit; + +import co.elastic.clients.elasticsearch.ElasticsearchClient; +import co.elastic.clients.json.jackson.JacksonJsonpMapper; +import co.elastic.clients.transport.rest_client.RestClientTransport; +import org.apache.http.HttpHost; +import org.apache.http.auth.AuthScope; +import org.apache.http.auth.UsernamePasswordCredentials; +import org.apache.http.impl.client.BasicCredentialsProvider; +import org.elasticsearch.client.RestClient; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.testcontainers.elasticsearch.ElasticsearchContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.ai.chat.memory.ChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link ElasticSearchChatMemoryRepository}. + * + * @author Fu Jian + * @since 1.1.0 + */ +@Testcontainers +class ElasticSearchChatMemoryRepositoryIT { + + @Container + static ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer( + "docker.elastic.co/elasticsearch/elasticsearch:8.10.2") + .withEnv("xpack.security.enabled", "false") + .withEnv("xpack.security.http.ssl.enabled", "false"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(ElasticSearchChatMemoryRepositoryIT.TestApplication.class); + + @Test + void ensureBeansGetsCreated() { + this.contextRunner.run(context -> { + ElasticSearchChatMemoryRepository memory = context.getBean(ElasticSearchChatMemoryRepository.class); + Assertions.assertNotNull(memory); + }); + } + + @ParameterizedTest + @CsvSource({ "Message from assistant,ASSISTANT", "Message from user,USER", "Message from system,SYSTEM" }) + void add_shouldInsertSingleMessage(String content, MessageType messageType) { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + assertThat(chatMemory).isInstanceOf(ElasticSearchChatMemoryRepository.class); + + var conversationId = UUID.randomUUID().toString(); + var message = switch (messageType) { + case ASSISTANT -> new AssistantMessage(content); + case USER -> new UserMessage(content); + case SYSTEM -> new SystemMessage(content); + default -> throw new IllegalArgumentException("Type not supported: " + messageType); + }; + + java.time.Instant beforeSave = java.time.Instant.now(); + chatMemory.saveAll(conversationId, List.of(message)); + java.time.Instant afterSave = java.time.Instant.now(); + + sleepForSearchable(); + assertThat(chatMemory.findConversationIds()).isNotEmpty(); + assertThat(chatMemory.findConversationIds()).contains(conversationId); + + List retrievedMessages = chatMemory.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(1); + assertThat(retrievedMessages.get(0).getText()).isEqualTo(content); + assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(messageType); + + // Verify timestamp was automatically added + Object timestampObj = retrievedMessages.get(0) + .getMetadata() + .get(ElasticSearchChatMemoryRepository.CONVERSATION_TS); + assertThat(timestampObj).as("Timestamp should be automatically added").isNotNull(); + assertThat(timestampObj).isInstanceOf(java.time.Instant.class); + java.time.Instant timestamp = (java.time.Instant) timestampObj; + assertThat(timestamp).isBetween(beforeSave.minusSeconds(1), afterSave.plusSeconds(1)); + }); + } + + private static void sleepForSearchable() throws InterruptedException { + TimeUnit.SECONDS.sleep(2); + } + + @Test + void shouldSaveAndRetrieveMultipleMessages() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + assertThat(chatMemory).isInstanceOf(ElasticSearchChatMemoryRepository.class); + + var conversationId = UUID.randomUUID().toString(); + + // Use 5 messages to verify sequenceNumber preserves order + List messages = List.of(new SystemMessage("System message"), new UserMessage("User message"), + new AssistantMessage("Assistant message"), new UserMessage("Second user message"), + new AssistantMessage("Second assistant message")); + + java.time.Instant beforeSave = java.time.Instant.now(); + chatMemory.saveAll(conversationId, messages); + java.time.Instant afterSave = java.time.Instant.now(); + + sleepForSearchable(); + + List retrievedMessages = chatMemory.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(5); + + // Verify sequenceNumber preserves message order + assertThat(retrievedMessages.get(0).getText()).isEqualTo("System message"); + assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(MessageType.SYSTEM); + + assertThat(retrievedMessages.get(1).getText()).isEqualTo("User message"); + assertThat(retrievedMessages.get(1).getMessageType()).isEqualTo(MessageType.USER); + + assertThat(retrievedMessages.get(2).getText()).isEqualTo("Assistant message"); + assertThat(retrievedMessages.get(2).getMessageType()).isEqualTo(MessageType.ASSISTANT); + + assertThat(retrievedMessages.get(3).getText()).isEqualTo("Second user message"); + assertThat(retrievedMessages.get(3).getMessageType()).isEqualTo(MessageType.USER); + + assertThat(retrievedMessages.get(4).getText()).isEqualTo("Second assistant message"); + assertThat(retrievedMessages.get(4).getMessageType()).isEqualTo(MessageType.ASSISTANT); + + // Verify timestamp consistency: all messages saved together should have the + // same timestamp + java.time.Instant firstTimestamp = null; + for (Message message : retrievedMessages) { + Object timestampObj = message.getMetadata().get(ElasticSearchChatMemoryRepository.CONVERSATION_TS); + assertThat(timestampObj).isNotNull(); + + java.time.Instant timestamp = (java.time.Instant) timestampObj; + assertThat(timestamp).isBetween(beforeSave.minusSeconds(1), afterSave.plusSeconds(1)); + + if (firstTimestamp == null) { + firstTimestamp = timestamp; + } + else { + // All messages in same saveAll should have identical timestamp + assertThat(timestamp).isEqualTo(firstTimestamp); + } + } + }); + } + + @Test + void shouldReplaceExistingMessages() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + var conversationId = UUID.randomUUID().toString(); + + // Save initial messages (3 messages) + List initialMessages = List.of(new UserMessage("Initial user message"), + new AssistantMessage("Initial assistant message"), new UserMessage("Initial third message")); + chatMemory.saveAll(conversationId, initialMessages); + + sleepForSearchable(); + + // Verify initial save and capture timestamp + List retrievedMessages = chatMemory.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(3); + java.time.Instant initialTimestamp = (java.time.Instant) retrievedMessages.get(0) + .getMetadata() + .get(ElasticSearchChatMemoryRepository.CONVERSATION_TS); + + // Wait to ensure timestamp difference + TimeUnit.MILLISECONDS.sleep(100); + + // Replace with new messages (2 messages - verify sequenceNumber restarts) + List newMessages = List.of(new SystemMessage("New system message"), + new UserMessage("New user message")); + chatMemory.saveAll(conversationId, newMessages); + + sleepForSearchable(); + + // Verify replacement + retrievedMessages = chatMemory.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(2); + assertThat(retrievedMessages.get(0).getText()).isEqualTo("New system message"); + assertThat(retrievedMessages.get(1).getText()).isEqualTo("New user message"); + + // Verify timestamp was updated (should be after initial timestamp) + java.time.Instant updatedTimestamp = (java.time.Instant) retrievedMessages.get(0) + .getMetadata() + .get(ElasticSearchChatMemoryRepository.CONVERSATION_TS); + assertThat(updatedTimestamp).isAfter(initialTimestamp); + }); + } + + @Test + void shouldDeleteConversation() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + var conversationId = UUID.randomUUID().toString(); + + // Save messages + List messages = List.of(new UserMessage("User message"), + new AssistantMessage("Assistant message")); + chatMemory.saveAll(conversationId, messages); + + sleepForSearchable(); + + // Verify messages exist + assertThat(chatMemory.findByConversationId(conversationId)).hasSize(2); + + // Delete conversation + chatMemory.deleteByConversationId(conversationId); + + sleepForSearchable(); + + // Verify messages are deleted + assertThat(chatMemory.findByConversationId(conversationId)).isEmpty(); + }); + } + + @Test + void shouldFindAllConversationIds() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + var conversationId1 = UUID.randomUUID().toString(); + var conversationId2 = UUID.randomUUID().toString(); + + // Save messages for two conversations + chatMemory.saveAll(conversationId1, List.of(new UserMessage("Message 1"))); + chatMemory.saveAll(conversationId2, List.of(new UserMessage("Message 2"))); + + sleepForSearchable(); + + // Verify both conversation IDs are found + List conversationIds = chatMemory.findConversationIds(); + assertThat(conversationIds).contains(conversationId1, conversationId2); + }); + } + + @Test + void shouldHandleEmptyConversation() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemoryRepository.class); + var conversationId = UUID.randomUUID().toString(); + + // Try to find messages for non-existent conversation + List messages = chatMemory.findByConversationId(conversationId); + assertThat(messages).isEmpty(); + + // Delete non-existent conversation (should not throw) + chatMemory.deleteByConversationId(conversationId); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration + static class TestApplication { + + @Bean + public ElasticsearchClient elasticsearchClient() throws IOException { + var credentialsProvider = new BasicCredentialsProvider(); + credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials("elastic", "changeme")); + + var restClient = RestClient.builder(HttpHost.create(elasticsearchContainer.getHttpHostAddress())).build(); + + var transport = new RestClientTransport(restClient, new JacksonJsonpMapper()); + return new ElasticsearchClient(transport); + } + + @Bean + public ElasticSearchChatMemoryRepositoryConfig elasticSearchChatMemoryRepositoryConfig( + ElasticsearchClient elasticsearchClient) { + return ElasticSearchChatMemoryRepositoryConfig.builder() + .withClient(elasticsearchClient) + .withIndexName("test-chat-memory") + .build(); + } + + @Bean + public ElasticSearchChatMemoryRepository elasticSearchChatMemoryRepository( + ElasticSearchChatMemoryRepositoryConfig config) { + return ElasticSearchChatMemoryRepository.create(config); + } + + } + +} diff --git a/pom.xml b/pom.xml index abb2bc543a3..8c38c62f7b0 100644 --- a/pom.xml +++ b/pom.xml @@ -43,6 +43,7 @@ memory/repository/spring-ai-model-chat-memory-repository-cassandra memory/repository/spring-ai-model-chat-memory-repository-cosmos-db + memory/repository/spring-ai-model-chat-memory-repository-elasticsearch memory/repository/spring-ai-model-chat-memory-repository-jdbc memory/repository/spring-ai-model-chat-memory-repository-neo4j @@ -90,6 +91,7 @@ auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cassandra auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db + auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-elasticsearch auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-jdbc auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j