Tutorial 17: Copilot & AI Integration¶

# portal/backend/api/chat.py (reference pattern)
from fastapi import APIRouter, Depends
from openai import AsyncAzureOpenAI
import os

router = APIRouter()

client = AsyncAzureOpenAI(
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
    api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-08-06"),
)

@router.post("/api/chat")
async def chat(request: dict):
    messages = request.get("messages", [])

    # Add system prompt with platform context
    system_message = {
        "role": "system",
        "content": (
            "You are an AI assistant for the CSA-in-a-Box data platform. "
            "You help users understand their data, write queries, and "
            "interpret analytics results. Be concise and accurate."
        ),
    }

    response = await client.chat.completions.create(
        model=os.environ["AZURE_OPENAI_DEPLOYMENT"],
        messages=[system_message] + messages,
        temperature=0.7,
        max_tokens=1024,
    )

    return {"reply": response.choices[0].message.content}

# scripts/create_search_index.py
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SimpleField,
    SearchableField,
    SearchFieldDataType,
)
from azure.core.credentials import AzureKeyCredential
import os

client = SearchIndexClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)

index = SearchIndex(
    name="datalake-metadata",
    fields=[
        SimpleField(name="id", type=SearchFieldDataType.String, key=True),
        SearchableField(name="table_name", type=SearchFieldDataType.String),
        SearchableField(name="layer", type=SearchFieldDataType.String, filterable=True),
        SearchableField(name="domain", type=SearchFieldDataType.String, filterable=True),
        SearchableField(name="description", type=SearchFieldDataType.String),
        SearchableField(name="columns", type=SearchFieldDataType.String),
        SimpleField(name="row_count", type=SearchFieldDataType.Int64),
        SimpleField(name="last_updated", type=SearchFieldDataType.DateTimeOffset),
    ],
)

client.create_or_update_index(index)
print(f"Index '{index.name}' created.")

# scripts/index_dbt_metadata.py
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
import json, os, glob

client = SearchClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    index_name="datalake-metadata",
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)

# Parse dbt manifest for model metadata
manifest_path = "target/manifest.json"
with open(manifest_path) as f:
    manifest = json.load(f)

documents = []
for node_id, node in manifest.get("nodes", {}).items():
    if node["resource_type"] != "model":
        continue
    columns = ", ".join(
        [f"{c['name']} ({c.get('description', 'N/A')})"
         for c in node.get("columns", {}).values()]
    )
    documents.append({
        "id": node_id.replace(".", "-"),
        "table_name": node["name"],
        "layer": node.get("tags", ["unknown"])[0] if node.get("tags") else "unknown",
        "domain": node.get("fqn", [""])[1] if len(node.get("fqn", [])) > 1 else "",
        "description": node.get("description", ""),
        "columns": columns,
        "row_count": 0,
    })

result = client.upload_documents(documents)
print(f"Indexed {len(documents)} models.")

# portal/backend/api/chat_grounded.py
from fastapi import APIRouter
from openai import AsyncAzureOpenAI
from azure.search.documents.aio import SearchClient
from azure.core.credentials import AzureKeyCredential
import os

router = APIRouter()

openai_client = AsyncAzureOpenAI(
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    api_version="2024-08-06",
)

search_client = SearchClient(
    endpoint=os.environ["SEARCH_ENDPOINT"],
    index_name="datalake-metadata",
    credential=AzureKeyCredential(os.environ["SEARCH_KEY"]),
)


@router.post("/api/chat/grounded")
async def grounded_chat(request: dict):
    user_message = request["messages"][-1]["content"]

    # 1. Retrieve relevant context from AI Search
    search_results = search_client.search(
        search_text=user_message,
        top=5,
        select=["table_name", "layer", "domain", "description", "columns"],
    )

    context_chunks = []
    async for result in search_results:
        context_chunks.append(
            f"Table: {result['table_name']} (Layer: {result['layer']}, "
            f"Domain: {result['domain']})\n"
            f"Description: {result['description']}\n"
            f"Columns: {result['columns']}"
        )

    grounding_context = "\n\n".join(context_chunks)

    # 2. Build augmented prompt
    system_message = {
        "role": "system",
        "content": (
            "You are an AI assistant for the CSA-in-a-Box data platform. "
            "Answer the user's question using ONLY the following data catalog context. "
            "If the context does not contain the answer, say so.\n\n"
            f"## Data Catalog Context\n{grounding_context}"
        ),
    }

    # 3. Send to GPT-4o
    response = await openai_client.chat.completions.create(
        model=os.environ["AZURE_OPENAI_DEPLOYMENT"],
        messages=[system_message] + request["messages"],
        temperature=0.3,  # Lower temperature for factual answers
        max_tokens=1024,
    )

    return {
        "reply": response.choices[0].message.content,
        "sources": [c.split("\n")[0] for c in context_chunks],
    }

"""
Custom data copilot built with Semantic Kernel.
References: examples/ai-agents/data-analyst-agent/agent.py
"""

import asyncio
import os
from semantic_kernel import Kernel
from semantic_kernel.connectors.ai.open_ai import AzureChatCompletion
from semantic_kernel.contents import ChatHistory
from semantic_kernel.functions import kernel_function


class DataPlatformPlugin:
    """Plugin that gives the agent access to data platform queries."""

    @kernel_function(
        name="query_gold_layer",
        description="Execute a SQL query against the Gold layer tables in Databricks.",
    )
    def query_gold_layer(self, query: str) -> str:
        """Execute a query against the Gold layer."""
        # In production, use Databricks SQL connector
        # For tutorial, return mock results
        return (
            f"Query executed: {query}\n"
            "Results: 42 rows returned.\n"
            "Sample: order_date=2024-01-15, region=West, total_revenue=14995.00"
        )

    @kernel_function(
        name="list_available_tables",
        description="List all available tables in the Gold layer.",
    )
    def list_available_tables(self) -> str:
        return (
            "Available Gold layer tables:\n"
            "- rpt_daily_sales (order_date, region, total_orders, total_revenue)\n"
            "- rpt_product_performance (product_name, category, total_units_sold)\n"
            "- dim_customers (customer_id, customer_name, region, customer_tier)\n"
            "- dim_products (product_id, product_name, category, price_tier)"
        )

    @kernel_function(
        name="assess_data_quality",
        description="Check data quality metrics for a given table.",
    )
    def assess_data_quality(self, table_name: str) -> str:
        return (
            f"Data quality report for {table_name}:\n"
            "- Row count: 1,247\n"
            "- Null percentage: 0.3%\n"
            "- Last refreshed: 2024-01-22 10:00 UTC\n"
            "- All schema tests passing"
        )


async def main():
    kernel = Kernel()

    # Add Azure OpenAI chat completion
    kernel.add_service(
        AzureChatCompletion(
            deployment_name=os.environ["AZURE_OPENAI_DEPLOYMENT"],
            endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
        )
    )

    # Add the data platform plugin
    kernel.add_plugin(DataPlatformPlugin(), plugin_name="data_platform")

    # Create chat history
    history = ChatHistory()
    history.add_system_message(
        "You are a data analyst assistant for the CSA-in-a-Box platform. "
        "Use the data_platform plugin to query tables, list available data, "
        "and check data quality. Always ground your answers in actual data."
    )

    print("Data Copilot ready. Type 'quit' to exit.\n")

    while True:
        user_input = input("You: ")
        if user_input.lower() == "quit":
            break

        history.add_user_message(user_input)

        result = await kernel.invoke_prompt(
            prompt="{{$history}}",
            history=history,
        )

        print(f"Agent: {result}\n")
        history.add_assistant_message(str(result))


if __name__ == "__main__":
    asyncio.run(main())

# portal/backend/middleware/content_safety.py
from azure.ai.contentsafety import ContentSafetyClient
from azure.ai.contentsafety.models import AnalyzeTextOptions, TextCategory
from azure.core.credentials import AzureKeyCredential
import os

safety_client = ContentSafetyClient(
    endpoint=os.environ["SAFETY_ENDPOINT"],
    credential=AzureKeyCredential(os.environ["SAFETY_KEY"]),
)

SEVERITY_THRESHOLD = 2  # 0=safe, 2=low, 4=medium, 6=high


def check_content_safety(text: str) -> dict:
    """Screen text for harmful content. Returns pass/fail with details."""
    result = safety_client.analyze_text(
        AnalyzeTextOptions(text=text)
    )

    categories_flagged = []
    for category in result.categories_analysis:
        if category.severity >= SEVERITY_THRESHOLD:
            categories_flagged.append({
                "category": category.category,
                "severity": category.severity,
            })

    return {
        "safe": len(categories_flagged) == 0,
        "flagged_categories": categories_flagged,
    }

@router.post("/api/chat/safe")
async def safe_chat(request: dict):
    user_message = request["messages"][-1]["content"]

    # Screen the user input
    input_check = check_content_safety(user_message)
    if not input_check["safe"]:
        return {
            "reply": "I'm unable to process that request. Please rephrase your question.",
            "safety_flag": input_check["flagged_categories"],
        }

    # Get AI response (using grounded chat from Step 4)
    response = await grounded_chat(request)

    # Screen the AI output
    output_check = check_content_safety(response["reply"])
    if not output_check["safe"]:
        return {
            "reply": "The response was filtered for safety. Please try a different question.",
            "safety_flag": output_check["flagged_categories"],
        }

    return response

# portal/backend/middleware/prompt_logger.py
import json
import logging
from datetime import datetime, timezone
from azure.monitor.opentelemetry import configure_azure_monitor
from opentelemetry import trace

# Configure Application Insights (if available)
# configure_azure_monitor(connection_string=os.environ.get("APPINSIGHTS_CONNECTION_STRING"))

logger = logging.getLogger("llmops")
tracer = trace.get_tracer(__name__)


def log_prompt_response(
    user_message: str,
    system_prompt: str,
    response: str,
    model: str,
    tokens_prompt: int,
    tokens_completion: int,
    latency_ms: float,
    grounding_sources: list[str] | None = None,
):
    """Log prompt/response pairs for LLMOps monitoring."""
    entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "model": model,
        "user_message_length": len(user_message),
        "response_length": len(response),
        "tokens_prompt": tokens_prompt,
        "tokens_completion": tokens_completion,
        "tokens_total": tokens_prompt + tokens_completion,
        "latency_ms": latency_ms,
        "grounding_sources": grounding_sources or [],
    }
    logger.info("llm_interaction", extra=entry)

    # Also emit as OpenTelemetry span
    with tracer.start_as_current_span("llm.chat") as span:
        span.set_attribute("llm.model", model)
        span.set_attribute("llm.tokens.prompt", tokens_prompt)
        span.set_attribute("llm.tokens.completion", tokens_completion)
        span.set_attribute("llm.latency_ms", latency_ms)

import time

@router.post("/api/chat/grounded")
async def grounded_chat_with_logging(request: dict):
    start = time.monotonic()
    user_message = request["messages"][-1]["content"]

    # ... (search + prompt construction from Step 4) ...

    response = await openai_client.chat.completions.create(
        model=os.environ["AZURE_OPENAI_DEPLOYMENT"],
        messages=[system_message] + request["messages"],
        temperature=0.3,
        max_tokens=1024,
    )

    latency_ms = (time.monotonic() - start) * 1000

    # Log the interaction
    log_prompt_response(
        user_message=user_message,
        system_prompt=system_message["content"],
        response=response.choices[0].message.content,
        model=response.model,
        tokens_prompt=response.usage.prompt_tokens,
        tokens_completion=response.usage.completion_tokens,
        latency_ms=latency_ms,
        grounding_sources=[c.split("\n")[0] for c in context_chunks],
    )

    return {
        "reply": response.choices[0].message.content,
        "sources": [c.split("\n")[0] for c in context_chunks],
        "usage": {
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens,
        },
    }