Integrating Azure OpenAI and Cosmos DB with Improved Quart App Functionality

azure-openai-features.md

@@ -0,0 +1,266 @@
+import copy
+import json
+import os
+import logging
+import uuid
+import httpx
+import asyncio
+from quart import (
+    Blueprint,
+    Quart,
+    jsonify,
+    make_response,
+    request,
+    send_from_directory,
+    render_template,
+    current_app,
+)
+
+from openai import AsyncAzureOpenAI
+from azure.identity.aio import (
+    DefaultAzureCredential,
+    get_bearer_token_provider
+)
+from backend.auth.auth_utils import get_authenticated_user_details
+from backend.security.ms_defender_utils import get_msdefender_user_json
+from backend.history.cosmosdbservice import CosmosConversationClient
+from backend.settings import (
+    app_settings,
+    MINIMUM_SUPPORTED_AZURE_OPENAI_PREVIEW_API_VERSION
+)
+from backend.utils import (
+    format_as_ndjson,
+    format_stream_response,
+    format_non_streaming_response,
+    convert_to_pf_format,
+    format_pf_non_streaming_response,
+)
+
+bp = Blueprint("routes", __name__, static_folder="static", template_folder="static")
+
+cosmos_db_ready = asyncio.Event()
+
+# Debug settings
+DEBUG = os.environ.get("DEBUG", "false")
+if DEBUG.lower() == "true":
+    logging.basicConfig(level=logging.DEBUG)
+
+USER_AGENT = "GitHubSampleWebApp/AsyncAzureOpenAI/1.0.0"
+
+# Frontend Settings via Environment Variables
+frontend_settings = {
+    "auth_enabled": app_settings.base_settings.auth_enabled,
+    "feedback_enabled": (
+        app_settings.chat_history and
+        app_settings.chat_history.enable_feedback
+    ),
+    "ui": {
+        "title": app_settings.ui.title,
+        "logo": app_settings.ui.logo,
+        "chat_logo": app_settings.ui.chat_logo or app_settings.ui.logo,
+        "chat_title": app_settings.ui.chat_title,
+        "chat_description": app_settings.ui.chat_description,
+        "show_share_button": app_settings.ui.show_share_button,
+        "show_chat_history_button": app_settings.ui.show_chat_history_button,
+    },
+    "sanitize_answer": app_settings.base_settings.sanitize_answer,
+}
+
+# Enable Microsoft Defender for Cloud Integration
+MS_DEFENDER_ENABLED = os.environ.get("MS_DEFENDER_ENABLED", "true").lower() == "true"
+
+# Initialize Azure OpenAI Client
+async def init_openai_client():
+    azure_openai_client = None
+    try:
+        # API version check
+        if (
+            app_settings.azure_openai.preview_api_version
+            < MINIMUM_SUPPORTED_AZURE_OPENAI_PREVIEW_API_VERSION
+        ):
+            raise ValueError(
+                f"The minimum supported Azure OpenAI preview API version is '{MINIMUM_SUPPORTED_AZURE_OPENAI_PREVIEW_API_VERSION}'"
+            )
+
+        # Endpoint
+        if (
+            not app_settings.azure_openai.endpoint and
+            not app_settings.azure_openai.resource
+        ):
+            raise ValueError(
+                "AZURE_OPENAI_ENDPOINT or AZURE_OPENAI_RESOURCE is required"
+            )
+
+        endpoint = (
+            app_settings.azure_openai.endpoint
+            if app_settings.azure_openai.endpoint
+            else f"https://{app_settings.azure_openai.resource}.openai.azure.com/"
+        )
+
+        # Authentication
+        aoai_api_key = app_settings.azure_openai.key
+        ad_token_provider = None
+        if not aoai_api_key:
+            logging.debug("No AZURE_OPENAI_KEY found, using Azure Entra ID auth")
+            async with DefaultAzureCredential() as credential:
+                ad_token_provider = get_bearer_token_provider(
+                    credential,
+                    "https://cognitiveservices.azure.com/.default"
+                )
+
+        # Deployment
+        deployment = app_settings.azure_openai.model
+        if not deployment:
+            raise ValueError("AZURE_OPENAI_MODEL is required")
+
+        # Default Headers
+        default_headers = {"x-ms-useragent": USER_AGENT}
+
+        azure_openai_client = AsyncAzureOpenAI(
+            api_version=app_settings.azure_openai.preview_api_version,
+            api_key=aoai_api_key,
+            azure_ad_token_provider=ad_token_provider,
+            default_headers=default_headers,
+            azure_endpoint=endpoint,
+        )
+
+        return azure_openai_client
+    except Exception as e:
+        logging.exception("Exception in Azure OpenAI initialization", e)
+        azure_openai_client = None
+        raise e
+
+async def init_cosmosdb_client():
+    cosmos_conversation_client = None
+    if app_settings.chat_history:
+        try:
+            cosmos_endpoint = (
+                f"https://{app_settings.chat_history.account}.documents.azure.com:443/"
+            )
+
+            if not app_settings.chat_history.account_key:
+                async with DefaultAzureCredential() as cred:
+                    credential = cred
+
+            else:
+                credential = app_settings.chat_history.account_key
+
+            cosmos_conversation_client = CosmosConversationClient(
+                cosmosdb_endpoint=cosmos_endpoint,
+                credential=credential,
+                database_name=app_settings.chat_history.database,
+                container_name=app_settings.chat_history.conversations_container,
+                enable_message_feedback=app_settings.chat_history.enable_feedback,
+            )
+        except Exception as e:
+            logging.exception("Exception in CosmosDB initialization", e)
+            cosmos_conversation_client = None
+            raise e
+    else:
+        logging.debug("CosmosDB not configured")
+
+    return cosmos_conversation_client
+
+def prepare_model_args(request_body, request_headers):
+    request_messages = request_body.get("messages", [])
+    messages = []
+    if not app_settings.datasource:
+        messages = [
+            {
+                "role": "system",
+                "content": app_settings.azure_openai.system_message
+            }
+        ]
+
+    for message in request_messages:
+        if message:
+            messages.append(
+                {
+                    "role": message["role"],
+                    "content": message["content"]
+                }
+            )
+
+    user_json = None
+    if (MS_DEFENDER_ENABLED):
+        authenticated_user_details = get_authenticated_user_details(request_headers)
+        conversation_id = request_body.get("conversation_id", None)        
+        user_json = get_msdefender_user_json(authenticated_user_details, request_headers, conversation_id)
+
+    model_args = {
+        "messages": messages,
+        "temperature": app_settings.azure_openai.temperature,
+        "max_tokens": app_settings.azure_openai.max_tokens,
+        "top_p": app_settings.azure_openai.top_p,
+        "stop": app_settings.azure_openai.stop_sequence,
+        "stream": app_settings.azure_openai.stream,
+        "model": app_settings.azure_openai.model,
+        "user": user_json
+    }
+
+    if app_settings.datasource:
+        model_args["extra_body"] = {
+            "data_sources": [
+                app_settings.datasource.construct_payload_configuration(
+                    request=request
+                )
+            ]
+        }
+
+    return model_args
+
+def create_app():
+    app = Quart(__name__)
+    app.register_blueprint(bp)
+    app.config["TEMPLATES_AUTO_RELOAD"] = True
+
+    @app.before_serving
+    async def init():
+        try:
+            app.cosmos_conversation_client = await init_cosmosdb_client()
+            cosmos_db_ready.set()
+            app.openai_client = await init_openai_client()
+        except Exception as e:
+            logging.exception("Failed to initialize services")
+            app.cosmos_conversation_client = None
+            app.openai_client = None
+            raise e
+
+    return app
+
+@bp.route("/")
+async def index():
+    return await render_template(
+        "index.html",
+        title=app_settings.ui.title,
+        favicon=app_settings.ui.favicon
+    )
+
+@bp.route("/favicon.ico")
+async def favicon():
+    return await bp.send_static_file("favicon.ico")
+
+@bp.route("/assets/<path:path>")
+async def assets(path):
+    return await send_from_directory("static/assets", path)
+
+@bp.route("/ai/ask", methods=["POST"])
+async def ask_ai():
+    data = await request.json
+    question = data.get("question", "")
+
+    if not question:
+        return jsonify({"error": "No question provided"}), 400
+
+    try:
+        openai_client = current_app.openai_client
+        response = await openai_client.completions.create(
+            model="text-davinci-003",
+            prompt=question,
+            max_tokens=150
+        )
+        answer = response.choices[0].text.strip()
+        return jsonify({"answer": answer}), 200
+    except Exception as e:
+        logging.exception("Error while processing AI request")
+        return jsonify({"error": "Failed to get response from AI service"}), 500