CosmosDB progress

Author: pamelafox

app/backend/chat_history/cosmosdb.py

@@ -1,8 +1,11 @@
+import logging
 import os
 import time
 from typing import Any, Dict, Union
 
+from azure.cosmos import exceptions
 from azure.cosmos.aio import ContainerProxy, CosmosClient
+from azure.cosmos.partition_key import PartitionKey
 from azure.identity.aio import AzureDeveloperCliCredential, ManagedIdentityCredential
 from quart import Blueprint, current_app, jsonify, request
 
@@ -15,9 +18,19 @@
 from decorators import authenticated
 from error import error_response
 
+logger = logging.getLogger("scripts")
+
 chat_history_cosmosdb_bp = Blueprint("chat_history_cosmos", __name__, static_folder="static")
 
 
+def make_partition_key(entra_id, session_id=None):
+    if entra_id and session_id:
+        # Need multihash for hierachical partitioning
+        return PartitionKey(path=["/entra_id", "/session_id"], kind="MultiHash")
+    else:
+        return PartitionKey(path="/entra_id")
+
+
 @chat_history_cosmosdb_bp.post("/chat_history")
 @authenticated
 async def post_chat_history(auth_claims: Dict[str, Any]):
@@ -34,44 +47,69 @@ async def post_chat_history(auth_claims: Dict[str, Any]):
 
     try:
         request_json = await request.get_json()
-        id = request_json.get("id")
+        session_id = request_json.get("id")
         answers = request_json.get("answers")
         title = answers[0][0][:50] + "..." if len(answers[0][0]) > 50 else answers[0][0]
         timestamp = int(time.time() * 1000)
 
         # Insert the session item:
-        await container.upsert_item(
-            {
-                "id": id,
-                "session_id": id,
-                "entra_oid": entra_oid,
-                "type": "session",
-                "title": title,
-                "timestamp": timestamp,
-            }
-        )
-
+        session_item = {
+            "id": id,
+            "session_id": session_id,
+            "entra_oid": entra_oid,
+            "type": "session",
+            "title": title,
+            "timestamp": timestamp,
+        }
+
+        message_items = []
         # Now insert a message item for each question/response pair:
         for ind, message_pair in enumerate(zip(answers[::2], answers[1::2])):
-            # TODO: Can I do a batch upsert?
-            await container.upsert_item(
+            # The id: what if you delete a message and then add a new one? The id will be the same.
+            # If we had delete mechanism, and you deleted item 5 in a history, then item 6 would still hang around
+            # and youd have two of item 6.
+            # abc-0
+            # abc-1
+            # abc-2 <-- DELETE
+            # abc-3
+            # One approach would be to delete EVERYTHING, then upsert everything.
+            # Another approach would be to delete item plus everything after, then upsert everything after.
+            # Or: Change the frontend?
+            # We can do this first, and change the frontend after
+            message_items.append(
                 {
-                    "id": f"{id}-{ind}",
+                    "id": f"{session_id}-{ind}",
                     "session_id": id,
                     "entra_oid": entra_oid,
                     "type": "message",
                     "question": message_pair[0],
                     "response": message_pair[1],
-                    "timestamp": timestamp,
+                    "timestamp": timestamp,  # <-- This is the timestamp of the session, not the message
                 }
             )
 
+        batch_operations = [("upsert", tuple([session_item] + message_items), {})]
+
+        try:
+            # Run that list of operations
+            batch_results = container.execute_item_batch(
+                batch_operations=batch_operations, partition_key=make_partition_key(entra_oid, session_id)
+            )
+            # Batch results are returned as a list of item operation results - or raise a CosmosBatchOperationError if
+            # one of the operations failed within your batch request.
+            print(f"\nResults for the batch operations: {batch_results}\n")
+        except exceptions.CosmosBatchOperationError as e:
+            error_operation_index = e.error_index
+            error_operation_response = e.operation_responses[error_operation_index]
+            error_operation = batch_operations[error_operation_index]
+            logger.error(f"Batch operation failed: {error_operation_response} for operation {error_operation}")
+            return jsonify({"error": "Batch operation failed"}), 400
         return jsonify({}), 201
     except Exception as error:
         return error_response(error, "/chat_history")
 
 
-@chat_history_cosmosdb_bp.post("/chat_history/items")
+@chat_history_cosmosdb_bp.get("/chat_history/items")
 @authenticated
 async def get_chat_history(auth_claims: Dict[str, Any]):
     if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
@@ -86,14 +124,15 @@ async def get_chat_history(auth_claims: Dict[str, Any]):
         return jsonify({"error": "User OID not found"}), 401
 
     try:
-        request_json = await request.get_json()
-        count = request_json.get("count", 20)
-        continuation_token = request_json.get("continuation_token")
+        # get the count and continuation token from the request URL
+        count = request.args.get("count", 10)
+        continuation_token = request.args.get("continuation_token")
 
         res = container.query_items(
+            # TODO: do we need distinct? per Mark's code - Mark says no!
             query="SELECT c.id, c.entra_oid, c.title, c.timestamp FROM c WHERE c.entra_oid = @entra_oid AND c.type = @type ORDER BY c.timestamp DESC",
             parameters=[dict(name="@entra_oid", value=entra_oid), dict(name="@type", value="session")],
-            partition_key=entra_oid,
+            partition_key=make_partition_key(entra_oid),
             max_item_count=count,
         )
 
@@ -142,6 +181,14 @@ async def get_chat_history_session(auth_claims: Dict[str, Any], item_id: str):
         return jsonify({"error": "User OID not found"}), 401
 
     try:
+        res = container.query_items(
+            # TODO: do we need distinct? per Mark's code
+            query="SELECT c.id, c.entra_oid, c.title, c.timestamp FROM c WHERE c.session_id = @session_id ORDER BY c.timestamp DESC",
+            parameters=[dict(name="@entra_oid", value=entra_oid), dict(name="@session_id", value=item_id)],
+            partition_key=make_partition_key(entra_oid, item_id),
+            # max_item_count=?
+        )
+
         res = await container.read_item(item=item_id, partition_key=entra_oid)
         return (
             jsonify(
@@ -175,6 +222,8 @@ async def delete_chat_history_session(auth_claims: Dict[str, Any], item_id: str)
 
     try:
         await container.delete_item(item=item_id, partition_key=entra_oid)
+        # Delete session, and all the message items associated with it
+        # TODO: Delete all the message items as well
         return jsonify({}), 204
     except Exception as error:
         return error_response(error, f"/chat_history/items/{item_id}")

app/frontend/src/api/api.ts

@@ -145,10 +145,15 @@ export async function postChatHistoryApi(item: any, idToken: string): Promise<an
 
 export async function getChatHistoryListApi(count: number, continuationToken: string | undefined, idToken: string): Promise<HistoryListApiResponse> {
     const headers = await getHeaders(idToken);
-    const response = await fetch("/chat_history/items", {
-        method: "POST",
-        headers: { ...headers, "Content-Type": "application/json" },
-        body: JSON.stringify({ count: count, continuation_token: continuationToken })
+    const url = new URL("/chat_history/items", BACKEND_URI);
+    url.searchParams.append("count", count.toString());
+    if (continuationToken) {
+        url.searchParams.append("continuationToken", continuationToken);
+    }
+
+    const response = await fetch(url.toString(), {
+        method: "GET",
+        headers: { ...headers, "Content-Type": "application/json" }
     });
 
     if (!response.ok) {

infra/main.bicep

@@ -815,6 +815,9 @@ module cosmosDb 'br/public:avm/res/document-db/database-account:0.6.1' = if (use
                 {
                   path: '/timestamp/?'
                 }
+                {
+                  path: '/type/?'
+                }
               ]
               excludedPaths: [
                 {