Enhance Neo4jChatMemoryRepository, expand integration tests, improve configuration, and update documentation

- Enhanced Neo4jChatMemoryRepository to correctly restore custom metadata for SystemMessage using SystemMessage.Builder.
- Refactored and clarified Neo4jChatMemoryRepository implementation code.
- Added comprehensive integration tests for Neo4jChatMemoryConfig and Neo4jChatMemoryRepository, including:
-- Index creation verification
-- Custom label support
-- Getter validation for all configuration properties
-- Tests for saving and retrieving SystemMessage metadata
-- Tests ensuring saveAll(conversationId, Collections.emptyList()) clears all messages and removes the conversation node
-- Tests for handling of messages with empty content and empty metadata
-- Improved overall test coverage for Neo4j persistence and configuration edge cases
- Fixed resource management bugs in test classes (ensured proper driver/session closure).
- Improved index creation logic in Neo4jChatMemoryConfig for reliability and logging.
- Updated documentation to include Neo4jChatMemoryRepository usage and configuration

Signed-off-by: enricorampazzo <[email protected]>
Signed-off-by: Mark Pollack <[email protected]>

Author: enricorampazzo

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryDataSourceScriptDatabaseInitializerPostgresqlTests.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2024-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.jdbc.autoconfigure;
+
+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 JdbcChatMemoryDataSourceScriptDatabaseInitializerPostgresqlTests {
+
+	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()),
+				String.format("spring.datasource.username=%s", postgresContainer.getUsername()),
+				String.format("spring.datasource.password=%s", postgresContainer.getPassword()));
+
+	@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/jdbc/schema-postgresql.sql");
+		});
+	}
+
+}

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-neo4j/src/main/java/org/springframework/ai/model/chat/memory/neo4j/autoconfigure/Neo4jChatMemoryAutoConfiguration.java

@@ -18,8 +18,8 @@
 
 import org.neo4j.driver.Driver;
 
-import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemory;
 import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
+import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepository;
 import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -29,19 +29,19 @@
 import org.springframework.context.annotation.Bean;
 
 /**
- * {@link AutoConfiguration Auto-configuration} for {@link Neo4jChatMemory}.
+ * {@link AutoConfiguration Auto-configuration} for {@link Neo4jChatMemoryRepository}.
  *
  * @author Enrico Rampazzo
  * @since 1.0.0
  */
 @AutoConfiguration(after = Neo4jAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
-@ConditionalOnClass({ Neo4jChatMemory.class, Driver.class })
+@ConditionalOnClass({ Neo4jChatMemoryRepository.class, Driver.class })
 @EnableConfigurationProperties(Neo4jChatMemoryProperties.class)
 public class Neo4jChatMemoryAutoConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean
-	public Neo4jChatMemory chatMemory(Neo4jChatMemoryProperties properties, Driver driver) {
+	public Neo4jChatMemoryRepository chatMemoryRepository(Neo4jChatMemoryProperties properties, Driver driver) {
 
 		var builder = Neo4jChatMemoryConfig.builder()
 			.withMediaLabel(properties.getMediaLabel())
@@ -52,7 +52,7 @@ public Neo4jChatMemory chatMemory(Neo4jChatMemoryProperties properties, Driver d
 			.withToolResponseLabel(properties.getToolResponseLabel())
 			.withDriver(driver);
 
-		return Neo4jChatMemory.create(builder.build());
+		return new Neo4jChatMemoryRepository(builder.build());
 	}
 
 }

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-neo4j/src/test/java/org/springframework/ai/model/chat/memory/neo4j/autoconfigure/Neo4jChatMemoryAutoConfigurationIT.java renamed to auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-neo4j/src/test/java/org/springframework/ai/model/chat/memory/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfigurationIT.java

@@ -23,12 +23,14 @@
 import java.util.UUID;
 
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.memory.ChatMemory;
+import org.springframework.ai.chat.memory.ChatMemoryRepository;
+import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryRepository;
 import org.testcontainers.containers.Neo4jContainer;
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 import org.testcontainers.utility.DockerImageName;
 
-import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemory;
 import org.springframework.ai.chat.memory.neo4j.Neo4jChatMemoryConfig;
 import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.Message;
@@ -51,7 +53,7 @@
  * @since 1.0.0
  */
 @Testcontainers
-class Neo4jChatMemoryAutoConfigurationIT {
+class Neo4jChatMemoryRepositoryAutoConfigurationIT {
 
 	static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse("neo4j");
 
@@ -67,31 +69,31 @@ class Neo4jChatMemoryAutoConfigurationIT {
 	@Test
 	void addAndGet() {
 		this.contextRunner.withPropertyValues("spring.neo4j.uri=" + neo4jContainer.getBoltUrl()).run(context -> {
-			Neo4jChatMemory memory = context.getBean(Neo4jChatMemory.class);
+			ChatMemoryRepository memory = context.getBean(ChatMemoryRepository.class);
 
 			String sessionId = UUID.randomUUID().toString();
-			assertThat(memory.get(sessionId, Integer.MAX_VALUE)).isEmpty();
+			assertThat(memory.findByConversationId(sessionId)).isEmpty();
 
 			UserMessage userMessage = new UserMessage("test question");
 
-			memory.add(sessionId, userMessage);
-			List<Message> messages = memory.get(sessionId, Integer.MAX_VALUE);
+			memory.saveAll(sessionId, List.of(userMessage));
+			List<Message> messages = memory.findByConversationId(sessionId);
 			assertThat(messages).hasSize(1);
 			assertThat(messages.get(0)).usingRecursiveAssertion().isEqualTo(userMessage);
 
-			memory.clear(sessionId);
-			assertThat(memory.get(sessionId, Integer.MAX_VALUE)).isEmpty();
+			memory.deleteByConversationId(sessionId);
+			assertThat(memory.findByConversationId(sessionId)).isEmpty();
 
 			AssistantMessage assistantMessage = new AssistantMessage("test answer", Map.of(),
 					List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")));
 
-			memory.add(sessionId, List.of(userMessage, assistantMessage));
-			messages = memory.get(sessionId, Integer.MAX_VALUE);
+			memory.saveAll(sessionId, List.of(userMessage, assistantMessage));
+			messages = memory.findByConversationId(sessionId);
 			assertThat(messages).hasSize(2);
-			assertThat(messages.get(1)).isEqualTo(userMessage);
+			assertThat(messages.get(0)).isEqualTo(userMessage);
 
-			assertThat(messages.get(0)).isEqualTo(assistantMessage);
-			memory.clear(sessionId);
+			assertThat(messages.get(1)).isEqualTo(assistantMessage);
+			memory.deleteByConversationId(sessionId);
 			MimeType textPlain = MimeType.valueOf("text/plain");
 			List<Media> media = List.of(
 					Media.builder()
@@ -102,28 +104,28 @@ void addAndGet() {
 						.build(),
 					Media.builder().data(URI.create("http://www.google.com")).mimeType(textPlain).build());
 			UserMessage userMessageWithMedia = UserMessage.builder().text("Message with media").media(media).build();
-			memory.add(sessionId, userMessageWithMedia);
+			memory.saveAll(sessionId, List.of(userMessageWithMedia));
 
-			messages = memory.get(sessionId, Integer.MAX_VALUE);
+			messages = memory.findByConversationId(sessionId);
 			assertThat(messages.size()).isEqualTo(1);
 			assertThat(messages.get(0)).isEqualTo(userMessageWithMedia);
 			assertThat(((UserMessage) messages.get(0)).getMedia()).hasSize(2);
 			assertThat(((UserMessage) messages.get(0)).getMedia()).usingRecursiveFieldByFieldElementComparator()
 				.isEqualTo(media);
-			memory.clear(sessionId);
+			memory.deleteByConversationId(sessionId);
 			ToolResponseMessage toolResponseMessage = new ToolResponseMessage(
 					List.of(new ToolResponse("id", "name", "responseData"),
 							new ToolResponse("id2", "name2", "responseData2")),
 					Map.of("id", "id", "metadataKey", "metadata"));
-			memory.add(sessionId, toolResponseMessage);
-			messages = memory.get(sessionId, Integer.MAX_VALUE);
+			memory.saveAll(sessionId, List.of(toolResponseMessage));
+			messages = memory.findByConversationId(sessionId);
 			assertThat(messages.size()).isEqualTo(1);
 			assertThat(messages.get(0)).isEqualTo(toolResponseMessage);
 
-			memory.clear(sessionId);
+			memory.deleteByConversationId(sessionId);
 			SystemMessage sm = new SystemMessage("this is a System message");
-			memory.add(sessionId, sm);
-			messages = memory.get(sessionId, Integer.MAX_VALUE);
+			memory.saveAll(sessionId, List.of(sm));
+			messages = memory.findByConversationId(sessionId);
 			assertThat(messages).hasSize(1);
 			assertThat(messages.get(0)).usingRecursiveAssertion().isEqualTo(sm);
 		});
@@ -148,7 +150,7 @@ void setCustomConfiguration() {
 					propertyBase.formatted("toolresponselabel", toolResponseLabel),
 					propertyBase.formatted("medialabel", mediaLabel))
 			.run(context -> {
-				Neo4jChatMemory chatMemory = context.getBean(Neo4jChatMemory.class);
+				Neo4jChatMemoryRepository chatMemory = context.getBean(Neo4jChatMemoryRepository.class);
 				Neo4jChatMemoryConfig config = chatMemory.getConfig();
 				assertThat(config.getMessageLabel()).isEqualTo(messageLabel);
 				assertThat(config.getMediaLabel()).isEqualTo(mediaLabel);

auto-configurations/models/embedding/observation/spring-ai-autoconfigure-model-embedding-observation/src/main/java/org/springframework/ai/model/embedding/observation/autoconfigure/EmbeddingObservationAutoConfiguration.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2023-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.model.embedding.observation.autoconfigure;
+
+import io.micrometer.core.instrument.MeterRegistry;
+
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.observation.EmbeddingModelMeterObservationHandler;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.boot.autoconfigure.AutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.context.annotation.Bean;
+
+/**
+ * Auto-configuration for Spring AI embedding model observations.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+@AutoConfiguration(
+		afterName = "org.springframework.boot.actuate.autoconfigure.observation.ObservationAutoConfiguration")
+@ConditionalOnClass(EmbeddingModel.class)
+public class EmbeddingObservationAutoConfiguration {
+
+	@Bean
+	@ConditionalOnMissingBean
+	@ConditionalOnBean(MeterRegistry.class)
+	EmbeddingModelMeterObservationHandler embeddingModelMeterObservationHandler(
+			ObjectProvider<MeterRegistry> meterRegistry) {
+		return new EmbeddingModelMeterObservationHandler(meterRegistry.getObject());
+	}
+
+}

auto-configurations/models/embedding/observation/spring-ai-autoconfigure-model-embedding-observation/src/test/java/org/springframework/ai/model/embedding/observation/autoconfigure/EmbeddingObservationAutoConfigurationTests.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2023-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.model.embedding.observation.autoconfigure;
+
+import io.micrometer.core.instrument.composite.CompositeMeterRegistry;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.embedding.observation.EmbeddingModelMeterObservationHandler;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests for {@link EmbeddingObservationAutoConfiguration}.
+ *
+ * @author Thomas Vitale
+ */
+class EmbeddingObservationAutoConfigurationTests {
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+		.withConfiguration(AutoConfigurations.of(EmbeddingObservationAutoConfiguration.class));
+
+	@Test
+	void meterObservationHandlerEnabled() {
+		this.contextRunner.withBean(CompositeMeterRegistry.class)
+			.run(context -> assertThat(context).hasSingleBean(EmbeddingModelMeterObservationHandler.class));
+	}
+
+	@Test
+	void meterObservationHandlerDisabled() {
+		this.contextRunner
+			.run(context -> assertThat(context).doesNotHaveBean(EmbeddingModelMeterObservationHandler.class));
+	}
+
+}

memory/spring-ai-model-chat-memory-neo4j/pom.xml

@@ -50,6 +50,7 @@
             <artifactId>spring-data-neo4j</artifactId>
         </dependency>
 
+        <!-- TESTING -->
         <dependency>
             <groupId>org.springframework.boot</groupId>
             <artifactId>spring-boot-starter-test</artifactId>
@@ -68,6 +69,12 @@
             <artifactId>spring-boot-testcontainers</artifactId>
             <scope>test</scope>
         </dependency>
+        
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>testcontainers</artifactId>
+            <scope>test</scope>
+        </dependency>
 
 		<dependency>
 			<groupId>org.neo4j.driver</groupId>
@@ -79,6 +86,12 @@
             <artifactId>neo4j</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>junit-jupiter</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>

memory/spring-ai-model-chat-memory-neo4j/src/main/java/org/springframework/ai/chat/memory/neo4j/Neo4jChatMemoryConfig.java

@@ -93,6 +93,33 @@ private Neo4jChatMemoryConfig(Builder builder) {
 		this.toolCallLabel = builder.toolCallLabel;
 		this.metadataLabel = builder.metadataLabel;
 		this.toolResponseLabel = builder.toolResponseLabel;
+		ensureIndexes();
+	}
+
+	/**
+	 * Ensures that indexes exist on conversationId for Session nodes and index for
+	 * Message nodes. This improves query performance for lookups and ordering.
+	 */
+	private void ensureIndexes() {
+		if (this.driver == null) {
+			logger.warn("Neo4j Driver is null, cannot ensure indexes.");
+			return;
+		}
+		try (var session = this.driver.session()) {
+			// Index for conversationId on Session nodes
+			String sessionIndexCypher = String.format(
+					"CREATE INDEX session_conversation_id_index IF NOT EXISTS FOR (n:%s) ON (n.conversationId)",
+					this.sessionLabel);
+			// Index for index on Message nodes
+			String messageIndexCypher = String
+				.format("CREATE INDEX message_index_index IF NOT EXISTS FOR (n:%s) ON (n.index)", this.messageLabel);
+			session.run(sessionIndexCypher);
+			session.run(messageIndexCypher);
+			logger.info("Ensured Neo4j indexes for conversationId and message index.");
+		}
+		catch (Exception e) {
+			logger.warn("Failed to ensure Neo4j indexes for chat memory: {}", e.getMessage());
+		}
 	}
 
 	public static Builder builder() {