Update partitionKey path definitions for cosmosDatabase containers. (m…

…icrosoft#240)

### Motivation and Context
Customer issue reported in
microsoft#238 related to
partition scheme in cosmosdb containers.

### Description
Update partition-key path definition for `ChatMessage`,
`ChatPartitipant`, and `MemorySource`.

Existing deployments will continue to work in their current mode.

CosmosDB containers must be removed and redeployed to realize the
peformance update related to the partition scheme.

### Contribution Checklist
- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [Contribution
Guidelines](https://github.com/microsoft/copilot-chat/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/copilot-chat/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

scripts/deploy/main.bicep

@@ -637,7 +637,7 @@ resource messageContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabases/co
       }
       partitionKey: {
         paths: [
-          '/id'
+          '/chatId'
         ]
         kind: 'Hash'
         version: 2
@@ -699,7 +699,7 @@ resource participantContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDatabase
       }
       partitionKey: {
         paths: [
-          '/id'
+          '/userId'
         ]
         kind: 'Hash'
         version: 2
@@ -730,7 +730,7 @@ resource memorySourcesContainer 'Microsoft.DocumentDB/databaseAccounts/sqlDataba
       }
       partitionKey: {
         paths: [
-          '/id'
+          '/chatId'
         ]
         kind: 'Hash'
         version: 2

webapi/Controllers/ChatHistoryController.cs

@@ -128,7 +128,7 @@ public async Task<IActionResult> CreateChatSessionAsync(
     public async Task<IActionResult> GetChatSessionByIdAsync(Guid chatId)
     {
         ChatSession? chat = null;
-        if (await this._sessionRepository.TryFindByIdAsync(chatId.ToString(), v => chat = v))
+        if (await this._sessionRepository.TryFindByIdAsync(chatId.ToString(), callback: v => chat = v))
         {
             return this.Ok(chat);
         }
@@ -156,7 +156,7 @@ public async Task<IActionResult> GetAllChatSessionsAsync(string userId)
         foreach (var chatParticipant in chatParticipants)
         {
             ChatSession? chat = null;
-            if (await this._sessionRepository.TryFindByIdAsync(chatParticipant.ChatId, v => chat = v))
+            if (await this._sessionRepository.TryFindByIdAsync(chatParticipant.ChatId, callback: v => chat = v))
             {
                 chats.Add(chat!);
             }
@@ -231,7 +231,7 @@ public async Task<IActionResult> EditChatSessionAsync(
         }
 
         ChatSession? chat = null;
-        if (await this._sessionRepository.TryFindByIdAsync(chatId, v => chat = v))
+        if (await this._sessionRepository.TryFindByIdAsync(chatId, callback: v => chat = v))
         {
             chat!.Title = chatParameters.Title ?? chat!.Title;
             chat!.SystemDescription = chatParameters.SystemDescription ?? chat!.SystemDescription;
@@ -260,7 +260,7 @@ public async Task<ActionResult<IEnumerable<MemorySource>>> GetSourcesAsync(
     {
         this._logger.LogInformation("Get imported sources of chat session {0}", chatId);
 
-        if (await this._sessionRepository.TryFindByIdAsync(chatId.ToString(), v => _ = v))
+        if (await this._sessionRepository.TryFindByIdAsync(chatId.ToString()))
         {
             var sources = await this._sourceRepository.FindByChatIdAsync(chatId.ToString());
             return this.Ok(sources);

webapi/Controllers/ChatMemoryController.cs

@@ -67,7 +67,7 @@ public async Task<IActionResult> GetSemanticMemoriesAsync(
         var sanitizedMemoryName = memoryName.Replace(Environment.NewLine, string.Empty, StringComparison.Ordinal);
 
         // Make sure the chat session exists.
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => _ = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId))
         {
             this._logger.LogWarning("Chat session: {0} does not exist.", sanitizedChatId);
             return this.BadRequest($"Chat session: {sanitizedChatId} does not exist.");

webapi/Controllers/ChatParticipantController.cs

@@ -64,7 +64,7 @@ public async Task<IActionResult> JoinChatAsync(
         string userId = authInfo.UserId;
 
         // Make sure the chat session exists.
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => _ = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId))
         {
             return this.BadRequest("Chat session does not exist.");
         }
@@ -97,7 +97,7 @@ public async Task<IActionResult> JoinChatAsync(
     public async Task<IActionResult> GetAllParticipantsAsync(Guid chatId)
     {
         // Make sure the chat session exists.
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId.ToString(), v => _ = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId.ToString()))
         {
             return this.NotFound("Chat session does not exist.");
         }

webapi/Models/Storage/ChatMessage.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Globalization;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using CopilotChat.WebApi.Models.Response;
 using CopilotChat.WebApi.Storage;
 
@@ -109,6 +110,12 @@ public enum ChatMessageType
     /// </summary>
     public Dictionary<string, int>? TokenUsage { get; set; }
 
+    /// <summary>
+    /// The partition key for the source.
+    /// </summary>
+    [JsonIgnore]
+    public string Partition => this.ChatId;
+
     /// <summary>
     /// Create a new chat message. Timestamp is automatically generated.
     /// </summary>

webapi/Models/Storage/ChatParticipant.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Text.Json.Serialization;
 using CopilotChat.WebApi.Storage;
 
 namespace CopilotChat.WebApi.Models.Storage;
@@ -26,6 +27,12 @@ public class ChatParticipant : IStorageEntity
     /// </summary>
     public string ChatId { get; set; }
 
+    /// <summary>
+    /// The partition key for the source.
+    /// </summary>
+    [JsonIgnore]
+    public string Partition => this.UserId;
+
     public ChatParticipant(string userId, string chatId)
     {
         this.Id = Guid.NewGuid().ToString();

webapi/Models/Storage/ChatSession.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Text.Json.Serialization;
 using CopilotChat.WebApi.Storage;
 
 namespace CopilotChat.WebApi.Models.Storage;
@@ -37,6 +38,12 @@ public class ChatSession : IStorageEntity
     /// </summary>
     public double MemoryBalance { get; set; } = 0.5;
 
+    /// <summary>
+    /// The partition key for the session.
+    /// </summary>
+    [JsonIgnore]
+    public string Partition => this.Id;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="ChatSession"/> class.
     /// </summary>

webapi/Models/Storage/MemorySource.cs

@@ -75,6 +75,12 @@ public class MemorySource : IStorageEntity
     [JsonPropertyName("tokens")]
     public long Tokens { get; set; } = 0;
 
+    /// <summary>
+    /// The partition key for the source.
+    /// </summary>
+    [JsonIgnore]
+    public string Partition => this.ChatId;
+
     /// <summary>
     /// Empty constructor for serialization.
     /// </summary>

webapi/Skills/ChatSkills/ChatSkill.cs

@@ -564,7 +564,7 @@ private async Task<string> AcquireExternalInformationAsync(SKContext context, st
     private async Task<ChatMessage> SaveNewMessageAsync(string message, string userId, string userName, string chatId, string type)
     {
         // Make sure the chat exists.
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => _ = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId))
         {
             throw new ArgumentException("Chat session does not exist.");
         }
@@ -597,7 +597,7 @@ private async Task<ChatMessage> SaveNewMessageAsync(string message, string userI
     private async Task<ChatMessage> SaveNewResponseAsync(string response, string prompt, string chatId, string userId, Dictionary<string, int>? tokenUsage)
     {
         // Make sure the chat exists.
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => _ = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId))
         {
             throw new ArgumentException("Chat session does not exist.");
         }
@@ -775,7 +775,7 @@ private async Task UpdateBotResponseStatusOnClient(string chatId, string status)
     private async Task SetSystemDescriptionAsync(string chatId)
     {
         ChatSession? chatSession = null;
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => chatSession = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, callback: v => chatSession = v))
         {
             throw new ArgumentException("Chat session does not exist.");
         }

webapi/Skills/ChatSkills/SemanticChatMemorySkill.cs

@@ -56,7 +56,7 @@ public async Task<string> QueryMemoriesAsync(
         ISemanticTextMemory textMemory)
     {
         ChatSession? chatSession = null;
-        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, v => chatSession = v))
+        if (!await this._chatSessionRepository.TryFindByIdAsync(chatId, callback: v => chatSession = v))
         {
             throw new ArgumentException($"Chat session {chatId} not found.");
         }

webapi/Storage/CosmosDbContext.cs

@@ -70,11 +70,11 @@ public async Task DeleteAsync(T entity)
             throw new ArgumentOutOfRangeException(nameof(entity.Id), "Entity Id cannot be null or empty.");
         }
 
-        await this._container.DeleteItemAsync<T>(entity.Id, new PartitionKey(entity.Id));
+        await this._container.DeleteItemAsync<T>(entity.Id, new PartitionKey(entity.Partition));
     }
 
     /// <inheritdoc/>
-    public async Task<T> ReadAsync(string entityId)
+    public async Task<T> ReadAsync(string entityId, string partitionKey)
     {
         if (string.IsNullOrWhiteSpace(entityId))
         {
@@ -83,7 +83,7 @@ public async Task<T> ReadAsync(string entityId)
 
         try
         {
-            var response = await this._container.ReadItemAsync<T>(entityId, new PartitionKey(entityId));
+            var response = await this._container.ReadItemAsync<T>(entityId, new PartitionKey(partitionKey));
             return response.Resource;
         }
         catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound)

webapi/Storage/FileSystemContext.cs

@@ -65,7 +65,7 @@ public Task DeleteAsync(T entity)
     }
 
     /// <inheritdoc/>
-    public Task<T> ReadAsync(string entityId)
+    public Task<T> ReadAsync(string entityId, string partitionKey)
     {
         if (string.IsNullOrWhiteSpace(entityId))
         {

webapi/Storage/IRepository.cs

@@ -32,14 +32,16 @@ public interface IRepository<T> where T : IStorageEntity
     /// Finds an entity by its id.
     /// </summary>
     /// <param name="id">Id of the entity.</param>
+    /// <param name="partition">Partition of the entity.</param>
     /// <returns>An entity</returns>
-    Task<T> FindByIdAsync(string id);
+    Task<T> FindByIdAsync(string id, string partition);
 
     /// <summary>
     /// Tries to find an entity by its id.
     /// </summary>
     /// <param name="id">Id of the entity.</param>
-    /// <param name="entity">The entity delegate. Note async methods don't support ref or out parameters.</param>
+    /// <param name="partition">Partition of the entity.</param>
+    /// <param name="callback">The entity delegate. Note async methods don't support ref or out parameters.</param>
     /// <returns>True if the entity was found, false otherwise.</returns>
-    Task<bool> TryFindByIdAsync(string id, Action<T?> entity);
+    Task<bool> TryFindByIdAsync(string id, string partition, Action<T?> callback);
 }

webapi/Storage/IStorageContext.cs

@@ -20,8 +20,9 @@ public interface IStorageContext<T> where T : IStorageEntity
     /// Read an entity from the storage context by id.
     /// </summary>
     /// <param name="entityId">The entity id.</param>
+    /// <param name="partitionKey">The entity partition</param>
     /// <returns>The entity.</returns>
-    Task<T> ReadAsync(string entityId);
+    Task<T> ReadAsync(string entityId, string partitionKey);
 
     /// <summary>
     /// Create an entity in the storage context.

webapi/Storage/IStorageEntity.cs

@@ -5,4 +5,6 @@ namespace CopilotChat.WebApi.Storage;
 public interface IStorageEntity
 {
     string Id { get; set; }
+
+    string Partition { get; }
 }

webapi/Storage/Repository.cs

@@ -42,22 +42,21 @@ public Task DeleteAsync(T entity)
     }
 
     /// <inheritdoc/>
-    public Task<T> FindByIdAsync(string id)
+    public Task<T> FindByIdAsync(string id, string? partition = null)
     {
-        return this.StorageContext.ReadAsync(id);
+        return this.StorageContext.ReadAsync(id, partition ?? id);
     }
 
     /// <inheritdoc/>
-    public async Task<bool> TryFindByIdAsync(string id, Action<T?> entity)
+    public async Task<bool> TryFindByIdAsync(string id, string? partition = null, Action<T?>? callback = null)
     {
         try
         {
-            entity(await this.FindByIdAsync(id));
+            callback?.Invoke(await this.FindByIdAsync(id, partition ?? id));
             return true;
         }
         catch (Exception ex) when (ex is ArgumentOutOfRangeException || ex is KeyNotFoundException)
         {
-            entity(default);
             return false;
         }
     }

webapi/Storage/VolatileContext.cs

@@ -61,7 +61,7 @@ public Task DeleteAsync(T entity)
     }
 
     /// <inheritdoc/>
-    public Task<T> ReadAsync(string entityId)
+    public Task<T> ReadAsync(string entityId, string partitionKey)
     {
         if (string.IsNullOrWhiteSpace(entityId))
         {