Skip to content

Commit f3d3324

Browse files
deejay1deejay1
committed
Create Mongo Chat Memory repository
Signed-off-by: Łukasz Jernaś <[email protected]>
1 parent ffd8f1f commit f3d3324

File tree

16 files changed

+924
-0
lines changed

16 files changed

+924
-0
lines changed

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/pom.xml

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
3+
xmlns="http://maven.apache.org/POM/4.0.0"
4+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
5+
<modelVersion>4.0.0</modelVersion>
6+
<parent>
7+
<groupId>org.springframework.ai</groupId>
8+
<artifactId>spring-ai-parent</artifactId>
9+
<version>1.0.0-SNAPSHOT</version>
10+
<relativePath>../../../../../pom.xml</relativePath>
11+
</parent>
12+
<artifactId>spring-ai-autoconfigure-model-chat-memory-mongodb</artifactId>
13+
<packaging>jar</packaging>
14+
<name>Spring AI MongoDB Chat Memory Auto Configuration</name>
15+
<description>Spring AI MongoDB Chat Memory Auto Configuration</description>
16+
<url>https://github.com/spring-projects/spring-ai</url>
17+
18+
<scm>
19+
<url>https://github.com/spring-projects/spring-ai</url>
20+
<connection>git://github.com/spring-projects/spring-ai.git</connection>
21+
<developerConnection>[email protected]:spring-projects/spring-ai.git</developerConnection>
22+
</scm>
23+
24+
25+
<dependencies>
26+
27+
<dependency>
28+
<groupId>org.springframework.ai</groupId>
29+
<artifactId>spring-ai-model-chat-memory-mongodb</artifactId>
30+
<version>${project.parent.version}</version>
31+
</dependency>
32+
33+
<dependency>
34+
<groupId>org.springframework.ai</groupId>
35+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
36+
<version>${project.parent.version}</version>
37+
</dependency>
38+
39+
<!-- Boot dependencies -->
40+
<dependency>
41+
<groupId>org.springframework.boot</groupId>
42+
<artifactId>spring-boot-starter</artifactId>
43+
</dependency>
44+
45+
<dependency>
46+
<groupId>org.springframework.boot</groupId>
47+
<artifactId>spring-boot-starter-data-mongodb</artifactId>
48+
</dependency>
49+
50+
<dependency>
51+
<groupId>org.springframework.boot</groupId>
52+
<artifactId>spring-boot-configuration-processor</artifactId>
53+
<optional>true</optional>
54+
</dependency>
55+
56+
<dependency>
57+
<groupId>org.springframework.boot</groupId>
58+
<artifactId>spring-boot-autoconfigure-processor</artifactId>
59+
<optional>true</optional>
60+
</dependency>
61+
62+
<!-- Test dependencies -->
63+
<dependency>
64+
<groupId>org.springframework.ai</groupId>
65+
<artifactId>spring-ai-test</artifactId>
66+
<version>${project.parent.version}</version>
67+
<scope>test</scope>
68+
</dependency>
69+
70+
<dependency>
71+
<groupId>org.springframework.ai</groupId>
72+
<artifactId>spring-ai-autoconfigure-model-chat-client</artifactId>
73+
<version>${project.parent.version}</version>
74+
<scope>test</scope>
75+
</dependency>
76+
77+
<dependency>
78+
<groupId>org.springframework.ai</groupId>
79+
<artifactId>spring-ai-openai</artifactId>
80+
<version>${project.parent.version}</version>
81+
<scope>test</scope>
82+
</dependency>
83+
84+
<dependency>
85+
<groupId>org.springframework.boot</groupId>
86+
<artifactId>spring-boot-starter-test</artifactId>
87+
<scope>test</scope>
88+
</dependency>
89+
90+
<dependency>
91+
<groupId>org.springframework.boot</groupId>
92+
<artifactId>spring-boot-testcontainers</artifactId>
93+
<scope>test</scope>
94+
</dependency>
95+
96+
<dependency>
97+
<groupId>org.testcontainers</groupId>
98+
<artifactId>junit-jupiter</artifactId>
99+
<scope>test</scope>
100+
</dependency>
101+
102+
<dependency>
103+
<groupId>org.testcontainers</groupId>
104+
<artifactId>mongodb</artifactId>
105+
<scope>test</scope>
106+
</dependency>
107+
108+
<dependency>
109+
<groupId>org.mockito</groupId>
110+
<artifactId>mockito-core</artifactId>
111+
<scope>test</scope>
112+
</dependency>
113+
</dependencies>
114+
</project>

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/src/main/java/org/springframework/ai/model/chat/memory/mongo/autoconfigure/MongoChatMemoryAutoConfiguration.java

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.mongo.autoconfigure;
18+
19+
import org.springframework.ai.chat.memory.mongo.MongoChatMemoryRepository;
20+
import org.springframework.ai.model.chat.memory.autoconfigure.ChatMemoryAutoConfiguration;
21+
import org.springframework.boot.autoconfigure.AutoConfiguration;
22+
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
23+
import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration;
24+
import org.springframework.boot.context.properties.EnableConfigurationProperties;
25+
import org.springframework.context.annotation.Bean;
26+
import org.springframework.data.mongodb.core.MongoTemplate;
27+
28+
/**
29+
* Spring Boot autoconfiguration for {@link MongoChatMemoryRepository}.
30+
*
31+
* @author Łukasz Jernaś
32+
* @since 1.0.0
33+
*/
34+
@AutoConfiguration(after = MongoDataAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
35+
@EnableConfigurationProperties(MongoChatMemoryProperties.class)
36+
public class MongoChatMemoryAutoConfiguration {
37+
38+
@Bean
39+
@ConditionalOnMissingBean
40+
MongoChatMemoryRepository chatMemoryRepository(MongoTemplate mongoTemplate) {
41+
return MongoChatMemoryRepository.builder().mongoTemplate(mongoTemplate).build();
42+
}
43+
44+
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/src/main/java/org/springframework/ai/model/chat/memory/mongo/autoconfigure/MongoChatMemoryIndexCreator.java

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.mongo.autoconfigure;
18+
19+
import org.slf4j.Logger;
20+
import org.slf4j.LoggerFactory;
21+
import org.springframework.ai.chat.memory.mongo.Conversation;
22+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
23+
import org.springframework.context.event.ContextRefreshedEvent;
24+
import org.springframework.context.event.EventListener;
25+
import org.springframework.data.domain.Sort;
26+
import org.springframework.data.mongodb.core.MongoTemplate;
27+
import org.springframework.data.mongodb.core.index.Index;
28+
import org.springframework.stereotype.Component;
29+
30+
/**
31+
* Class responsible for creating MongoDB proper indices for the ChatMemory. Creates a
32+
* main index on the conversationId and timestamp fields, and a TTL index on the timestamp
33+
* field if the TTL is set in properties.
34+
*
35+
* @author Łukasz Jernaś
36+
* @see MongoChatMemoryProperties
37+
* @since 1.0.0
38+
*/
39+
@Component
40+
@ConditionalOnProperty(value = "spring.ai.chat.memory.mongodb.create-indices", havingValue = "true")
41+
public class MongoChatMemoryIndexCreator {
42+
43+
private static final Logger logger = LoggerFactory.getLogger(MongoChatMemoryIndexCreator.class);
44+
45+
private final MongoTemplate mongoTemplate;
46+
47+
private final MongoChatMemoryProperties mongoChatMemoryProperties;
48+
49+
public MongoChatMemoryIndexCreator(MongoTemplate mongoTemplate,
50+
MongoChatMemoryProperties mongoChatMemoryProperties) {
51+
this.mongoTemplate = mongoTemplate;
52+
this.mongoChatMemoryProperties = mongoChatMemoryProperties;
53+
}
54+
55+
@EventListener(ContextRefreshedEvent.class)
56+
public void initIndicesAfterStartup() {
57+
logger.info("Creating MongoDB indices for ChatMemory");
58+
// Create a main index
59+
mongoTemplate.indexOps(Conversation.class)
60+
.ensureIndex(new Index().on("conversationId", Sort.Direction.ASC).on("timestamp", Sort.Direction.DESC));
61+
62+
createOrUpdateTtlIndex();
63+
}
64+
65+
private void createOrUpdateTtlIndex() {
66+
if (!this.mongoChatMemoryProperties.getTtl().isZero()) {
67+
// Check for existing TTL index
68+
mongoTemplate.indexOps(Conversation.class).getIndexInfo().forEach(idx -> {
69+
if (idx.getExpireAfter().isPresent()
70+
&& !idx.getExpireAfter().get().equals(this.mongoChatMemoryProperties.getTtl())) {
71+
logger.warn("Dropping existing TTL index, because TTL is different");
72+
mongoTemplate.indexOps(Conversation.class).dropIndex(idx.getName());
73+
}
74+
});
75+
mongoTemplate.indexOps(Conversation.class)
76+
.ensureIndex(new Index().on("timestamp", Sort.Direction.ASC)
77+
.expire(this.mongoChatMemoryProperties.getTtl()));
78+
}
79+
}
80+
81+
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/src/main/java/org/springframework/ai/model/chat/memory/mongo/autoconfigure/MongoChatMemoryProperties.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.mongo.autoconfigure;
18+
19+
import org.springframework.boot.context.properties.ConfigurationProperties;
20+
21+
import java.time.Duration;
22+
23+
/**
24+
* @author Łukasz Jernaś
25+
* @since 1.0.0
26+
*/
27+
@ConfigurationProperties(MongoChatMemoryProperties.CONFIG_PREFIX)
28+
public class MongoChatMemoryProperties {
29+
30+
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.mongodb";
31+
32+
/**
33+
* If the indexes should be automatically created on app startup. Note: Changing the
34+
* TTL value will drop the TTL index and recreate it.
35+
*/
36+
private boolean createIndices = false;
37+
38+
/**
39+
* The time to live (TTL) for the conversation documents in the database. The default
40+
* value is 0, which means that the documents will not expire.
41+
*/
42+
private Duration ttl = Duration.ZERO;
43+
44+
public Duration getTtl() {
45+
return ttl;
46+
}
47+
48+
public void setTtl(Duration ttl) {
49+
this.ttl = ttl;
50+
}
51+
52+
public boolean isCreateIndices() {
53+
return createIndices;
54+
}
55+
56+
public void setCreateIndices(boolean createIndices) {
57+
this.createIndices = createIndices;
58+
}
59+
60+
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#
2+
# Copyright 2024-2025 the original author or authors.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# https://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
org.springframework.ai.model.chat.memory.mongo.autoconfigure.MongoChatMemoryAutoConfiguration

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-mongodb/src/test/java/org/springframework/ai/model/chat/memory/mongo/autoconfigure/MongoChatMemoryAutoConfigurationIT.java

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.model.chat.memory.mongo.autoconfigure;
18+
19+
import org.junit.jupiter.api.Test;
20+
import org.springframework.ai.chat.memory.ChatMemoryRepository;
21+
import org.springframework.ai.chat.messages.SystemMessage;
22+
import org.springframework.beans.factory.annotation.Autowired;
23+
import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration;
24+
import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration;
25+
import org.springframework.boot.test.context.SpringBootTest;
26+
import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
27+
import org.springframework.data.mongodb.core.MongoTemplate;
28+
import org.springframework.test.context.TestPropertySource;
29+
import org.testcontainers.containers.MongoDBContainer;
30+
import org.testcontainers.junit.jupiter.Container;
31+
32+
import java.util.List;
33+
import java.util.UUID;
34+
35+
import static org.assertj.core.api.Assertions.assertThat;
36+
37+
@SpringBootTest(classes = { MongoAutoConfiguration.class, MongoDataAutoConfiguration.class,
38+
MongoChatMemoryAutoConfiguration.class })
39+
@TestPropertySource(properties = "spring.data.mongodb.uri=spring.data.mongodb.uri=%s/ai_test")
40+
class MongoChatMemoryAutoConfigurationIT {
41+
42+
@Autowired
43+
private ChatMemoryRepository chatMemoryRepository;
44+
45+
@Autowired
46+
private MongoTemplate mongoTemplate;
47+
48+
@Container
49+
@ServiceConnection
50+
static MongoDBContainer mongoDbContainer = new MongoDBContainer("mongo:8.0.6");
51+
52+
@Test
53+
void allMethodsShouldExecute() {
54+
var conversationId = UUID.randomUUID().toString();
55+
var systemMessage = new SystemMessage("Some system message");
56+
57+
chatMemoryRepository.saveAll(conversationId, List.of(systemMessage));
58+
59+
assertThat(chatMemoryRepository.findConversationIds().contains(conversationId)).isTrue();
60+
61+
assertThat(chatMemoryRepository.findByConversationId(conversationId).size()).isEqualTo(1);
62+
63+
chatMemoryRepository.deleteByConversationId(conversationId);
64+
65+
assertThat(chatMemoryRepository.findByConversationId(conversationId).size()).isEqualTo(0);
66+
67+
}
68+
69+
}

0 commit comments

Comments
 (0)