Querying chat messages with partition key (microsoft#346)

### Motivation and Context

<!-- Thank you for your contribution to the chat-copilot repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

Some changes to complement [partition key
change](https://github.com/microsoft/chat-copilot/pull/240/files#diff-f9c37b8a51287ee8b7d42811b02d2ee33e4bcb565a155b358216519494eb38a9):
- Documentation updates to make partition key requirements more clear
- Using chatId as partition key when updating chat message state 
- Using partition key on upsert

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

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

webapi/README.md

@@ -3,9 +3,11 @@
 This directory contains the source code for Chat Copilot's backend web API service. The front end web application component can be found in the [webapp/](../webapp/) directory.
 
 ## Running the Chat Copilot sample
+
 To configure and run either the full Chat Copilot application or only the backend API, please view the [main instructions](../README.md#instructions).
 
 # (Under Development)
+
 The following material is under development and may not be complete or accurate.
 
 ## Visual Studio Code
@@ -20,14 +22,13 @@ The following material is under development and may not be complete or accurate.
 2. In Solution Explorer, right-click on `CopilotChatWebApi` and select `Set as Startup Project`.
 3. Start debugging by pressing `F5` or selecting the menu item `Debug`->`Start Debugging`.
 
-1. **(Optional)** To enable support for uploading image file formats such as png, jpg and tiff, there are two options within the `OcrSupport` section of `./appsettings.json`, the Tesseract open source library and Azure Form Recognizer.
+4. **(Optional)** To enable support for uploading image file formats such as png, jpg and tiff, there are two options within the `OcrSupport` section of `./appsettings.json`, the Tesseract open source library and Azure Form Recognizer.
    - **Tesseract** we have included the [Tesseract](https://www.nuget.org/packages/Tesseract) nuget package.
      - You will need to obtain one or more [tessdata language data files](https://github.com/tesseract-ocr/tessdata) such as `eng.traineddata` and add them to your `./data` directory or the location specified in the `OcrSupport:Tesseract:FilePath` location in `./appsettings.json`.
      - Set the `Copy to Output Directory` value to `Copy if newer`.
    - **Azure Form Recognizer** we have included the [Azure.AI.FormRecognizer](https://www.nuget.org/packages/Azure.AI.FormRecognizer) nuget package.
      - You will need to obtain an [Azure Form Recognizer](https://azure.microsoft.com/en-us/services/form-recognizer/) resource and add the `OcrSupport:AzureFormRecognizer:Endpoint` and `OcrSupport:AzureFormRecognizer:Key` values to the `./appsettings.json` file.
 
-
 ## Enabling Sequential Planner
 
 If you want to use SequentialPlanner (multi-step) instead ActionPlanner (single-step), we recommend using `gpt-4` or `gpt-3.5-turbo` as the planner model. **SequentialPlanner works best with `gpt-4`.** Using `gpt-3.5-turbo` will require using a relevancy filter.
@@ -49,6 +50,24 @@ To enable sequential planner,
         \* The `RelevancyThreshold` is a number from 0 to 1 that represents how similar a goal is to a function's name/description/inputs. You want to tune that value when using SequentialPlanner to help keep things scoped while not missing on on things that are relevant or including too many things that really aren't. `0.75` is an arbitrary threshold and we recommend developers play around with this number to see what best fits their scenarios.
 1. Restart the `webapi` - Copilot Chat should be now running locally with SequentialPlanner.
 
+## (Optional) Enabling Cosmos Chat Store.
+
+[Azure Cosmos DB](https://learn.microsoft.com/en-us/azure/cosmos-db/introduction) can be used as a persistent chat store for Chat Copilot. Chat stores are used for storing chat sessions, participants, and messages.
+
+### Prerequisites
+
+#### 1. Containers and PartitionKeys
+
+In an effort to optimize performance, each container must be created with a specific partition key:
+| Store | ContainerName | PartitionKey |
+| ----- | ------------- | ------------ |
+| Chat Sessions | chatsessions | /id (default) |
+| Chat Messages | chatmessages | /chatId |
+| Chat Memory Sources | chatmemorysources | /chatId |
+| Chat Partipants | chatparticipants | /userId |
+
+> For existing customers using CosmosDB before [Release 0.3](https://github.com/microsoft/chat-copilot/releases/tag/0.3), our recommendation is to remove the existing Cosmos DB containers and redeploy to realize the performance update related to the partition schema. To preserve existing chats, containers can be migrated as described [here](https://learn.microsoft.com/en-us/azure/cosmos-db/intra-account-container-copy#copy-a-container).
+
 ## (Optional) Enabling the Qdrant Memory Store
 
 By default, the service uses an in-memory volatile memory store that, when the service stops or restarts, forgets all memories.

webapi/Skills/ChatSkills/ChatSkill.cs

@@ -302,7 +302,7 @@ public async Task<SKContext> ChatAsync(
         if (!string.IsNullOrWhiteSpace(planJson) &&
             !string.IsNullOrEmpty(messageId))
         {
-            await this.UpdateChatMessageContentAsync(planJson, messageId, cancellationToken);
+            await this.UpdateChatMessageContentAsync(planJson, messageId, chatId, cancellationToken);
         }
 
         ChatMessage chatMessage;
@@ -634,14 +634,18 @@ private async Task<ChatMessage> SaveNewResponseAsync(string response, string pro
     /// Updates previously saved response in the chat history.
     /// </summary>
     /// <param name="updatedResponse">Updated response from the chat.</param>
-    /// <param name="messageId">The chat message ID</param>
+    /// <param name="messageId">The chat message ID.</param>
+    /// <param name="chatId">The chat ID that's used as the partition Id.</param>
     /// <param name="cancellationToken">The cancellation token.</param>
-    private async Task UpdateChatMessageContentAsync(string updatedResponse, string messageId, CancellationToken cancellationToken)
+    private async Task UpdateChatMessageContentAsync(string updatedResponse, string messageId, string chatId, CancellationToken cancellationToken)
     {
-        // Make sure the chat exists.
-        var chatMessage = await this._chatMessageRepository.FindByIdAsync(messageId);
-        chatMessage.Content = updatedResponse;
+        ChatMessage? chatMessage = null;
+        if (!await this._chatMessageRepository.TryFindByIdAsync(messageId, chatId, callback: v => chatMessage = v))
+        {
+            throw new ArgumentException($"Chat message {messageId} does not exist.");
+        }
 
+        chatMessage!.Content = updatedResponse;
         await this._chatMessageRepository.UpsertAsync(chatMessage);
     }
 

webapi/Storage/CosmosDbContext.cs

@@ -59,7 +59,7 @@ public async Task CreateAsync(T entity)
             throw new ArgumentOutOfRangeException(nameof(entity.Id), "Entity Id cannot be null or empty.");
         }
 
-        await this._container.CreateItemAsync(entity);
+        await this._container.CreateItemAsync(entity, new PartitionKey(entity.Partition));
     }
 
     /// <inheritdoc/>
@@ -100,7 +100,7 @@ public async Task UpsertAsync(T entity)
             throw new ArgumentOutOfRangeException(nameof(entity.Id), "Entity Id cannot be null or empty.");
         }
 
-        await this._container.UpsertItemAsync(entity);
+        await this._container.UpsertItemAsync(entity, new PartitionKey(entity.Partition));
     }
 
     public void Dispose()

webapi/Storage/IRepository.cs

@@ -32,15 +32,15 @@ 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>
+    /// <param name="partition">Partition key value of the entity.</param>
     /// <returns>An entity</returns>
     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="partition">Partition of the entity.</param>
+    /// <param name="partition">Partition key value 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, string partition, Action<T?> callback);

webapi/Storage/IStorageEntity.cs

@@ -4,7 +4,13 @@ namespace CopilotChat.WebApi.Storage;
 
 public interface IStorageEntity
 {
+    /// <summary>
+    /// Unique ID of the entity.
+    /// </summary>
     string Id { get; set; }
 
+    /// <summary>
+    /// Partition key value.
+    /// </summary>
     string Partition { get; }
 }

webapi/appsettings.json

@@ -121,6 +121,8 @@
     },
     "Cosmos": {
       "Database": "CopilotChat",
+      // IMPORTANT: Each container requires a specific partition key. Ensure these are set correctly in your CosmosDB instance.
+      // See details at ./README.md#1-containers-and-partitionkeys
       "ChatSessionsContainer": "chatsessions",
       "ChatMessagesContainer": "chatmessages",
       "ChatMemorySourcesContainer": "chatmemorysources",