diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfiguration.java b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfiguration.java index 970cb6be91c..5694456cbf0 100644 --- a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfiguration.java +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfiguration.java @@ -19,7 +19,6 @@ import org.neo4j.driver.Driver; import org.springframework.ai.chat.memory.repository.neo4j.Neo4jChatMemoryRepository; -import org.springframework.ai.chat.memory.repository.neo4j.Neo4jChatMemoryRepositoryConfig; import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; @@ -43,17 +42,15 @@ public class Neo4jChatMemoryRepositoryAutoConfiguration { @ConditionalOnMissingBean public Neo4jChatMemoryRepository neo4jChatMemoryRepository(Neo4jChatMemoryRepositoryProperties properties, Driver driver) { - - var builder = Neo4jChatMemoryRepositoryConfig.builder() - .withMediaLabel(properties.getMediaLabel()) - .withMessageLabel(properties.getMessageLabel()) - .withMetadataLabel(properties.getMetadataLabel()) - .withSessionLabel(properties.getSessionLabel()) - .withToolCallLabel(properties.getToolCallLabel()) - .withToolResponseLabel(properties.getToolResponseLabel()) - .withDriver(driver); - - return new Neo4jChatMemoryRepository(builder.build()); + return Neo4jChatMemoryRepository.builder() + .driver(driver) + .mediaLabel(properties.getMediaLabel()) + .messageLabel(properties.getMessageLabel()) + .metadataLabel(properties.getMetadataLabel()) + .sessionLabel(properties.getSessionLabel()) + .toolCallLabel(properties.getToolCallLabel()) + .toolResponseLabel(properties.getToolResponseLabel()) + .build(); } } diff --git a/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java index 21cdd80a54e..a26ec1e9f58 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java @@ -25,6 +25,7 @@ import java.util.UUID; import java.util.stream.Collectors; +import org.neo4j.driver.Driver; import org.neo4j.driver.Session; import org.neo4j.driver.Transaction; import org.neo4j.driver.TransactionContext; @@ -38,6 +39,7 @@ import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.content.Media; import org.springframework.ai.content.MediaContent; +import org.springframework.util.Assert; import org.springframework.util.MimeType; /** @@ -52,7 +54,7 @@ public final class Neo4jChatMemoryRepository implements ChatMemoryRepository { private final Neo4jChatMemoryRepositoryConfig config; - public Neo4jChatMemoryRepository(Neo4jChatMemoryRepositoryConfig config) { + private Neo4jChatMemoryRepository(Neo4jChatMemoryRepositoryConfig config) { this.config = config; } @@ -326,4 +328,56 @@ private List> convertMediaToMap(List media) { return mediaMaps; } + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + + private final Neo4jChatMemoryRepositoryConfig.Builder builder = Neo4jChatMemoryRepositoryConfig.builder(); + + private Builder() { + } + + public Builder driver(Driver driver) { + this.builder.withDriver(driver); + return this; + } + + public Builder sessionLabel(String sessionLabel) { + this.builder.withSessionLabel(sessionLabel); + return this; + } + + public Builder toolCallLabel(String toolCallLabel) { + this.builder.withToolCallLabel(toolCallLabel); + return this; + } + + public Builder metadataLabel(String metadataLabel) { + this.builder.withMetadataLabel(metadataLabel); + return this; + } + + public Builder messageLabel(String messageLabel) { + this.builder.withMessageLabel(messageLabel); + return this; + } + + public Builder toolResponseLabel(String toolResponseLabel) { + this.builder.withToolResponseLabel(toolResponseLabel); + return this; + } + + public Builder mediaLabel(String mediaLabel) { + this.builder.withMediaLabel(mediaLabel); + return this; + } + + public Neo4jChatMemoryRepository build() { + Assert.notNull(this.builder.getDriver(), "Driver cannot be null"); + return new Neo4jChatMemoryRepository(this.builder.build()); + } + } + } diff --git a/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryIT.java b/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryIT.java index 83ff42a71ae..cdbd8ffdb17 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryIT.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryIT.java @@ -67,7 +67,7 @@ class Neo4jChatMemoryRepositoryIT { .withoutAuthentication() .withExposedPorts(7474, 7687); - private ChatMemoryRepository chatMemoryRepository; + private Neo4jChatMemoryRepository chatMemoryRepository; private Driver driver; @@ -76,8 +76,8 @@ class Neo4jChatMemoryRepositoryIT { @BeforeEach void setUp() { this.driver = Neo4jDriverFactory.create(neo4jContainer.getBoltUrl()); - this.config = Neo4jChatMemoryRepositoryConfig.builder().withDriver(this.driver).build(); - this.chatMemoryRepository = new Neo4jChatMemoryRepository(this.config); + this.chatMemoryRepository = Neo4jChatMemoryRepository.builder().driver(driver).build(); + this.config = chatMemoryRepository.getConfig(); } @AfterEach 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 435751c8b56..e8b9fe8948a 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 @@ -222,7 +222,7 @@ If you'd rather create the `CassandraChatMemoryRepository` manually, you can do [source,java] ---- ChatMemoryRepository chatMemoryRepository = CassandraChatMemoryRepository - .create(CassandraChatMemoryConfig.builder().withCqlSession(cqlSession)); + .create(CassandraChatMemoryRepositoryConfig.builder().withCqlSession(cqlSession).build()); ChatMemory chatMemory = MessageWindowChatMemory.builder() .chatMemoryRepository(chatMemoryRepository)