diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/pom.xml b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/pom.xml new file mode 100644 index 00000000000..535fd9f44ed --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/pom.xml @@ -0,0 +1,98 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 1.1.0-SNAPSHOT + ../../../../../../pom.xml + + + spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db + Spring AI Auto Configuration - Chat Memory Repository - CosmosDB + Spring AI Auto Configuration for CosmosDB Chat Memory Repository + + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + org.springframework.ai + spring-ai-model-chat-memory-repository-cosmos-db + ${project.version} + + + + org.springframework.ai + spring-ai-autoconfigure-model-chat-memory + ${project.parent.version} + + + + org.springframework.boot + spring-boot-autoconfigure + + + + org.springframework.boot + spring-boot-configuration-processor + true + + + + org.springframework.boot + spring-boot-autoconfigure-processor + true + + + + com.azure + azure-spring-data-cosmos + ${azure-cosmos.version} + true + + + + com.azure + azure-identity + ${azure-identity.version} + + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.springframework.ai + spring-ai-test + ${project.version} + test + + + diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfiguration.java b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfiguration.java new file mode 100644 index 00000000000..1a3ca9e9711 --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfiguration.java @@ -0,0 +1,95 @@ +/* + * Copyright 2023-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.repository.cosmosdb.autoconfigure; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.identity.DefaultAzureCredentialBuilder; + +import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepository; +import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepositoryConfig; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; + +/** + * {@link AutoConfiguration Auto-configuration} for {@link CosmosDBChatMemoryRepository}. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +@AutoConfiguration +@ConditionalOnClass({ CosmosDBChatMemoryRepository.class, CosmosAsyncClient.class }) +@EnableConfigurationProperties(CosmosDBChatMemoryRepositoryProperties.class) +public class CosmosDBChatMemoryRepositoryAutoConfiguration { + + private final String agentSuffix = "SpringAI-CDBNoSQL-ChatMemoryRepository"; + + @Bean + @ConditionalOnMissingBean + @ConditionalOnProperty(prefix = "spring.ai.chat.memory.repository.cosmosdb", name = "endpoint") + public CosmosAsyncClient cosmosClient(CosmosDBChatMemoryRepositoryProperties properties) { + if (properties.getEndpoint() == null || properties.getEndpoint().isEmpty()) { + throw new IllegalArgumentException( + "Cosmos DB endpoint must be provided via spring.ai.chat.memory.repository.cosmosdb.endpoint property"); + } + + String mode = properties.getConnectionMode(); + if (mode == null) { + properties.setConnectionMode("gateway"); + } + else if (!mode.equals("direct") && !mode.equals("gateway")) { + throw new IllegalArgumentException("Connection mode must be either 'direct' or 'gateway'"); + } + + CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(properties.getEndpoint()) + .userAgentSuffix(this.agentSuffix); + + if (properties.getKey() == null || properties.getKey().isEmpty()) { + builder.credential(new DefaultAzureCredentialBuilder().build()); + } + else { + builder.key(properties.getKey()); + } + + return ("direct".equals(properties.getConnectionMode()) ? builder.directMode() : builder.gatewayMode()) + .buildAsyncClient(); + } + + @Bean + @ConditionalOnMissingBean + public CosmosDBChatMemoryRepositoryConfig cosmosDBChatMemoryRepositoryConfig( + CosmosDBChatMemoryRepositoryProperties properties, CosmosAsyncClient cosmosAsyncClient) { + + return CosmosDBChatMemoryRepositoryConfig.builder() + .withCosmosClient(cosmosAsyncClient) + .withDatabaseName(properties.getDatabaseName()) + .withContainerName(properties.getContainerName()) + .withPartitionKeyPath(properties.getPartitionKeyPath()) + .build(); + } + + @Bean + @ConditionalOnMissingBean + public CosmosDBChatMemoryRepository cosmosDBChatMemoryRepository(CosmosDBChatMemoryRepositoryConfig config) { + return CosmosDBChatMemoryRepository.create(config); + } + +} diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryProperties.java b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryProperties.java new file mode 100644 index 00000000000..56e813d263b --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryProperties.java @@ -0,0 +1,93 @@ +/* + * Copyright 2023-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.repository.cosmosdb.autoconfigure; + +import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepositoryConfig; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * Configuration properties for CosmosDB chat memory. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +@ConfigurationProperties(CosmosDBChatMemoryRepositoryProperties.CONFIG_PREFIX) +public class CosmosDBChatMemoryRepositoryProperties { + + public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.cosmosdb"; + + private String endpoint; + + private String key; + + private String connectionMode = "gateway"; + + private String databaseName = CosmosDBChatMemoryRepositoryConfig.DEFAULT_DATABASE_NAME; + + private String containerName = CosmosDBChatMemoryRepositoryConfig.DEFAULT_CONTAINER_NAME; + + private String partitionKeyPath = CosmosDBChatMemoryRepositoryConfig.DEFAULT_PARTITION_KEY_PATH; + + public String getEndpoint() { + return this.endpoint; + } + + public void setEndpoint(String endpoint) { + this.endpoint = endpoint; + } + + public String getKey() { + return this.key; + } + + public void setKey(String key) { + this.key = key; + } + + public String getConnectionMode() { + return this.connectionMode; + } + + public void setConnectionMode(String connectionMode) { + this.connectionMode = connectionMode; + } + + public String getDatabaseName() { + return this.databaseName; + } + + public void setDatabaseName(String databaseName) { + this.databaseName = databaseName; + } + + public String getContainerName() { + return this.containerName; + } + + public void setContainerName(String containerName) { + this.containerName = containerName; + } + + public String getPartitionKeyPath() { + return this.partitionKeyPath; + } + + public void setPartitionKeyPath(String partitionKeyPath) { + this.partitionKeyPath = partitionKeyPath; + } + +} diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports new file mode 100644 index 00000000000..411c013dcb8 --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -0,0 +1 @@ +org.springframework.ai.model.chat.memory.repository.cosmosdb.autoconfigure.CosmosDBChatMemoryRepositoryAutoConfiguration diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfigurationIT.java b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfigurationIT.java new file mode 100644 index 00000000000..5b3d6f1d9da --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryAutoConfigurationIT.java @@ -0,0 +1,113 @@ +/* + * Copyright 2023-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.repository.cosmosdb.autoconfigure; + +import java.util.List; +import java.util.UUID; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link CosmosDBChatMemoryRepositoryAutoConfiguration}. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+") +class CosmosDBChatMemoryRepositoryAutoConfigurationIT { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(CosmosDBChatMemoryRepositoryAutoConfiguration.class)) + .withPropertyValues( + "spring.ai.chat.memory.repository.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT")) + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.database-name=test-database") + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.container-name=autoconfig-test-container"); + + @Test + void addAndGet() { + this.contextRunner.run(context -> { + CosmosDBChatMemoryRepository memory = context.getBean(CosmosDBChatMemoryRepository.class); + + String conversationId = UUID.randomUUID().toString(); + assertThat(memory.findByConversationId(conversationId)).isEmpty(); + + memory.saveAll(conversationId, List.of(new UserMessage("test question"))); + + assertThat(memory.findByConversationId(conversationId)).hasSize(1); + assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER); + assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question"); + + memory.deleteByConversationId(conversationId); + assertThat(memory.findByConversationId(conversationId)).isEmpty(); + + memory.saveAll(conversationId, + List.of(new UserMessage("test question"), new AssistantMessage("test answer"))); + + assertThat(memory.findByConversationId(conversationId)).hasSize(2); + assertThat(memory.findByConversationId(conversationId).get(0).getMessageType()).isEqualTo(MessageType.USER); + assertThat(memory.findByConversationId(conversationId).get(0).getText()).isEqualTo("test question"); + assertThat(memory.findByConversationId(conversationId).get(1).getMessageType()) + .isEqualTo(MessageType.ASSISTANT); + assertThat(memory.findByConversationId(conversationId).get(1).getText()).isEqualTo("test answer"); + }); + } + + @Test + void propertiesConfiguration() { + this.contextRunner + .withPropertyValues( + "spring.ai.chat.memory.repository.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT")) + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.database-name=test-database") + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.container-name=custom-testcontainer") + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.partition-key-path=/customPartitionKey") + .run(context -> { + CosmosDBChatMemoryRepositoryProperties properties = context + .getBean(CosmosDBChatMemoryRepositoryProperties.class); + assertThat(properties.getEndpoint()).isEqualTo(System.getenv("AZURE_COSMOSDB_ENDPOINT")); + assertThat(properties.getDatabaseName()).isEqualTo("test-database"); + assertThat(properties.getContainerName()).isEqualTo("custom-testcontainer"); + assertThat(properties.getPartitionKeyPath()).isEqualTo("/customPartitionKey"); + }); + } + + @Test + void findConversationIds() { + this.contextRunner.run(context -> { + CosmosDBChatMemoryRepository memory = context.getBean(CosmosDBChatMemoryRepository.class); + + String conversationId1 = UUID.randomUUID().toString(); + String conversationId2 = UUID.randomUUID().toString(); + + memory.saveAll(conversationId1, List.of(new UserMessage("test question 1"))); + memory.saveAll(conversationId2, List.of(new UserMessage("test question 2"))); + + List conversationIds = memory.findConversationIds(); + assertThat(conversationIds).contains(conversationId1, conversationId2); + }); + } + +} diff --git a/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryPropertiesTest.java b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryPropertiesTest.java new file mode 100644 index 00000000000..03c972f3e3a --- /dev/null +++ b/auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/model/chat/memory/repository/cosmosdb/autoconfigure/CosmosDBChatMemoryRepositoryPropertiesTest.java @@ -0,0 +1,73 @@ +/* + * Copyright 2023-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.repository.cosmosdb.autoconfigure; + +import org.junit.jupiter.api.Test; + +import org.springframework.ai.chat.memory.repository.cosmosdb.CosmosDBChatMemoryRepositoryConfig; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Configuration; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link CosmosDBChatMemoryRepositoryProperties}. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +class CosmosDBChatMemoryRepositoryPropertiesTest { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestConfiguration.class); + + @Test + void defaultProperties() { + this.contextRunner.run(context -> { + CosmosDBChatMemoryRepositoryProperties properties = context + .getBean(CosmosDBChatMemoryRepositoryProperties.class); + assertThat(properties.getDatabaseName()) + .isEqualTo(CosmosDBChatMemoryRepositoryConfig.DEFAULT_DATABASE_NAME); + assertThat(properties.getContainerName()) + .isEqualTo(CosmosDBChatMemoryRepositoryConfig.DEFAULT_CONTAINER_NAME); + assertThat(properties.getPartitionKeyPath()) + .isEqualTo(CosmosDBChatMemoryRepositoryConfig.DEFAULT_PARTITION_KEY_PATH); + }); + } + + @Test + void customProperties() { + this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.database-name=custom-db") + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.container-name=custom-container") + .withPropertyValues("spring.ai.chat.memory.repository.cosmosdb.partition-key-path=/custom-partition-key") + .run(context -> { + CosmosDBChatMemoryRepositoryProperties properties = context + .getBean(CosmosDBChatMemoryRepositoryProperties.class); + assertThat(properties.getDatabaseName()).isEqualTo("custom-db"); + assertThat(properties.getContainerName()).isEqualTo("custom-container"); + assertThat(properties.getPartitionKeyPath()).isEqualTo("/custom-partition-key"); + }); + } + + @Configuration + @EnableConfigurationProperties(CosmosDBChatMemoryRepositoryProperties.class) + static class TestConfiguration { + + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/pom.xml b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/pom.xml new file mode 100644 index 00000000000..1872d93dc1a --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/pom.xml @@ -0,0 +1,76 @@ + + + + + 4.0.0 + + org.springframework.ai + spring-ai-parent + 1.1.0-SNAPSHOT + ../../../pom.xml + + + spring-ai-model-chat-memory-repository-cosmos-db + Spring AI Azure Cosmos DB Chat Memory Repository + Spring AI Azure Cosmos DB Chat Memory Repository implementation + + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + org.springframework.ai + spring-ai-model + ${project.version} + + + + com.azure + azure-spring-data-cosmos + ${azure-cosmos.version} + + + + com.azure + azure-identity + ${azure-identity.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.springframework.ai + spring-ai-test + ${project.version} + test + + + + + diff --git a/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepository.java new file mode 100644 index 00000000000..a399608861b --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepository.java @@ -0,0 +1,243 @@ +/* + * Copyright 2023-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.chat.memory.repository.cosmosdb; + +import java.time.Instant; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.stream.Collectors; + +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.models.CosmosBulkOperations; +import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.util.CosmosPagedFlux; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; + +import org.springframework.ai.chat.memory.ChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.util.Assert; + +/** + * An implementation of {@link ChatMemoryRepository} for Azure Cosmos DB. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +public final class CosmosDBChatMemoryRepository implements ChatMemoryRepository { + + public static final String CONVERSATION_TS = CosmosDBChatMemoryRepository.class.getSimpleName() + + "_message_timestamp"; + + private static final Logger logger = LoggerFactory.getLogger(CosmosDBChatMemoryRepository.class); + + private final CosmosAsyncContainer container; + + private CosmosDBChatMemoryRepository(CosmosDBChatMemoryRepositoryConfig config) { + Assert.notNull(config, "config cannot be null"); + this.container = config.getContainer(); + } + + public static CosmosDBChatMemoryRepository create(CosmosDBChatMemoryRepositoryConfig config) { + return new CosmosDBChatMemoryRepository(config); + } + + @Override + public List findConversationIds() { + logger.info("Finding all conversation IDs from Cosmos DB"); + + String query = "SELECT DISTINCT c.conversationId FROM c"; + SqlQuerySpec querySpec = new SqlQuerySpec(query); + + CosmosPagedFlux results = this.container.queryItems(querySpec, new CosmosQueryRequestOptions(), + Object.class); + + List conversationDocs = results.byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + if (conversationDocs == null) { + return Collections.emptyList(); + } + + return conversationDocs.stream() + .filter(Map.class::isInstance) + .map(doc -> (Map) doc) + .map(doc -> (String) doc.get("conversationId")) + .distinct() + .collect(Collectors.toList()); + } + + @Override + public List findByConversationId(String conversationId) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + logger.info("Finding messages for conversation: {}", conversationId); + + String query = "SELECT * FROM c WHERE c.conversationId = @conversationId ORDER BY c._ts ASC"; + SqlParameter param = new SqlParameter("@conversationId", conversationId); + SqlQuerySpec querySpec = new SqlQuerySpec(query, List.of(param)); + + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions() + .setPartitionKey(new PartitionKey(conversationId)); + + CosmosPagedFlux results = this.container.queryItems(querySpec, options, Object.class); + + List messageDocs = results.byPage().flatMapIterable(FeedResponse::getResults).collectList().block(); + + if (messageDocs == null) { + return Collections.emptyList(); + } + + @SuppressWarnings("unchecked") + List messages = messageDocs.stream() + .filter(Map.class::isInstance) + .map(doc -> (Map) doc) + .map(this::mapToMessage) + .collect(Collectors.toList()); + + return messages; + } + + @Override + public void saveAll(String conversationId, List messages) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + Assert.notNull(messages, "messages cannot be null"); + Assert.noNullElements(messages, "messages cannot contain null elements"); + + logger.info("Saving {} messages for conversation: {}", messages.size(), conversationId); + + // First delete existing messages for this conversation + deleteByConversationId(conversationId); + + // Then save the new messages + Instant timestamp = Instant.now(); + + for (int i = 0; i < messages.size(); i++) { + Message message = messages.get(i); + Map doc = createMessageDocument(conversationId, message, timestamp, i); + + this.container.createItem(doc, new PartitionKey(conversationId), new CosmosItemRequestOptions()).block(); + } + } + + @Override + public void deleteByConversationId(String conversationId) { + Assert.hasText(conversationId, "conversationId cannot be null or empty"); + logger.info("Deleting messages for conversation: {}", conversationId); + + String query = "SELECT c.id FROM c WHERE c.conversationId = @conversationId"; + SqlParameter param = new SqlParameter("@conversationId", conversationId); + SqlQuerySpec querySpec = new SqlQuerySpec(query, List.of(param)); + + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions() + .setPartitionKey(new PartitionKey(conversationId)); + + CosmosPagedFlux results = this.container.queryItems(querySpec, options, Object.class); + + List items = results.byPage().flatMapIterable(FeedResponse::getResults).collectList().block(); + + if (items == null || items.isEmpty()) { + return; + } + + @SuppressWarnings("unchecked") + List operations = items.stream() + .filter(Map.class::isInstance) + .map(item -> (Map) item) + .map(item -> CosmosBulkOperations.getDeleteItemOperation((String) item.get("id"), + new PartitionKey(conversationId))) + .collect(Collectors.toList()); + + this.container.executeBulkOperations(Flux.fromIterable(operations)).collectList().block(); + } + + private Map createMessageDocument(String conversationId, Message message, Instant timestamp, + int sequenceNumber) { + Map doc = new HashMap<>(); + doc.put("id", UUID.randomUUID().toString()); + doc.put("conversationId", conversationId); + doc.put("messageType", message.getMessageType().name()); + doc.put("content", message.getText()); + doc.put("sequenceNumber", sequenceNumber); + + // Add timestamp from metadata or use provided timestamp + Instant messageTimestamp = (Instant) message.getMetadata().get(CONVERSATION_TS); + if (messageTimestamp == null) { + messageTimestamp = timestamp; + message.getMetadata().put(CONVERSATION_TS, messageTimestamp); + } + doc.put("messageTimestamp", messageTimestamp.toEpochMilli()); + + // Store any additional metadata + Map filteredMetadata = message.getMetadata() + .entrySet() + .stream() + .filter(entry -> !CONVERSATION_TS.equals(entry.getKey())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + if (!filteredMetadata.isEmpty()) { + doc.put("metadata", filteredMetadata); + } + + return doc; + } + + private Message mapToMessage(Map doc) { + String content = (String) doc.get("content"); + String messageTypeStr = (String) doc.get("messageType"); + MessageType messageType = MessageType.valueOf(messageTypeStr); + + // Reconstruct metadata + Map metadata = new HashMap<>(); + if (doc.containsKey("messageTimestamp")) { + Long timestampMillis = ((Number) doc.get("messageTimestamp")).longValue(); + metadata.put(CONVERSATION_TS, Instant.ofEpochMilli(timestampMillis)); + } + + // Add any additional metadata that was stored + @SuppressWarnings("unchecked") + Map additionalMetadata = (Map) doc.get("metadata"); + if (additionalMetadata != null) { + metadata.putAll(additionalMetadata); + } + + return switch (messageType) { + case ASSISTANT -> new AssistantMessage(content, metadata); + case USER -> UserMessage.builder().text(content).metadata(metadata).build(); + case SYSTEM -> SystemMessage.builder().text(content).metadata(metadata).build(); + case TOOL -> new ToolResponseMessage(List.of(), metadata); + default -> throw new IllegalStateException(String.format("Unknown message type: %s", messageTypeStr)); + }; + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryConfig.java b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryConfig.java new file mode 100644 index 00000000000..b96324cffbd --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryConfig.java @@ -0,0 +1,131 @@ +/* + * Copyright 2023-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.chat.memory.repository.cosmosdb; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; + +import org.springframework.util.Assert; + +/** + * Configuration for the CosmosDB Chat Memory store. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +public final class CosmosDBChatMemoryRepositoryConfig { + + public static final String DEFAULT_DATABASE_NAME = "springai"; + + public static final String DEFAULT_CONTAINER_NAME = "chat_memory"; + + public static final String DEFAULT_PARTITION_KEY_PATH = "/conversationId"; + + private final CosmosAsyncClient cosmosClient; + + private final String databaseName; + + private final String containerName; + + private final String partitionKeyPath; + + private CosmosAsyncContainer container; + + private CosmosDBChatMemoryRepositoryConfig(Builder builder) { + this.cosmosClient = builder.cosmosClient; + this.databaseName = builder.databaseName; + this.containerName = builder.containerName; + this.partitionKeyPath = builder.partitionKeyPath; + this.initializeContainer(); + } + + public static Builder builder() { + return new Builder(); + } + + public CosmosAsyncContainer getContainer() { + return this.container; + } + + public String getDatabaseName() { + return this.databaseName; + } + + public String getContainerName() { + return this.containerName; + } + + public String getPartitionKeyPath() { + return this.partitionKeyPath; + } + + private void initializeContainer() { + // Create database if it doesn't exist + this.cosmosClient.createDatabaseIfNotExists(this.databaseName).block(); + CosmosAsyncDatabase database = this.cosmosClient.getDatabase(this.databaseName); + + // Create container if it doesn't exist + database.createContainerIfNotExists(this.containerName, this.partitionKeyPath).block(); + this.container = database.getContainer(this.containerName); + } + + public static final class Builder { + + private CosmosAsyncClient cosmosClient; + + private String databaseName = DEFAULT_DATABASE_NAME; + + private String containerName = DEFAULT_CONTAINER_NAME; + + private String partitionKeyPath = DEFAULT_PARTITION_KEY_PATH; + + private Builder() { + } + + public Builder withCosmosClient(CosmosAsyncClient cosmosClient) { + this.cosmosClient = cosmosClient; + return this; + } + + public Builder withDatabaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + public Builder withContainerName(String containerName) { + this.containerName = containerName; + return this; + } + + public Builder withPartitionKeyPath(String partitionKeyPath) { + this.partitionKeyPath = partitionKeyPath; + return this; + } + + public CosmosDBChatMemoryRepositoryConfig build() { + Assert.notNull(this.cosmosClient, "CosmosAsyncClient cannot be null"); + Assert.hasText(this.databaseName, "databaseName cannot be null or empty"); + Assert.hasText(this.containerName, "containerName cannot be null or empty"); + Assert.hasText(this.partitionKeyPath, "partitionKeyPath cannot be null or empty"); + + return new CosmosDBChatMemoryRepositoryConfig(this); + } + + } + +} diff --git a/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryIT.java b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryIT.java new file mode 100644 index 00000000000..bffa5894d4d --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/test/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepositoryIT.java @@ -0,0 +1,215 @@ +/* + * Copyright 2023-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.chat.memory.repository.cosmosdb; + +import java.util.List; +import java.util.UUID; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.identity.DefaultAzureCredentialBuilder; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; + +import org.springframework.ai.chat.memory.ChatMemoryRepository; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests for {@link CosmosDBChatMemoryRepository}. + * + * @author Theo van Kraay + * @since 1.0.0 + */ +@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+") +class CosmosDBChatMemoryRepositoryIT { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(CosmosDBChatMemoryRepositoryIT.TestApplication.class); + + private ChatMemoryRepository chatMemoryRepository; + + @BeforeEach + public void setup() { + this.contextRunner.run(context -> this.chatMemoryRepository = context.getBean(ChatMemoryRepository.class)); + } + + @Test + void ensureBeansGetsCreated() { + this.contextRunner.run(context -> { + CosmosDBChatMemoryRepository memory = context.getBean(CosmosDBChatMemoryRepository.class); + Assertions.assertNotNull(memory); + }); + } + + @ParameterizedTest + @CsvSource({ "Message from assistant,ASSISTANT", "Message from user,USER", "Message from system,SYSTEM" }) + void add_shouldInsertSingleMessage(String content, MessageType messageType) { + var conversationId = UUID.randomUUID().toString(); + var message = switch (messageType) { + case ASSISTANT -> new AssistantMessage(content); + case USER -> new UserMessage(content); + case SYSTEM -> new SystemMessage(content); + default -> throw new IllegalArgumentException("Type not supported: " + messageType); + }; + + this.chatMemoryRepository.saveAll(conversationId, List.of(message)); + assertThat(this.chatMemoryRepository.findConversationIds()).isNotEmpty(); + assertThat(this.chatMemoryRepository.findConversationIds()).contains(conversationId); + + List retrievedMessages = this.chatMemoryRepository.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(1); + assertThat(retrievedMessages.get(0).getText()).isEqualTo(content); + assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(messageType); + } + + @Test + void shouldSaveAndRetrieveMultipleMessages() { + var conversationId = UUID.randomUUID().toString(); + + List messages = List.of(new SystemMessage("System message"), new UserMessage("User message"), + new AssistantMessage("Assistant message")); + + this.chatMemoryRepository.saveAll(conversationId, messages); + + List retrievedMessages = this.chatMemoryRepository.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(3); + + // Messages should be in the same order they were saved + assertThat(retrievedMessages.get(0).getText()).isEqualTo("System message"); + assertThat(retrievedMessages.get(0).getMessageType()).isEqualTo(MessageType.SYSTEM); + + assertThat(retrievedMessages.get(1).getText()).isEqualTo("User message"); + assertThat(retrievedMessages.get(1).getMessageType()).isEqualTo(MessageType.USER); + + assertThat(retrievedMessages.get(2).getText()).isEqualTo("Assistant message"); + assertThat(retrievedMessages.get(2).getMessageType()).isEqualTo(MessageType.ASSISTANT); + } + + @Test + void shouldReplaceExistingMessages() { + var conversationId = UUID.randomUUID().toString(); + + // Save initial messages + List initialMessages = List.of(new UserMessage("Initial user message"), + new AssistantMessage("Initial assistant message")); + this.chatMemoryRepository.saveAll(conversationId, initialMessages); + + // Verify initial save + List retrievedMessages = this.chatMemoryRepository.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(2); + + // Replace with new messages + List newMessages = List.of(new SystemMessage("New system message"), + new UserMessage("New user message")); + this.chatMemoryRepository.saveAll(conversationId, newMessages); + + // Verify replacement + retrievedMessages = this.chatMemoryRepository.findByConversationId(conversationId); + assertThat(retrievedMessages).hasSize(2); + assertThat(retrievedMessages.get(0).getText()).isEqualTo("New system message"); + assertThat(retrievedMessages.get(1).getText()).isEqualTo("New user message"); + } + + @Test + void shouldDeleteConversation() { + var conversationId = UUID.randomUUID().toString(); + + // Save messages + List messages = List.of(new UserMessage("User message"), new AssistantMessage("Assistant message")); + this.chatMemoryRepository.saveAll(conversationId, messages); + + // Verify messages exist + assertThat(this.chatMemoryRepository.findByConversationId(conversationId)).hasSize(2); + + // Delete conversation + this.chatMemoryRepository.deleteByConversationId(conversationId); + + // Verify messages are deleted + assertThat(this.chatMemoryRepository.findByConversationId(conversationId)).isEmpty(); + } + + @Test + void shouldFindAllConversationIds() { + var conversationId1 = UUID.randomUUID().toString(); + var conversationId2 = UUID.randomUUID().toString(); + + // Save messages for two conversations + this.chatMemoryRepository.saveAll(conversationId1, List.of(new UserMessage("Message 1"))); + this.chatMemoryRepository.saveAll(conversationId2, List.of(new UserMessage("Message 2"))); + + // Verify both conversation IDs are found + List conversationIds = this.chatMemoryRepository.findConversationIds(); + assertThat(conversationIds).contains(conversationId1, conversationId2); + } + + @Test + void shouldHandleEmptyConversation() { + var conversationId = UUID.randomUUID().toString(); + + // Try to find messages for non-existent conversation + List messages = this.chatMemoryRepository.findByConversationId(conversationId); + assertThat(messages).isEmpty(); + + // Delete non-existent conversation (should not throw) + this.chatMemoryRepository.deleteByConversationId(conversationId); + } + + @SpringBootConfiguration + @EnableAutoConfiguration + static class TestApplication { + + @Bean + public CosmosAsyncClient cosmosAsyncClient() { + return new CosmosClientBuilder().endpoint(System.getenv("AZURE_COSMOSDB_ENDPOINT")) + .credential(new DefaultAzureCredentialBuilder().build()) + .userAgentSuffix("SpringAI-CDBNoSQL-ChatMemoryRepository") + .gatewayMode() + .buildAsyncClient(); + } + + @Bean + public CosmosDBChatMemoryRepositoryConfig cosmosDBChatMemoryRepositoryConfig( + CosmosAsyncClient cosmosAsyncClient) { + return CosmosDBChatMemoryRepositoryConfig.builder() + .withCosmosClient(cosmosAsyncClient) + .withDatabaseName("test-database") + .withContainerName("chat-memory-test-container") + .build(); + } + + @Bean + public CosmosDBChatMemoryRepository cosmosDBChatMemoryRepository(CosmosDBChatMemoryRepositoryConfig config) { + return CosmosDBChatMemoryRepository.create(config); + } + + } + +} diff --git a/pom.xml b/pom.xml index 5031dc90eed..49491b6700c 100644 --- a/pom.xml +++ b/pom.xml @@ -42,6 +42,7 @@ advisors/spring-ai-advisors-vector-store memory/repository/spring-ai-model-chat-memory-repository-cassandra + memory/repository/spring-ai-model-chat-memory-repository-cosmos-db memory/repository/spring-ai-model-chat-memory-repository-jdbc memory/repository/spring-ai-model-chat-memory-repository-neo4j @@ -88,6 +89,7 @@ auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cassandra + auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-jdbc auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 0f5e1e7e26a..0cc9e93e73a 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -205,6 +205,12 @@ ${project.version} + + org.springframework.ai + spring-ai-model-chat-memory-repository-cosmos-db + ${project.version} + + org.springframework.ai spring-ai-model-chat-memory-repository-jdbc @@ -490,6 +496,12 @@ ${project.version} + + org.springframework.ai + spring-ai-autoconfigure-model-chat-memory-repository-cosmos-db + ${project.version} + + org.springframework.ai spring-ai-autoconfigure-model-chat-memory-repository-jdbc 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 00a41dab460..04c239cde10 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 @@ -323,6 +323,90 @@ ChatMemory chatMemory = MessageWindowChatMemory.builder() 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. +=== CosmosDBChatMemoryRepository + +`CosmosDBChatMemoryRepository` is a built-in implementation that uses Azure Cosmos DB NoSQL API to store messages. It is suitable for applications that require a globally distributed, highly scalable document database for chat memory persistence. The repository uses the conversation ID as the partition key to ensure efficient data distribution and fast retrieval. + +First, add the following dependency to your project: + +[tabs] +====== +Maven:: ++ +[source, xml] +---- + + org.springframework.ai + spring-ai-starter-model-chat-memory-repository-cosmos-db + +---- + +Gradle:: ++ +[source,groovy] +---- +dependencies { + implementation 'org.springframework.ai:spring-ai-starter-model-chat-memory-repository-cosmos-db' +} +---- +====== + +Spring AI provides auto-configuration for the `CosmosDBChatMemoryRepository`, which you can use directly in your application. + +[source,java] +---- +@Autowired +CosmosDBChatMemoryRepository chatMemoryRepository; + +ChatMemory chatMemory = MessageWindowChatMemory.builder() + .chatMemoryRepository(chatMemoryRepository) + .maxMessages(10) + .build(); +---- + +If you'd rather create the `CosmosDBChatMemoryRepository` manually, you can do so by providing a `CosmosDBChatMemoryRepositoryConfig` instance: + +[source,java] +---- +ChatMemoryRepository chatMemoryRepository = CosmosDBChatMemoryRepository + .create(CosmosDBChatMemoryRepositoryConfig.builder() + .withCosmosClient(cosmosAsyncClient) + .withDatabaseName("chat-memory-db") + .withContainerName("conversations") + .build()); + +ChatMemory chatMemory = MessageWindowChatMemory.builder() + .chatMemoryRepository(chatMemoryRepository) + .maxMessages(10) + .build(); +---- + +==== Configuration Properties + +[cols="2,5,1",stripes=even] +|=== +|Property | Description | Default Value +| `spring.ai.chat.memory.repository.cosmosdb.endpoint` | Azure Cosmos DB endpoint URI. Required for auto-configuration. | +| `spring.ai.chat.memory.repository.cosmosdb.key` | Azure Cosmos DB primary or secondary key. If not provided, Azure Identity authentication will be used. | +| `spring.ai.chat.memory.repository.cosmosdb.connection-mode` | Connection mode for Cosmos DB client (`direct` or `gateway`). | `gateway` +| `spring.ai.chat.memory.repository.cosmosdb.database-name` | Name of the Cosmos DB database. | `SpringAIChatMemory` +| `spring.ai.chat.memory.repository.cosmosdb.container-name` | Name of the Cosmos DB container. | `ChatMemory` +| `spring.ai.chat.memory.repository.cosmosdb.partition-key-path` | Partition key path for the container. | `/conversationId` +|=== + +==== Authentication + +The Cosmos DB Chat Memory Repository supports two authentication methods: + +1. **Key-based authentication**: Provide the `spring.ai.chat.memory.repository.cosmosdb.key` property with your Cosmos DB primary or secondary key. +2. **Azure Identity authentication**: When no key is provided, the repository uses Azure Identity (`DefaultAzureCredential`) to authenticate with managed identity, service principal, or other Azure credential sources. + +==== Schema Initialization + +The auto-configuration will automatically create the specified database and container if they don't exist. The container is configured with the conversation ID as the partition key (`/conversationId`) to ensure optimal performance for chat memory operations. No manual schema setup is required. + +You can customize the database and container names using the configuration properties mentioned above. + == Memory in Chat Client When using the ChatClient API, you can provide a `ChatMemory` implementation to maintain conversation context across multiple interactions. diff --git a/src/checkstyle/checkstyle-suppressions.xml b/src/checkstyle/checkstyle-suppressions.xml index 33be2bf0fc2..77e6fd16481 100644 --- a/src/checkstyle/checkstyle-suppressions.xml +++ b/src/checkstyle/checkstyle-suppressions.xml @@ -48,6 +48,7 @@ +