From 831f3951d98c65b9c88c37bf681fe9fee68a958a Mon Sep 17 00:00:00 2001 From: leijendary Date: Fri, 24 Jan 2025 20:06:07 +0100 Subject: [PATCH] feat: JDBC implementation of ChatMemory Signed-off-by: leijendary feat: JDBC implementation of ChatMemory Signed-off-by: leijendary --- README.md | 6 +- .../spring-ai-chat-memory-jdbc/README.md | 1 + .../spring-ai-chat-memory-jdbc/pom.xml | 107 +++++++++ .../ai/chat/memory/JdbcChatMemory.java | 108 ++++++++++ .../ai/chat/memory/JdbcChatMemoryConfig.java | 66 ++++++ .../aot/hint/JdbcChatMemoryRuntimeHints.java | 28 +++ .../resources/META-INF/spring/aot.factories | 2 + .../ai/chat/memory/schema-drop-mariadb.sql | 1 + .../ai/chat/memory/schema-drop-postgresql.sql | 1 + .../ai/chat/memory/schema-mariadb.sql | 10 + .../ai/chat/memory/schema-postgresql.sql | 9 + .../chat/memory/JdbcChatMemoryConfigTest.java | 34 +++ .../ai/chat/memory/JdbcChatMemoryIT.java | 203 ++++++++++++++++++ .../hint/JdbcChatMemoryRuntimeHintsTest.java | 82 +++++++ pom.xml | 7 +- spring-ai-bom/pom.xml | 6 + .../modules/ROOT/pages/api/chatclient.adoc | 11 +- spring-ai-spring-boot-autoconfigure/pom.xml | 8 + .../jdbc/JdbcChatMemoryAutoConfiguration.java | 65 ++++++ ...ryDataSourceScriptDatabaseInitializer.java | 35 +++ .../memory/jdbc/JdbcChatMemoryProperties.java | 31 +++ ...ot.autoconfigure.AutoConfiguration.imports | 1 + .../JdbcChatMemoryAutoConfigurationIT.java | 102 +++++++++ ...aSourceScriptDatabaseInitializerTests.java | 52 +++++ .../jdbc/JdbcChatMemoryPropertiesTests.java | 43 ++++ .../pom.xml | 58 +++++ 26 files changed, 1071 insertions(+), 6 deletions(-) create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/README.md create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/pom.xml create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemory.java create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfig.java create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHints.java create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/META-INF/spring/aot.factories create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-mariadb.sql create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-postgresql.sql create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-mariadb.sql create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-postgresql.sql create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfigTest.java create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryIT.java create mode 100644 chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHintsTest.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfiguration.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializer.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryProperties.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfigurationIT.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializerTests.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryPropertiesTests.java create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-chat-memory-jdbc/pom.xml diff --git a/README.md b/README.md index e4b11b62c41..462e9eadc92 100644 --- a/README.md +++ b/README.md @@ -101,9 +101,9 @@ One way to run integration tests on part of the code is to first do a quick comp ```shell ./mvnw clean install -DskipTests -Dmaven.javadoc.skip=true ``` -Then run the integration test for a specifi module using the `-pl` option +Then run the integration test for a specific module using the `-pl` option ```shell -./mvnw verify -Pintegration-tests -pl spring-ai-spring-boot-autoconfigure +./mvnw verify -Pintegration-tests -pl spring-ai-spring-boot-autoconfigure ``` ### Documentation @@ -134,4 +134,4 @@ To build with checkstyles enabled. Checkstyles are currently disabled, but you can enable them by doing the following: ```shell ./mvnw clean package -DskipTests -Ddisable.checks=false -``` \ No newline at end of file +``` diff --git a/chat-memory/spring-ai-chat-memory-jdbc/README.md b/chat-memory/spring-ai-chat-memory-jdbc/README.md new file mode 100644 index 00000000000..8e100ad20a3 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/README.md @@ -0,0 +1 @@ +[Chat Memory Documentation](https://docs.spring.io/spring-ai/reference/api/chatclient.html#_chat_memory) diff --git a/chat-memory/spring-ai-chat-memory-jdbc/pom.xml b/chat-memory/spring-ai-chat-memory-jdbc/pom.xml new file mode 100644 index 00000000000..61e1737e263 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/pom.xml @@ -0,0 +1,107 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-chat-memory-jdbc + jar + Spring AI Chat Memory JDBC + Spring AI Chat Memory implementation with JDBC + 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 + + + + 17 + 17 + + + + + org.springframework.ai + spring-ai-core + ${project.parent.version} + + + + com.zaxxer + HikariCP + + + + org.springframework + spring-jdbc + + + + org.postgresql + postgresql + ${postgresql.version} + true + + + + org.mariadb.jdbc + mariadb-java-client + ${mariadb.version} + true + + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + testcontainers + test + + + + org.testcontainers + postgresql + test + + + + org.testcontainers + mariadb + test + + + + org.testcontainers + junit-jupiter + test + + + diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemory.java b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemory.java new file mode 100644 index 00000000000..eb54c089b52 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemory.java @@ -0,0 +1,108 @@ +/* + * Copyright 2024-2024 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; + +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.List; + +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.UserMessage; +import org.springframework.jdbc.core.BatchPreparedStatementSetter; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.jdbc.core.RowMapper; + +/** + * An implementation of {@link ChatMemory} for JDBC. Creating an instance of + * JdbcChatMemory example: + * JdbcChatMemory.create(JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build()); + * + * @author Jonathan Leijendekker + * @since 1.0.0 + */ +public class JdbcChatMemory implements ChatMemory { + + private static final String QUERY_ADD = """ + INSERT INTO ai_chat_memory (conversation_id, content, type) VALUES (?, ?, ?)"""; + + private static final String QUERY_GET = """ + SELECT content, type FROM ai_chat_memory WHERE conversation_id = ? ORDER BY "timestamp" DESC LIMIT ?"""; + + private static final String QUERY_CLEAR = "DELETE FROM ai_chat_memory WHERE conversation_id = ?"; + + private final JdbcTemplate jdbcTemplate; + + public JdbcChatMemory(JdbcChatMemoryConfig config) { + this.jdbcTemplate = config.getJdbcTemplate(); + } + + public static JdbcChatMemory create(JdbcChatMemoryConfig config) { + return new JdbcChatMemory(config); + } + + @Override + public void add(String conversationId, List messages) { + this.jdbcTemplate.batchUpdate(QUERY_ADD, new AddBatchPreparedStatement(conversationId, messages)); + } + + @Override + public List get(String conversationId, int lastN) { + return this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId, lastN); + } + + @Override + public void clear(String conversationId) { + this.jdbcTemplate.update(QUERY_CLEAR, conversationId); + } + + private record AddBatchPreparedStatement(String conversationId, + List messages) implements BatchPreparedStatementSetter { + @Override + public void setValues(PreparedStatement ps, int i) throws SQLException { + var message = this.messages.get(i); + + ps.setString(1, this.conversationId); + ps.setString(2, message.getText()); + ps.setString(3, message.getMessageType().name()); + } + + @Override + public int getBatchSize() { + return this.messages.size(); + } + } + + private static class MessageRowMapper implements RowMapper { + + @Override + public Message mapRow(ResultSet rs, int i) throws SQLException { + var content = rs.getString(1); + var type = MessageType.valueOf(rs.getString(2)); + + return switch (type) { + case USER -> new UserMessage(content); + case ASSISTANT -> new AssistantMessage(content); + default -> null; + }; + } + + } + +} diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfig.java b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfig.java new file mode 100644 index 00000000000..ad212bd6869 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfig.java @@ -0,0 +1,66 @@ +/* + * Copyright 2024-2024 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; + +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.util.Assert; + +/** + * Configuration for {@link JdbcChatMemory}. + * + * @author Jonathan Leijendekker + * @since 1.0.0 + */ +public final class JdbcChatMemoryConfig { + + private final JdbcTemplate jdbcTemplate; + + private JdbcChatMemoryConfig(Builder builder) { + this.jdbcTemplate = builder.jdbcTemplate; + } + + public static Builder builder() { + return new Builder(); + } + + JdbcTemplate getJdbcTemplate() { + return this.jdbcTemplate; + } + + public static final class Builder { + + private JdbcTemplate jdbcTemplate; + + private Builder() { + } + + public Builder jdbcTemplate(JdbcTemplate jdbcTemplate) { + Assert.notNull(jdbcTemplate, "jdbc template must not be null"); + + this.jdbcTemplate = jdbcTemplate; + return this; + } + + public JdbcChatMemoryConfig build() { + Assert.notNull(this.jdbcTemplate, "jdbc template must not be null"); + + return new JdbcChatMemoryConfig(this); + } + + } + +} diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHints.java b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHints.java new file mode 100644 index 00000000000..9b961138726 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHints.java @@ -0,0 +1,28 @@ +package org.springframework.ai.chat.memory.aot.hint; + +import javax.sql.DataSource; + +import org.springframework.aot.hint.MemberCategory; +import org.springframework.aot.hint.RuntimeHints; +import org.springframework.aot.hint.RuntimeHintsRegistrar; + +/** + * A {@link RuntimeHintsRegistrar} for JDBC Chat Memory hints + * + * @author Jonathan Leijendekker + */ +class JdbcChatMemoryRuntimeHints implements RuntimeHintsRegistrar { + + @Override + public void registerHints(RuntimeHints hints, ClassLoader classLoader) { + hints.reflection() + .registerType(DataSource.class, (hint) -> hint.withMembers(MemberCategory.INVOKE_DECLARED_METHODS)); + + hints.resources() + .registerPattern("org/springframework/ai/chat/memory/schema-drop-mariadb.sql") + .registerPattern("org/springframework/ai/chat/memory/schema-drop-postgresql.sql") + .registerPattern("org/springframework/ai/chat/memory/schema-mariadb.sql") + .registerPattern("org/springframework/ai/chat/memory/schema-postgresql.sql"); + } + +} diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/META-INF/spring/aot.factories b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/META-INF/spring/aot.factories new file mode 100644 index 00000000000..1877d6377a4 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/META-INF/spring/aot.factories @@ -0,0 +1,2 @@ +org.springframework.aot.hint.RuntimeHintsRegistrar=\ +org.springframework.ai.chat.memory.aot.hint.JdbcChatMemoryRuntimeHints diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-mariadb.sql b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-mariadb.sql new file mode 100644 index 00000000000..72f313114ba --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-mariadb.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS ai_chat_memory; diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-postgresql.sql b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-postgresql.sql new file mode 100644 index 00000000000..72f313114ba --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-drop-postgresql.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS ai_chat_memory; diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-mariadb.sql b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-mariadb.sql new file mode 100644 index 00000000000..c5024336ef1 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-mariadb.sql @@ -0,0 +1,10 @@ +CREATE TABLE IF NOT EXISTS ai_chat_memory ( + conversation_id VARCHAR(36) NOT NULL, + content TEXT NOT NULL, + type VARCHAR(10) NOT NULL, + `timestamp` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT type_check CHECK (type IN ('USER', 'ASSISTANT')) +); + +CREATE INDEX IF NOT EXISTS ai_chat_memory_conversation_id_timestamp_idx +ON ai_chat_memory(conversation_id, `timestamp` DESC); diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-postgresql.sql b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-postgresql.sql new file mode 100644 index 00000000000..c8d5909f04e --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/schema-postgresql.sql @@ -0,0 +1,9 @@ +CREATE TABLE IF NOT EXISTS ai_chat_memory ( + conversation_id VARCHAR(36) NOT NULL, + content TEXT NOT NULL, + type VARCHAR(10) NOT NULL CHECK (type IN ('USER', 'ASSISTANT')), + "timestamp" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS ai_chat_memory_conversation_id_timestamp_idx +ON ai_chat_memory(conversation_id, "timestamp" DESC); diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfigTest.java b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfigTest.java new file mode 100644 index 00000000000..5758c1195a6 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryConfigTest.java @@ -0,0 +1,34 @@ +package org.springframework.ai.chat.memory; + +import org.junit.jupiter.api.Test; + +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; + +/** + * @author Jonathan Leijendekker + */ +class JdbcChatMemoryConfigTest { + + @Test + void setValues() { + var jdbcTemplate = mock(JdbcTemplate.class); + var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build(); + + assertThat(config.getJdbcTemplate()).isEqualTo(jdbcTemplate); + } + + @Test + void setJdbcTemplateToNull_shouldThrow() { + assertThatThrownBy(() -> JdbcChatMemoryConfig.builder().jdbcTemplate(null)); + } + + @Test + void buildWithNullJdbcTemplate_shouldThrow() { + assertThatThrownBy(() -> JdbcChatMemoryConfig.builder().build()); + } + +} diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryIT.java b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryIT.java new file mode 100644 index 00000000000..bd8b852a0d2 --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/JdbcChatMemoryIT.java @@ -0,0 +1,203 @@ +/* + * Copyright 2024-2024 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; + +import java.sql.Timestamp; +import java.util.List; +import java.util.UUID; + +import javax.sql.DataSource; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +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.UserMessage; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties; +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Primary; +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonathan Leijendekker + */ +@Testcontainers +class JdbcChatMemoryIT { + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("postgres:17") + .withDatabaseName("chat_memory_test") + .withUsername("postgres") + .withPassword("postgres"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class) + .withPropertyValues( + // JdbcTemplate configuration + String.format("app.datasource.url=%s", postgresContainer.getJdbcUrl()), + String.format("app.datasource.username=%s", postgresContainer.getUsername()), + String.format("app.datasource.password=%s", postgresContainer.getPassword()), + "app.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + @Test + void correctChatMemoryInstance() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemory.class); + + assertThat(chatMemory).isInstanceOf(JdbcChatMemory.class); + }); + } + + @ParameterizedTest + @CsvSource({ "Message from assistant,ASSISTANT", "Message from user,USER" }) + void add_shouldInsertSingleMessage(String content, MessageType messageType) { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemory.class); + var conversationId = UUID.randomUUID().toString(); + var message = switch (messageType) { + case ASSISTANT -> new AssistantMessage(content + " - " + conversationId); + case USER -> new UserMessage(content + " - " + conversationId); + default -> throw new IllegalArgumentException("Type not supported: " + messageType); + }; + + chatMemory.add(conversationId, message); + + var jdbcTemplate = context.getBean(JdbcTemplate.class); + var query = "SELECT conversation_id, content, type, \"timestamp\" FROM ai_chat_memory WHERE conversation_id = ?"; + var result = jdbcTemplate.queryForMap(query, conversationId); + + assertThat(result.size()).isEqualTo(4); + assertThat(result.get("conversation_id")).isEqualTo(conversationId); + assertThat(result.get("content")).isEqualTo(message.getText()); + assertThat(result.get("type")).isEqualTo(messageType.name()); + assertThat(result.get("timestamp")).isInstanceOf(Timestamp.class); + }); + } + + @Test + void add_shouldInsertMessages() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemory.class); + var conversationId = UUID.randomUUID().toString(); + var messages = List.of(new AssistantMessage("Message from assistant - " + conversationId), + new UserMessage("Message from user - " + conversationId)); + + chatMemory.add(conversationId, messages); + + var jdbcTemplate = context.getBean(JdbcTemplate.class); + var query = "SELECT conversation_id, content, type, \"timestamp\" FROM ai_chat_memory WHERE conversation_id = ?"; + var results = jdbcTemplate.queryForList(query, conversationId); + + assertThat(results.size()).isEqualTo(messages.size()); + + for (var i = 0; i < messages.size(); i++) { + var message = messages.get(i); + var result = results.get(i); + + assertThat(result.get("conversation_id")).isNotNull(); + assertThat(result.get("conversation_id")).isEqualTo(conversationId); + assertThat(result.get("content")).isEqualTo(message.getText()); + assertThat(result.get("type")).isEqualTo(message.getMessageType().name()); + assertThat(result.get("timestamp")).isInstanceOf(Timestamp.class); + } + }); + } + + @Test + void get_shouldReturnMessages() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemory.class); + var conversationId = UUID.randomUUID().toString(); + var messages = List.of(new AssistantMessage("Message from assistant 1 - " + conversationId), + new AssistantMessage("Message from assistant 2 - " + conversationId), + new UserMessage("Message from user - " + conversationId)); + + chatMemory.add(conversationId, messages); + + var results = chatMemory.get(conversationId, Integer.MAX_VALUE); + + assertThat(results.size()).isEqualTo(messages.size()); + assertThat(results).isEqualTo(messages); + }); + } + + @Test + void clear_shouldDeleteMessages() { + this.contextRunner.run(context -> { + var chatMemory = context.getBean(ChatMemory.class); + var conversationId = UUID.randomUUID().toString(); + var messages = List.of(new AssistantMessage("Message from assistant - " + conversationId), + new UserMessage("Message from user - " + conversationId)); + + chatMemory.add(conversationId, messages); + + chatMemory.clear(conversationId); + + var jdbcTemplate = context.getBean(JdbcTemplate.class); + var count = jdbcTemplate.queryForObject("SELECT COUNT(*) FROM ai_chat_memory WHERE conversation_id = ?", + Integer.class, conversationId); + + assertThat(count).isZero(); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + static class TestApplication { + + @Bean + public ChatMemory chatMemory(JdbcTemplate jdbcTemplate) { + var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build(); + + return JdbcChatMemory.create(config); + } + + @Bean + public JdbcTemplate jdbcTemplate(DataSource dataSource) { + return new JdbcTemplate(dataSource); + } + + @Bean + @Primary + @ConfigurationProperties("app.datasource") + public DataSourceProperties dataSourceProperties() { + return new DataSourceProperties(); + } + + @Bean + public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) { + return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build(); + } + + } + +} diff --git a/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHintsTest.java b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHintsTest.java new file mode 100644 index 00000000000..b78b9fd0e1f --- /dev/null +++ b/chat-memory/spring-ai-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/aot/hint/JdbcChatMemoryRuntimeHintsTest.java @@ -0,0 +1,82 @@ +/* + * Copyright 2024-2024 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.aot.hint; + +import java.io.IOException; +import java.util.Arrays; +import java.util.stream.Stream; + +import javax.sql.DataSource; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import org.springframework.aot.hint.RuntimeHints; +import org.springframework.aot.hint.RuntimeHintsRegistrar; +import org.springframework.aot.hint.predicate.RuntimeHintsPredicates; +import org.springframework.core.io.Resource; +import org.springframework.core.io.support.PathMatchingResourcePatternResolver; +import org.springframework.core.io.support.SpringFactoriesLoader; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonathan Leijendekker + */ +class JdbcChatMemoryRuntimeHintsTest { + + private final RuntimeHints hints = new RuntimeHints(); + + private final JdbcChatMemoryRuntimeHints jdbcChatMemoryRuntimeHints = new JdbcChatMemoryRuntimeHints(); + + @Test + void aotFactoriesContainsRegistrar() { + var match = SpringFactoriesLoader.forResourceLocation("META-INF/spring/aot.factories") + .load(RuntimeHintsRegistrar.class) + .stream() + .anyMatch((registrar) -> registrar instanceof JdbcChatMemoryRuntimeHints); + + assertThat(match).isTrue(); + } + + @ParameterizedTest + @MethodSource("getSchemaFileNames") + void jdbcSchemasHasHints(String schemaFileName) { + this.jdbcChatMemoryRuntimeHints.registerHints(this.hints, getClass().getClassLoader()); + + var predicate = RuntimeHintsPredicates.resource() + .forResource("org/springframework/ai/chat/memory/" + schemaFileName); + + assertThat(predicate).accepts(this.hints); + } + + @Test + void dataSourceHasHints() { + this.jdbcChatMemoryRuntimeHints.registerHints(this.hints, getClass().getClassLoader()); + + assertThat(RuntimeHintsPredicates.reflection().onType(DataSource.class)).accepts(this.hints); + } + + private static Stream getSchemaFileNames() throws IOException { + var resources = new PathMatchingResourcePatternResolver() + .getResources("classpath*:org/springframework/ai/chat/memory/schema-*.sql"); + + return Arrays.stream(resources).map(Resource::getFilename); + } + +} diff --git a/pom.xml b/pom.xml index 5d2159e1775..2d7b6f2b4f8 100644 --- a/pom.xml +++ b/pom.xml @@ -65,6 +65,8 @@ vector-stores/spring-ai-typesense-store vector-stores/spring-ai-weaviate-store + chat-memory/spring-ai-chat-memory-jdbc + spring-ai-spring-boot-starters/spring-ai-starter-aws-opensearch-store spring-ai-spring-boot-starters/spring-ai-starter-azure-cosmos-db-store spring-ai-spring-boot-starters/spring-ai-starter-azure-store @@ -111,6 +113,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-azure-openai spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai spring-ai-spring-boot-starters/spring-ai-starter-bedrock-converse + spring-ai-spring-boot-starters/spring-ai-starter-chat-memory-jdbc spring-ai-spring-boot-starters/spring-ai-starter-huggingface spring-ai-spring-boot-starters/spring-ai-starter-minimax spring-ai-spring-boot-starters/spring-ai-starter-mistral-ai @@ -648,6 +651,9 @@ --> + + org.springframework.ai.chat.memory/**/*IT.java + org.springframework.ai.anthropic/**/*IT.java org.springframework.ai.azure.openai/**/*IT.java @@ -671,7 +677,6 @@ org.springframework.ai.vectorstore**/CosmosDB**IT.java org.springframework.ai.vectorstore.azure/**IT.java - org.springframework.ai.chat.memory/**/Cassandra**IT.java org.springframework.ai.vectorstore**/Cassandra**IT.java org.springframework.ai.chroma/**IT.java diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index fb8fee83f4e..73acadbb4a6 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -581,6 +581,12 @@ ${project.version} + + org.springframework.ai + spring-ai-chat-memory-jdbc-spring-boot-starter + ${project.version} + + diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc index 94ef91768a8..9ca3b90cd27 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc @@ -374,7 +374,7 @@ Refer to the xref:_retrieval_augmented_generation[Retrieval Augmented Generation The interface `ChatMemory` represents a storage for chat conversation history. It provides methods to add messages to a conversation, retrieve messages from a conversation, and clear the conversation history. -There are currently two implementations, `InMemoryChatMemory` and `CassandraChatMemory`, that provide storage for chat conversation history, in-memory and persisted with `time-to-live`, correspondingly. +There are currently three implementations, `InMemoryChatMemory`, `CassandraChatMemory`, and `JdbcChatMemory` that provide storage for chat conversation history, in-memory and persisted with `time-to-live`, correspondingly. To create a `CassandraChatMemory` with `time-to-live`: @@ -383,11 +383,18 @@ To create a `CassandraChatMemory` with `time-to-live`: CassandraChatMemory.create(CassandraChatMemoryConfig.builder().withTimeToLive(Duration.ofDays(1)).build()); ---- +To create a `JdbcChatMemory`: + +[source,java] +---- +JdbcChatMemory.create(JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build()); +---- + The following advisor implementations use the `ChatMemory` interface to advice the prompt with conversation history which differ in the details of how the memory is added to the prompt * `MessageChatMemoryAdvisor` : Memory is retrieved and added as a collection of messages to the prompt * `PromptChatMemoryAdvisor` : Memory is retrieved and added into the prompt's system text. -* `VectorStoreChatMemoryAdvisor` : The constructor `VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId, int chatHistoryWindowSize, int order)` This constructor allows you to: +* `VectorStoreChatMemoryAdvisor` : The constructor `VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId, int chatHistoryWindowSize, int order)` This constructor allows you to: . Specify the VectorStore instance used for managing and querying documents. . Set a default conversation ID to be used if none is provided in the context. diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 92ac01362c6..c23bb889a8c 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -412,6 +412,14 @@ true + + + org.springframework.ai + spring-ai-chat-memory-jdbc + ${project.parent.version} + true + + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfiguration.java new file mode 100644 index 00000000000..3442f44ad45 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfiguration.java @@ -0,0 +1,65 @@ +/* + * Copyright 2024-2024 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.autoconfigure.chat.memory.jdbc; + +import javax.sql.DataSource; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.chat.memory.JdbcChatMemory; +import org.springframework.ai.chat.memory.JdbcChatMemoryConfig; +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.condition.ConditionalOnProperty; +import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer; +import org.springframework.context.annotation.Bean; +import org.springframework.jdbc.core.JdbcTemplate; + +/** + * @author Jonathan Leijendekker + * @since 1.0.0 + */ +@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class) +@ConditionalOnClass({ JdbcChatMemory.class, DataSource.class, JdbcTemplate.class }) +@EnableConfigurationProperties(JdbcChatMemoryProperties.class) +public class JdbcChatMemoryAutoConfiguration { + + private static final Logger logger = LoggerFactory.getLogger(JdbcChatMemoryAutoConfiguration.class); + + @Bean + @ConditionalOnMissingBean + public JdbcChatMemory chatMemory(JdbcTemplate jdbcTemplate) { + var config = JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build(); + + return JdbcChatMemory.create(config); + } + + @Bean + @ConditionalOnMissingBean + @ConditionalOnProperty(value = "spring.ai.chat.memory.jdbc.initialize-schema", havingValue = "true", + matchIfMissing = true) + public DataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) { + logger.debug("Initializing JdbcChatMemory schema"); + + return new JdbcChatMemoryDataSourceScriptDatabaseInitializer(dataSource); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializer.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializer.java new file mode 100644 index 00000000000..de3b628faf2 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializer.java @@ -0,0 +1,35 @@ +package org.springframework.ai.autoconfigure.chat.memory.jdbc; + +import java.util.List; + +import javax.sql.DataSource; + +import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer; +import org.springframework.boot.jdbc.init.PlatformPlaceholderDatabaseDriverResolver; +import org.springframework.boot.sql.init.DatabaseInitializationMode; +import org.springframework.boot.sql.init.DatabaseInitializationSettings; + +class JdbcChatMemoryDataSourceScriptDatabaseInitializer extends DataSourceScriptDatabaseInitializer { + + private static final String SCHEMA_LOCATION = "classpath:org/springframework/ai/chat/memory/schema-@@platform@@.sql"; + + public JdbcChatMemoryDataSourceScriptDatabaseInitializer(DataSource dataSource) { + super(dataSource, getSettings(dataSource)); + } + + static DatabaseInitializationSettings getSettings(DataSource dataSource) { + var settings = new DatabaseInitializationSettings(); + settings.setSchemaLocations(resolveSchemaLocations(dataSource)); + settings.setMode(DatabaseInitializationMode.ALWAYS); + settings.setContinueOnError(true); + + return settings; + } + + static List resolveSchemaLocations(DataSource dataSource) { + var platformResolver = new PlatformPlaceholderDatabaseDriverResolver(); + + return platformResolver.resolveAll(dataSource, SCHEMA_LOCATION); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryProperties.java new file mode 100644 index 00000000000..f13aabe5b25 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryProperties.java @@ -0,0 +1,31 @@ +/* + * Copyright 2024-2024 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.autoconfigure.chat.memory.jdbc; + +import org.springframework.ai.autoconfigure.chat.memory.CommonChatMemoryProperties; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Jonathan Leijendekker + * @since 1.0.0 + */ +@ConfigurationProperties(JdbcChatMemoryProperties.CONFIG_PREFIX) +public class JdbcChatMemoryProperties extends CommonChatMemoryProperties { + + public static final String CONFIG_PREFIX = "spring.ai.chat.memory.jdbc"; + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index f3e5633efc0..2a7e2ae859d 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -66,3 +66,4 @@ org.springframework.ai.autoconfigure.minimax.MiniMaxAutoConfiguration org.springframework.ai.autoconfigure.vertexai.embedding.VertexAiEmbeddingAutoConfiguration org.springframework.ai.autoconfigure.chat.memory.cassandra.CassandraChatMemoryAutoConfiguration org.springframework.ai.autoconfigure.vectorstore.observation.VectorStoreObservationAutoConfiguration +org.springframework.ai.autoconfigure.chat.memory.jdbc.JdbcChatMemoryAutoConfiguration diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfigurationIT.java new file mode 100644 index 00000000000..ae3898f1f55 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryAutoConfigurationIT.java @@ -0,0 +1,102 @@ +/* + * Copyright 2024-2024 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.autoconfigure.chat.memory.jdbc; + +import java.util.List; +import java.util.UUID; + +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.ai.chat.memory.JdbcChatMemory; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonathan Leijendekker + */ +@Testcontainers +class JdbcChatMemoryAutoConfigurationIT { + + static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("postgres:17"); + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>(DEFAULT_IMAGE_NAME) + .withDatabaseName("chat_memory_auto_configuration_test") + .withUsername("postgres") + .withPassword("postgres"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(JdbcChatMemoryAutoConfiguration.class, + JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class)) + .withPropertyValues(String.format("spring.datasource.url=%s", postgresContainer.getJdbcUrl())) + .withPropertyValues(String.format("spring.datasource.username=%s", postgresContainer.getUsername())) + .withPropertyValues(String.format("spring.datasource.password=%s", postgresContainer.getPassword())) + .withPropertyValues("spring.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + @Test + void jdbcChatMemoryScriptDatabaseInitializer_shouldBeLoaded() { + this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true").run(context -> { + assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue(); + }); + } + + @Test + void jdbcChatMemoryScriptDatabaseInitializer_shouldNotBeLoaded() { + this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=false").run(context -> { + assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isFalse(); + }); + } + + @Test + void addGetAndClear_shouldAllExecute() { + this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true").run(context -> { + var chatMemory = context.getBean(JdbcChatMemory.class); + var conversationId = UUID.randomUUID().toString(); + var userMessage = new UserMessage("Message from the user"); + + chatMemory.add(conversationId, userMessage); + + assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).hasSize(1); + assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEqualTo(List.of(userMessage)); + + chatMemory.clear(conversationId); + + assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEmpty(); + + var multipleMessages = List.of(new UserMessage("Message from the user 1"), + new AssistantMessage("Message from the assistant 1")); + + chatMemory.add(conversationId, multipleMessages); + + assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).hasSize(multipleMessages.size()); + assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEqualTo(multipleMessages); + }); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializerTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializerTests.java new file mode 100644 index 00000000000..6bce964e07f --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryDataSourceScriptDatabaseInitializerTests.java @@ -0,0 +1,52 @@ +package org.springframework.ai.autoconfigure.chat.memory.jdbc; + +import javax.sql.DataSource; + +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonathan Leijendekker + */ +@Testcontainers +class JdbcChatMemoryDataSourceScriptDatabaseInitializerTests { + + static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("postgres:17"); + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>(DEFAULT_IMAGE_NAME) + .withDatabaseName("chat_memory_initializer_test") + .withUsername("postgres") + .withPassword("postgres"); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(JdbcChatMemoryAutoConfiguration.class, + JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class)) + .withPropertyValues(String.format("spring.datasource.url=%s", postgresContainer.getJdbcUrl())) + .withPropertyValues(String.format("spring.datasource.username=%s", postgresContainer.getUsername())) + .withPropertyValues(String.format("spring.datasource.password=%s", postgresContainer.getPassword())) + .withPropertyValues("spring.datasource.type=com.zaxxer.hikari.HikariDataSource"); + + @Test + void getSettings_shouldHaveSchemaLocations() { + this.contextRunner.run(context -> { + var dataSource = context.getBean(DataSource.class); + var settings = JdbcChatMemoryDataSourceScriptDatabaseInitializer.getSettings(dataSource); + + assertThat(settings.getSchemaLocations()) + .containsOnly("classpath:org/springframework/ai/chat/memory/schema-postgresql.sql"); + }); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryPropertiesTests.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryPropertiesTests.java new file mode 100644 index 00000000000..0c3dc4c91af --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/chat/memory/jdbc/JdbcChatMemoryPropertiesTests.java @@ -0,0 +1,43 @@ +/* + * Copyright 2024-2024 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.autoconfigure.chat.memory.jdbc; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonathan Leijendekker + */ +class JdbcChatMemoryPropertiesTests { + + @Test + void defaultValues() { + var props = new JdbcChatMemoryProperties(); + + assertThat(props.isInitializeSchema()).isTrue(); + } + + @Test + void customValues() { + var props = new JdbcChatMemoryProperties(); + props.setInitializeSchema(false); + + assertThat(props.isInitializeSchema()).isFalse(); + } + +} diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-chat-memory-jdbc/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-chat-memory-jdbc/pom.xml new file mode 100644 index 00000000000..08c373c0088 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-chat-memory-jdbc/pom.xml @@ -0,0 +1,58 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-chat-memory-jdbc-spring-boot-starter + jar + Spring AI Starter - Chat Memory JDBC + Spring AI Chat Memory JDBC Auto Configuration + 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.boot + spring-boot-starter + + + + org.springframework.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.ai + spring-ai-chat-memory-jdbc + ${project.parent.version} + + + +