Skip to content

Commit 1a35e8a

Browse files
TheovanKraaysobychacko
authored andcommitted
cosmos db chat memory repository
spring-javaformat:apply add docs and move checkstyle Signed-off-by: Theo van Kraay <[email protected]>
1 parent 54f5127 commit 1a35e8a

File tree

14 files changed

+1237
-0
lines changed

14 files changed

+1237
-0
lines changed
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!--
3+
~ Copyright 2023-2025 the original author or authors.
4+
~
5+
~ Licensed under the Apache License, Version 2.0 (the "License");
6+
~ you may not use this file except in compliance with the License.
7+
~ You may obtain a copy of the License at
8+
~
9+
~ https://www.apache.org/licenses/LICENSE-2.0
10+
~
11+
~ Unless required by applicable law or agreed to in writing, software
12+
~ distributed under the License is distributed on an "AS IS" BASIS,
13+
~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
~ See the License for the specific language governing permissions and
15+
~ limitations under the License.
16+
-->
17+
18+
<project xmlns="http://maven.apache.org/POM/4.0.0"
19+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
20+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
21+
<modelVersion>4.0.0</modelVersion>
22+
<parent>
23+
<groupId>org.springframework.ai</groupId>
24+
<artifactId>spring-ai-parent</artifactId>
25+
<version>1.1.0-SNAPSHOT</version>
26+
<relativePath>../../../../../../pom.xml</relativePath>
27+
</parent>
28+
29+
<artifactId>spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db</artifactId>
30+
<name>Spring AI Auto Configuration - Chat Memory Repository - CosmosDB</name>
31+
<description>Spring AI Auto Configuration for CosmosDB Chat Memory Repository</description>
32+
33+
<url>https://github.com/spring-projects/spring-ai</url>
34+
35+
<scm>
36+
<url>https://github.com/spring-projects/spring-ai</url>
37+
<connection>git://github.com/spring-projects/spring-ai.git</connection>
38+
<developerConnection>[email protected]:spring-projects/spring-ai.git</developerConnection>
39+
</scm>
40+
41+
<dependencies>
42+
<dependency>
43+
<groupId>org.springframework.ai</groupId>
44+
<artifactId>spring-ai-model-chat-memory-repository-cosmos-db</artifactId>
45+
<version>${project.version}</version>
46+
</dependency>
47+
48+
<dependency>
49+
<groupId>org.springframework.ai</groupId>
50+
<artifactId>spring-ai-autoconfigure-model-chat-memory</artifactId>
51+
<version>${project.parent.version}</version>
52+
</dependency>
53+
54+
<dependency>
55+
<groupId>org.springframework.boot</groupId>
56+
<artifactId>spring-boot-autoconfigure</artifactId>
57+
</dependency>
58+
59+
<dependency>
60+
<groupId>org.springframework.boot</groupId>
61+
<artifactId>spring-boot-configuration-processor</artifactId>
62+
<optional>true</optional>
63+
</dependency>
64+
65+
<dependency>
66+
<groupId>org.springframework.boot</groupId>
67+
<artifactId>spring-boot-autoconfigure-processor</artifactId>
68+
<optional>true</optional>
69+
</dependency>
70+
71+
<dependency>
72+
<groupId>com.azure</groupId>
73+
<artifactId>azure-spring-data-cosmos</artifactId>
74+
<version>${azure-cosmos.version}</version>
75+
<optional>true</optional>
76+
</dependency>
77+
78+
<dependency>
79+
<groupId>com.azure</groupId>
80+
<artifactId>azure-identity</artifactId>
81+
<version>${azure-identity.version}</version>
82+
</dependency>
83+
84+
<!-- Test dependencies -->
85+
<dependency>
86+
<groupId>org.springframework.boot</groupId>
87+
<artifactId>spring-boot-starter-test</artifactId>
88+
<scope>test</scope>
89+
</dependency>
90+
91+
<dependency>
92+
<groupId>org.springframework.ai</groupId>
93+
<artifactId>spring-ai-test</artifactId>
94+
<version>${project.version}</version>
95+
<scope>test</scope>
96+
</dependency>
97+
</dependencies>
98+
</project>
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
/*
2+
* Copyright 2023-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.repository.cosmosdb.autoconfigure;
18+
19+
import com.azure.cosmos.CosmosAsyncClient;
20+
import com.azure.cosmos.CosmosClientBuilder;
21+
import com.azure.identity.DefaultAzureCredentialBuilder;
22+
23+
import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepository;
24+
import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepositoryConfig;
25+
import org.springframework.boot.autoconfigure.AutoConfiguration;
26+
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
27+
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
28+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
29+
import org.springframework.boot.context.properties.EnableConfigurationProperties;
30+
import org.springframework.context.annotation.Bean;
31+
32+
/**
33+
* {@link AutoConfiguration Auto-configuration} for {@link CosmosDBChatMemoryRepository}.
34+
*
35+
* @author Theo van Kraay
36+
* @since 1.0.0
37+
*/
38+
@AutoConfiguration
39+
@ConditionalOnClass({ CosmosDBChatMemoryRepository.class, CosmosAsyncClient.class })
40+
@EnableConfigurationProperties(CosmosDBChatMemoryRepositoryProperties.class)
41+
public class CosmosDBChatMemoryRepositoryAutoConfiguration {
42+
43+
private final String agentSuffix = "SpringAI-CDBNoSQL-ChatMemoryRepository";
44+
45+
@Bean
46+
@ConditionalOnMissingBean
47+
@ConditionalOnProperty(prefix = "spring.ai.chat.memory.repository.cosmosdb", name = "endpoint")
48+
public CosmosAsyncClient cosmosClient(CosmosDBChatMemoryRepositoryProperties properties) {
49+
if (properties.getEndpoint() == null || properties.getEndpoint().isEmpty()) {
50+
throw new IllegalArgumentException(
51+
"Cosmos DB endpoint must be provided via spring.ai.chat.memory.repository.cosmosdb.endpoint property");
52+
}
53+
54+
String mode = properties.getConnectionMode();
55+
if (mode == null) {
56+
properties.setConnectionMode("gateway");
57+
}
58+
else if (!mode.equals("direct") && !mode.equals("gateway")) {
59+
throw new IllegalArgumentException("Connection mode must be either 'direct' or 'gateway'");
60+
}
61+
62+
CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(properties.getEndpoint())
63+
.userAgentSuffix(this.agentSuffix);
64+
65+
if (properties.getKey() == null || properties.getKey().isEmpty()) {
66+
builder.credential(new DefaultAzureCredentialBuilder().build());
67+
}
68+
else {
69+
builder.key(properties.getKey());
70+
}
71+
72+
return ("direct".equals(properties.getConnectionMode()) ? builder.directMode() : builder.gatewayMode())
73+
.buildAsyncClient();
74+
}
75+
76+
@Bean
77+
@ConditionalOnMissingBean
78+
public CosmosDBChatMemoryRepositoryConfig cosmosDBChatMemoryRepositoryConfig(
79+
CosmosDBChatMemoryRepositoryProperties properties, CosmosAsyncClient cosmosAsyncClient) {
80+
81+
return CosmosDBChatMemoryRepositoryConfig.builder()
82+
.withCosmosClient(cosmosAsyncClient)
83+
.withDatabaseName(properties.getDatabaseName())
84+
.withContainerName(properties.getContainerName())
85+
.withPartitionKeyPath(properties.getPartitionKeyPath())
86+
.build();
87+
}
88+
89+
@Bean
90+
@ConditionalOnMissingBean
91+
public CosmosDBChatMemoryRepository cosmosDBChatMemoryRepository(CosmosDBChatMemoryRepositoryConfig config) {
92+
return CosmosDBChatMemoryRepository.create(config);
93+
}
94+
95+
}
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
* Copyright 2023-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.repository.cosmosdb.autoconfigure;
18+
19+
import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepositoryConfig;
20+
import org.springframework.boot.context.properties.ConfigurationProperties;
21+
22+
/**
23+
* Configuration properties for CosmosDB chat memory.
24+
*
25+
* @author Theo van Kraay
26+
* @since 1.0.0
27+
*/
28+
@ConfigurationProperties(CosmosDBChatMemoryRepositoryProperties.CONFIG_PREFIX)
29+
public class CosmosDBChatMemoryRepositoryProperties {
30+
31+
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.cosmosdb";
32+
33+
private String endpoint;
34+
35+
private String key;
36+
37+
private String connectionMode = "gateway";
38+
39+
private String databaseName = CosmosDBChatMemoryRepositoryConfig.DEFAULT_DATABASE_NAME;
40+
41+
private String containerName = CosmosDBChatMemoryRepositoryConfig.DEFAULT_CONTAINER_NAME;
42+
43+
private String partitionKeyPath = CosmosDBChatMemoryRepositoryConfig.DEFAULT_PARTITION_KEY_PATH;
44+
45+
public String getEndpoint() {
46+
return this.endpoint;
47+
}
48+
49+
public void setEndpoint(String endpoint) {
50+
this.endpoint = endpoint;
51+
}
52+
53+
public String getKey() {
54+
return this.key;
55+
}
56+
57+
public void setKey(String key) {
58+
this.key = key;
59+
}
60+
61+
public String getConnectionMode() {
62+
return this.connectionMode;
63+
}
64+
65+
public void setConnectionMode(String connectionMode) {
66+
this.connectionMode = connectionMode;
67+
}
68+
69+
public String getDatabaseName() {
70+
return this.databaseName;
71+
}
72+
73+
public void setDatabaseName(String databaseName) {
74+
this.databaseName = databaseName;
75+
}
76+
77+
public String getContainerName() {
78+
return this.containerName;
79+
}
80+
81+
public void setContainerName(String containerName) {
82+
this.containerName = containerName;
83+
}
84+
85+
public String getPartitionKeyPath() {
86+
return this.partitionKeyPath;
87+
}
88+
89+
public void setPartitionKeyPath(String partitionKeyPath) {
90+
this.partitionKeyPath = partitionKeyPath;
91+
}
92+
93+
}
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
org.springframework.ai.model.chat.memory.repository.cosmosdb.autoconfigure.CosmosDBChatMemoryRepositoryAutoConfiguration

0 commit comments

Comments
 (0)