From f963b0283c1aa842c2329d2b1eb51e7439a031cb Mon Sep 17 00:00:00 2001
From: Chris <66376200+crickman@users.noreply.github.com>
Date: Mon, 28 Aug 2023 13:49:40 -0700
Subject: [PATCH] Update partitionKey path definitions for cosmosDatabase
 containers. (#240)

### Motivation and Context
Customer issue reported in
https://github.com/microsoft/chat-copilot/issues/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 :smile:
---
 scripts/deploy/main.bicep                           | 6 +++---
 webapi/Controllers/ChatHistoryController.cs         | 8 ++++----
 webapi/Controllers/ChatMemoryController.cs          | 2 +-
 webapi/Controllers/ChatParticipantController.cs     | 4 ++--
 webapi/Models/Storage/ChatMessage.cs                | 7 +++++++
 webapi/Models/Storage/ChatParticipant.cs            | 7 +++++++
 webapi/Models/Storage/ChatSession.cs                | 7 +++++++
 webapi/Models/Storage/MemorySource.cs               | 6 ++++++
 webapi/Skills/ChatSkills/ChatSkill.cs               | 6 +++---
 webapi/Skills/ChatSkills/SemanticChatMemorySkill.cs | 2 +-
 webapi/Storage/CosmosDbContext.cs                   | 6 +++---
 webapi/Storage/FileSystemContext.cs                 | 2 +-
 webapi/Storage/IRepository.cs                       | 8 +++++---
 webapi/Storage/IStorageContext.cs                   | 3 ++-
 webapi/Storage/IStorageEntity.cs                    | 2 ++
 webapi/Storage/Repository.cs                        | 9 ++++-----
 webapi/Storage/VolatileContext.cs                   | 2 +-
 17 files changed, 59 insertions(+), 28 deletions(-)

diff --git a/scripts/deploy/main.bicep b/scripts/deploy/main.bicep
index 2eb9636d9..4df201eac 100644
--- a/scripts/deploy/main.bicep
+++ b/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
diff --git a/webapi/Controllers/ChatHistoryController.cs b/webapi/Controllers/ChatHistoryController.cs
index 58f1078c7..d2c0f0903 100644
--- a/webapi/Controllers/ChatHistoryController.cs
+++ b/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);
diff --git a/webapi/Controllers/ChatMemoryController.cs b/webapi/Controllers/ChatMemoryController.cs
index 3f8e8d276..2190c93eb 100644
--- a/webapi/Controllers/ChatMemoryController.cs
+++ b/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.");
diff --git a/webapi/Controllers/ChatParticipantController.cs b/webapi/Controllers/ChatParticipantController.cs
index 18629ab71..62b5b6ee8 100644
--- a/webapi/Controllers/ChatParticipantController.cs
+++ b/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.");
         }
diff --git a/webapi/Models/Storage/ChatMessage.cs b/webapi/Models/Storage/ChatMessage.cs
index 32ff183ca..1ecdc6294 100644
--- a/webapi/Models/Storage/ChatMessage.cs
+++ b/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>
diff --git a/webapi/Models/Storage/ChatParticipant.cs b/webapi/Models/Storage/ChatParticipant.cs
index 732c94494..11bd86e05 100644
--- a/webapi/Models/Storage/ChatParticipant.cs
+++ b/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();
diff --git a/webapi/Models/Storage/ChatSession.cs b/webapi/Models/Storage/ChatSession.cs
index 639c5ef86..329358742 100644
--- a/webapi/Models/Storage/ChatSession.cs
+++ b/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>
diff --git a/webapi/Models/Storage/MemorySource.cs b/webapi/Models/Storage/MemorySource.cs
index 1b11003aa..a9488ba46 100644
--- a/webapi/Models/Storage/MemorySource.cs
+++ b/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>
diff --git a/webapi/Skills/ChatSkills/ChatSkill.cs b/webapi/Skills/ChatSkills/ChatSkill.cs
index 2c11b1d95..fdb46c71a 100644
--- a/webapi/Skills/ChatSkills/ChatSkill.cs
+++ b/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.");
         }
diff --git a/webapi/Skills/ChatSkills/SemanticChatMemorySkill.cs b/webapi/Skills/ChatSkills/SemanticChatMemorySkill.cs
index 25f5f43fa..ae464a398 100644
--- a/webapi/Skills/ChatSkills/SemanticChatMemorySkill.cs
+++ b/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.");
         }
diff --git a/webapi/Storage/CosmosDbContext.cs b/webapi/Storage/CosmosDbContext.cs
index 6a8e6264d..f3ccb8bc5 100644
--- a/webapi/Storage/CosmosDbContext.cs
+++ b/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)
diff --git a/webapi/Storage/FileSystemContext.cs b/webapi/Storage/FileSystemContext.cs
index 8946231d7..44803f1b8 100644
--- a/webapi/Storage/FileSystemContext.cs
+++ b/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))
         {
diff --git a/webapi/Storage/IRepository.cs b/webapi/Storage/IRepository.cs
index fb2e1e7f4..7e564afc8 100644
--- a/webapi/Storage/IRepository.cs
+++ b/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);
 }
diff --git a/webapi/Storage/IStorageContext.cs b/webapi/Storage/IStorageContext.cs
index 54e350d37..c8f55dc06 100644
--- a/webapi/Storage/IStorageContext.cs
+++ b/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.
diff --git a/webapi/Storage/IStorageEntity.cs b/webapi/Storage/IStorageEntity.cs
index e1aa6ab72..91d5880ef 100644
--- a/webapi/Storage/IStorageEntity.cs
+++ b/webapi/Storage/IStorageEntity.cs
@@ -5,4 +5,6 @@ namespace CopilotChat.WebApi.Storage;
 public interface IStorageEntity
 {
     string Id { get; set; }
+
+    string Partition { get; }
 }
diff --git a/webapi/Storage/Repository.cs b/webapi/Storage/Repository.cs
index d6add3cad..4cdad43fa 100644
--- a/webapi/Storage/Repository.cs
+++ b/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;
         }
     }
diff --git a/webapi/Storage/VolatileContext.cs b/webapi/Storage/VolatileContext.cs
index 4cf7a82aa..122b639d0 100644
--- a/webapi/Storage/VolatileContext.cs
+++ b/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))
         {