Added Cosmos DB chat history feature to the backend

Author: fujita-h

app/backend/app.py

@@ -16,6 +16,7 @@
     SpeechSynthesizer,
 )
 from azure.core.exceptions import ResourceNotFoundError
+from azure.cosmos.aio import ContainerProxy, CosmosClient
 from azure.identity.aio import (
     AzureDeveloperCliCredential,
     ManagedIdentityCredential,
@@ -60,7 +61,9 @@
     CONFIG_BLOB_CONTAINER_CLIENT,
     CONFIG_CHAT_APPROACH,
     CONFIG_CHAT_HISTORY_BROWSER_ENABLED,
+    CONFIG_CHAT_HISTORY_COSMOS_ENABLED,
     CONFIG_CHAT_VISION_APPROACH,
+    CONFIG_COSMOS_HISTORY_CONTAINER,
     CONFIG_CREDENTIAL,
     CONFIG_GPT4V_DEPLOYED,
     CONFIG_INGESTER,
@@ -224,7 +227,10 @@ async def chat(auth_claims: Dict[str, Any]):
         # else creates a new session_id depending on the chat history options enabled.
         session_state = request_json.get("session_state")
         if session_state is None:
-            session_state = create_session_id(current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED])
+            session_state = create_session_id(
+                current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED],
+                current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED],
+            )
         result = await approach.run(
             request_json["messages"],
             context=context,
@@ -255,7 +261,10 @@ async def chat_stream(auth_claims: Dict[str, Any]):
         # else creates a new session_id depending on the chat history options enabled.
         session_state = request_json.get("session_state")
         if session_state is None:
-            session_state = create_session_id(current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED])
+            session_state = create_session_id(
+                current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED],
+                current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED],
+            )
         result = await approach.run_stream(
             request_json["messages"],
             context=context,
@@ -289,6 +298,7 @@ def config():
             "showSpeechOutputBrowser": current_app.config[CONFIG_SPEECH_OUTPUT_BROWSER_ENABLED],
             "showSpeechOutputAzure": current_app.config[CONFIG_SPEECH_OUTPUT_AZURE_ENABLED],
             "showChatHistoryBrowser": current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED],
+            "showChatHistoryCosmos": current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED],
         }
     )
 
@@ -397,6 +407,129 @@ async def list_uploaded(auth_claims: dict[str, Any]):
     return jsonify(files), 200
 
 
+@bp.post("/chat_history")
+@authenticated
+async def post_chat_history(auth_claims: Dict[str, Any]):
+    if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    container = cast(ContainerProxy, current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER])
+    if not container:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    entra_oid = auth_claims.get("oid")
+    if not entra_oid:
+        return jsonify({"error": "User OID not found"}), 401
+
+    try:
+        request_json = await request.get_json()
+        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]
+
+        await container.upsert_item({"id": id, "entra_oid": entra_oid, "title": title, "answers": answers})
+
+        return jsonify({}), 201
+    except Exception as error:
+        return error_response(error, "/chat_history")
+
+
+@bp.post("/chat_history/items")
+@authenticated
+async def get_chat_history(auth_claims: Dict[str, Any]):
+    if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    container = cast(ContainerProxy, current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER])
+    if not container:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    entra_oid = auth_claims.get("oid")
+    if not entra_oid:
+        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")
+
+        res = container.query_items(
+            query="SELECT c.id, c.entra_oid, c.title, c._ts FROM c WHERE c.entra_oid = @entra_oid ORDER BY c._ts DESC",
+            parameters=[dict(name="@entra_oid", value=entra_oid)],
+            max_item_count=count,
+        )
+
+        # set the continuation token for the next page
+        pager = res.by_page(continuation_token)
+
+        # Get the first page, and the continuation token
+        try:
+            page = await pager.__anext__()
+            continuation_token = pager.continuation_token  # type: ignore
+
+            items = []
+            async for item in page:
+                items.append(item)
+
+        # If there are no page, StopAsyncIteration is raised
+        except StopAsyncIteration:
+            items = []
+            continuation_token = None
+
+        return jsonify({"items": items, "continuation_token": continuation_token}), 200
+
+    except Exception as error:
+        return error_response(error, "/chat_history/items")
+
+
+@bp.get("/chat_history/items/<path>")
+@authenticated_path
+async def get_chat_history_session(path: str, auth_claims: Dict[str, Any]):
+    if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    container = cast(ContainerProxy, current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER])
+    if not container:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    if not path:
+        return jsonify({"error": "Invalid path"}), 400
+
+    entra_oid = auth_claims.get("oid")
+    if not entra_oid:
+        return jsonify({"error": "User OID not found"}), 401
+
+    try:
+        res = await container.read_item(item=path, partition_key=entra_oid)
+        return jsonify(res), 200
+    except Exception as error:
+        return error_response(error, f"/chat_history/items/{path}")
+
+
+@bp.delete("/chat_history/items/<path>")
+@authenticated_path
+async def delete_chat_history_session(path: str, auth_claims: Dict[str, Any]):
+    if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    container = cast(ContainerProxy, current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER])
+    if not container:
+        return jsonify({"error": "Chat history not enabled"}), 405
+
+    if not path:
+        return jsonify({"error": "Invalid path"}), 400
+
+    entra_oid = auth_claims.get("oid")
+    if not entra_oid:
+        return jsonify({"error": "User OID not found"}), 401
+
+    try:
+        await container.delete_item(item=path, partition_key=entra_oid)
+        return jsonify({}), 200
+    except Exception as error:
+        return error_response(error, f"/chat_history/items/{path}")
+
+
 @bp.before_app_serving
 async def setup_clients():
     # Replace these with your own values, either in environment variables or directly here
@@ -452,7 +585,12 @@ async def setup_clients():
     USE_SPEECH_INPUT_BROWSER = os.getenv("USE_SPEECH_INPUT_BROWSER", "").lower() == "true"
     USE_SPEECH_OUTPUT_BROWSER = os.getenv("USE_SPEECH_OUTPUT_BROWSER", "").lower() == "true"
     USE_SPEECH_OUTPUT_AZURE = os.getenv("USE_SPEECH_OUTPUT_AZURE", "").lower() == "true"
+
     USE_CHAT_HISTORY_BROWSER = os.getenv("USE_CHAT_HISTORY_BROWSER", "").lower() == "true"
+    USE_CHAT_HISTORY_COSMOS = os.getenv("USE_CHAT_HISTORY_COSMOS", "").lower() == "true"
+    AZURE_COSMOSDB_ACCOUNT = os.getenv("AZURE_COSMOSDB_ACCOUNT")
+    AZURE_CHAT_HISTORY_DATABASE = os.getenv("AZURE_CHAT_HISTORY_DATABASE")
+    AZURE_CHAT_HISTORY_CONTAINER = os.getenv("AZURE_CHAT_HISTORY_CONTAINER")
 
     # WEBSITE_HOSTNAME is always set by App Service, RUNNING_IN_PRODUCTION is set in main.bicep
     RUNNING_ON_AZURE = os.getenv("WEBSITE_HOSTNAME") is not None or os.getenv("RUNNING_IN_PRODUCTION") is not None
@@ -556,6 +694,21 @@ async def setup_clients():
         )
         current_app.config[CONFIG_INGESTER] = ingester
 
+    if USE_CHAT_HISTORY_COSMOS:
+        current_app.logger.info("USE_CHAT_HISTORY_COSMOS is true, setting up CosmosDB client")
+        if not AZURE_COSMOSDB_ACCOUNT:
+            raise ValueError("AZURE_COSMOSDB_ACCOUNT must be set when USE_CHAT_HISTORY_COSMOS is true")
+        if not AZURE_CHAT_HISTORY_DATABASE:
+            raise ValueError("AZURE_CHAT_HISTORY_DATABASE must be set when USE_CHAT_HISTORY_COSMOS is true")
+        if not AZURE_CHAT_HISTORY_CONTAINER:
+            raise ValueError("AZURE_CHAT_HISTORY_CONTAINER must be set when USE_CHAT_HISTORY_COSMOS is true")
+        cosmos_client = CosmosClient(
+            url=f"https://{AZURE_COSMOSDB_ACCOUNT}.documents.azure.com:443/", credential=azure_credential
+        )
+        cosmos_db = cosmos_client.get_database_client(AZURE_CHAT_HISTORY_DATABASE)
+        cosmos_container = cosmos_db.get_container_client(AZURE_CHAT_HISTORY_CONTAINER)
+        current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER] = cosmos_container
+
     # Used by the OpenAI SDK
     openai_client: AsyncOpenAI
 
@@ -624,6 +777,7 @@ async def setup_clients():
     current_app.config[CONFIG_SPEECH_OUTPUT_BROWSER_ENABLED] = USE_SPEECH_OUTPUT_BROWSER
     current_app.config[CONFIG_SPEECH_OUTPUT_AZURE_ENABLED] = USE_SPEECH_OUTPUT_AZURE
     current_app.config[CONFIG_CHAT_HISTORY_BROWSER_ENABLED] = USE_CHAT_HISTORY_BROWSER
+    current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED] = USE_CHAT_HISTORY_COSMOS
 
     # Various approaches to integrate GPT and external knowledge, most applications will use a single one of these patterns
     # or some derivative, here we include several for exploration purposes

app/backend/config.py

@@ -23,3 +23,5 @@
 CONFIG_SPEECH_SERVICE_TOKEN = "speech_service_token"
 CONFIG_SPEECH_SERVICE_VOICE = "speech_service_voice"
 CONFIG_CHAT_HISTORY_BROWSER_ENABLED = "chat_history_browser_enabled"
+CONFIG_CHAT_HISTORY_COSMOS_ENABLED = "chat_history_cosmos_enabled"
+CONFIG_COSMOS_HISTORY_CONTAINER = "cosmos_history_container"

app/backend/core/sessionhelper.py

@@ -2,7 +2,11 @@
 from typing import Union
 
 
-def create_session_id(config_chat_history_browser_enabled: bool) -> Union[str, None]:
+def create_session_id(
+    config_chat_history_cosmos_enabled: bool, config_chat_history_browser_enabled: bool
+) -> Union[str, None]:
+    if config_chat_history_cosmos_enabled:
+        return str(uuid.uuid4())
     if config_chat_history_browser_enabled:
         return str(uuid.uuid4())
     return None

app/backend/requirements.in

@@ -7,6 +7,7 @@ tiktoken
 tenacity
 azure-ai-documentintelligence
 azure-cognitiveservices-speech
+azure-cosmos
 azure-search-documents==11.6.0b6
 azure-storage-blob
 azure-storage-file-datalake

app/backend/requirements.txt

@@ -34,6 +34,7 @@ azure-core==1.30.2
     # via
     #   azure-ai-documentintelligence
     #   azure-core-tracing-opentelemetry
+    #   azure-cosmos
     #   azure-identity
     #   azure-monitor-opentelemetry
     #   azure-monitor-opentelemetry-exporter
@@ -44,6 +45,8 @@ azure-core==1.30.2
     #   msrest
 azure-core-tracing-opentelemetry==1.0.0b11
     # via azure-monitor-opentelemetry
+azure-cosmos==4.7.0
+    # via -r requirements.in
 azure-identity==1.17.1
     # via
     #   -r requirements.in
@@ -402,7 +405,9 @@ typing-extensions==4.12.2
     # via
     #   azure-ai-documentintelligence
     #   azure-core
+    #   azure-cosmos
     #   azure-identity
+    #   azure-search-documents
     #   azure-storage-blob
     #   azure-storage-file-datalake
     #   openai