diff --git a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java index 31b7e6274c..ae36f3d11b 100644 --- a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java +++ b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java @@ -46,6 +46,7 @@ * * @author Enrico Rampazzo * @author Michael J. Simons + * @author Soby Chacko * @since 1.0.0 */ @@ -138,30 +139,29 @@ public void saveAll(String conversationId, List messages) { @Override public void deleteByConversationId(String conversationId) { - // First delete all messages and related nodes String deleteMessagesStatement = """ - MATCH (s:%s {id:$conversationId})-[r:HAS_MESSAGE]->(m:%s) - OPTIONAL MATCH (m)-[:HAS_METADATA]->(metadata:%s) - OPTIONAL MATCH (m)-[:HAS_MEDIA]->(media:%s) - OPTIONAL MATCH (m)-[:HAS_TOOL_RESPONSE]-(tr:%s) - OPTIONAL MATCH (m)-[:HAS_TOOL_CALL]->(tc:%s) + MATCH (s:$($sessionLabel) {id:$conversationId})-[r:HAS_MESSAGE]->(m:$($messageLabel)) + OPTIONAL MATCH (m)-[:HAS_METADATA]->(metadata:$($metadataLabel)) + OPTIONAL MATCH (m)-[:HAS_MEDIA]->(media:$($mediaLabel)) + OPTIONAL MATCH (m)-[:HAS_TOOL_RESPONSE]-(tr:$($toolResponseLabel)) + OPTIONAL MATCH (m)-[:HAS_TOOL_CALL]->(tc:$($toolCallLabel)) DETACH DELETE m, metadata, media, tr, tc - """.formatted(this.config.getSessionLabel(), this.config.getMessageLabel(), - this.config.getMetadataLabel(), this.config.getMediaLabel(), this.config.getToolResponseLabel(), - this.config.getToolCallLabel()); + """; - // Then delete the conversation node itself String deleteConversationStatement = """ - MATCH (s:%s {id:$conversationId}) + MATCH (s:$($sessionLabel) {id:$conversationId}) DETACH DELETE s - """.formatted(this.config.getSessionLabel()); + """; + + Map params = Map.of("conversationId", conversationId, "sessionLabel", + this.config.getSessionLabel(), "messageLabel", this.config.getMessageLabel(), "metadataLabel", + this.config.getMetadataLabel(), "mediaLabel", this.config.getMediaLabel(), "toolResponseLabel", + this.config.getToolResponseLabel(), "toolCallLabel", this.config.getToolCallLabel()); try (Session s = this.config.getDriver().session()) { try (Transaction t = s.beginTransaction()) { - // First delete messages - t.run(deleteMessagesStatement, Map.of("conversationId", conversationId)); - // Then delete the conversation node - t.run(deleteConversationStatement, Map.of("conversationId", conversationId)); + t.run(deleteMessagesStatement, params); + t.run(deleteConversationStatement, params); t.commit(); } } diff --git a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryConfig.java b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryConfig.java index 3086111f89..3a91097329 100644 --- a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryConfig.java +++ b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryConfig.java @@ -16,6 +16,8 @@ package org.springframework.ai.chat.memory.repository.neo4j; +import java.util.regex.Pattern; + import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.jspecify.annotations.Nullable; @@ -27,6 +29,7 @@ * Configuration for the Neo4j Chat Memory store. * * @author Enrico Rampazzo + * @author Soby Chacko */ public final class Neo4jChatMemoryRepositoryConfig { @@ -46,6 +49,8 @@ public final class Neo4jChatMemoryRepositoryConfig { private static final Log logger = LogFactory.getLog(Neo4jChatMemoryRepositoryConfig.class); + private static final Pattern SAFE_LABEL = Pattern.compile("[\\p{Alpha}_][\\p{Alnum}_]*"); + private final Driver driver; private final String sessionLabel; @@ -97,9 +102,20 @@ private Neo4jChatMemoryRepositoryConfig(Builder builder) { this.toolCallLabel = builder.toolCallLabel; this.metadataLabel = builder.metadataLabel; this.toolResponseLabel = builder.toolResponseLabel; + validateLabels(); ensureIndexes(); } + private void validateLabels() { + for (String label : new String[] { this.sessionLabel, this.messageLabel, this.metadataLabel, this.mediaLabel, + this.toolCallLabel, this.toolResponseLabel }) { + if (!SAFE_LABEL.matcher(label).matches()) { + throw new IllegalArgumentException("Invalid Neo4j node label: '" + label + + "'. Labels must start with a letter or underscore and contain only letters, digits, or underscores."); + } + } + } + /** * Ensures that indexes exist on conversationId for Session nodes and index for * Message nodes. This improves query performance for lookups and ordering. diff --git a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4JChatMemoryRepositoryConfigIT.java b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4JChatMemoryRepositoryConfigIT.java index 5fa46ca6ea..5fa4c6ba93 100644 --- a/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4JChatMemoryRepositoryConfigIT.java +++ b/memory-repositories/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4JChatMemoryRepositoryConfigIT.java @@ -16,9 +16,15 @@ package org.springframework.ai.chat.memory.repository.neo4j; +import java.util.function.BiFunction; +import java.util.stream.Stream; + import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import org.neo4j.driver.Driver; import org.neo4j.driver.GraphDatabase; import org.neo4j.driver.Result; @@ -28,7 +34,13 @@ import org.testcontainers.junit.jupiter.Testcontainers; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +/** + * Integration tests for {@link Neo4jChatMemoryRepositoryConfig}. + * + * @author Soby Chacko + */ @Testcontainers class Neo4JChatMemoryRepositoryConfigIT { @@ -87,6 +99,34 @@ void builderShouldSetCustomLabels() { assertThat(config.getMessageLabel()).isEqualTo(customMessageLabel); } + @ParameterizedTest + @ValueSource(strings = { "1Session", "Session Label", "Session-Label", "Session`", "Session)" }) + void shouldRejectMalformedSessionLabel(String invalidLabel) { + assertThatThrownBy(() -> Neo4jChatMemoryRepositoryConfig.builder() + .withDriver(driver) + .withSessionLabel(invalidLabel) + .build()).isInstanceOf(IllegalArgumentException.class).hasMessageContaining("Invalid Neo4j node label"); + } + + static Stream> allLabelSetters() { + return Stream.of(Neo4jChatMemoryRepositoryConfig.Builder::withSessionLabel, + Neo4jChatMemoryRepositoryConfig.Builder::withMessageLabel, + Neo4jChatMemoryRepositoryConfig.Builder::withMetadataLabel, + Neo4jChatMemoryRepositoryConfig.Builder::withMediaLabel, + Neo4jChatMemoryRepositoryConfig.Builder::withToolCallLabel, + Neo4jChatMemoryRepositoryConfig.Builder::withToolResponseLabel); + } + + @ParameterizedTest + @MethodSource("allLabelSetters") + void shouldRejectSpecialCharactersInAnyLabel( + BiFunction setter) { + assertThatThrownBy( + () -> setter.apply(Neo4jChatMemoryRepositoryConfig.builder().withDriver(driver), "Bad`) Label").build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Invalid Neo4j node label"); + } + @Test void gettersShouldReturnConfiguredValues() { Neo4jChatMemoryRepositoryConfig config = Neo4jChatMemoryRepositoryConfig.builder() diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat-memory.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat-memory.adoc index 828ce772bf..5c86d2a0fe 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat-memory.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat-memory.adoc @@ -364,6 +364,9 @@ ChatMemory chatMemory = MessageWindowChatMemory.builder() | `spring.ai.chat.memory.repository.neo4j.media-label` | The label for the nodes that store media associated with a message | `Media` |=== +NOTE: Label values must be valid Neo4j identifiers: they must start with a letter or underscore and contain only letters, digits, or underscores. +Supplying a value that does not meet these requirements will cause an `IllegalArgumentException` to be thrown at application startup. + ==== Index Initialization The Neo4j repository will automatically ensure that indexes are created for conversation IDs and message indices to optimize performance. If you use custom labels, indexes will be created for those labels as well. No schema initialization is required, but you should ensure your Neo4j instance is accessible to your application.