LAB.01 – Build an Internal Agent App with Databricks Apps.txt

import gradio as gr
import logging
from model_serving_utils import (
    endpoint_supports_feedback, 
    query_endpoint, 
    _get_endpoint_task_type,
)
import os

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Ensure environment variable is set correctly
SERVING_ENDPOINT = os.getenv('SERVING_ENDPOINT')
assert SERVING_ENDPOINT, "Unable to determine serving endpoint. Set SERVING_ENDPOINT environment variable."

ENDPOINT_SUPPORTS_FEEDBACK = endpoint_supports_feedback(SERVING_ENDPOINT)

def query_llm(message, history):
    if not message.strip():
        return "ERROR: The question should not be empty"

    # Convert Gradio history to OpenAI-style
    message_history = []
    for user_msg, assistant_msg in history:
        message_history.append({"role": "user", "content": user_msg})
        message_history.append({"role": "assistant", "content": assistant_msg})

    message_history.append({"role": "user", "content": message})

    try:
        logger.info(f"Sending request to Brixo Marketing Agent: {SERVING_ENDPOINT}")
        messages, request_id = query_endpoint(
            endpoint_name=SERVING_ENDPOINT,
            messages=message_history,
            return_traces=ENDPOINT_SUPPORTS_FEEDBACK
        )
        return messages[-1]
    except Exception as e:
        logger.error(f"Error querying model: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"

# Building the Dashboard UI based on the attachment
with gr.Blocks(title="Brixo Marketing Agent Dashboard") as demo:
    gr.Markdown("# Brixo Marketing Agent")
    gr.Markdown(
        "This agent helps you as a marketing agent. It can answer questions about our sales and stores. "
        "Also, you can ask the agent to generate custom marketing content based on sales and reviews."
    )
    
    with gr.Row():
        with gr.Column(scale=3):
            gr.ChatInterface(
                fn=query_llm,
                examples=[
                    "Write an Instagram message for the customers of my Seattle store.",
                    "Which cookies are best sellers in Seattle?",
                    "How many stores do we have in Seattle?"
                ],
                type="messages" # Modern Gradio chat format
            )
        
        # Adding a placeholder for the "Sales Data" sidebar seen in your image
        with gr.Column(scale=1):
            gr.Markdown("### Sales Data")
            gr.DataFrame(
                value=[["Seattle", 1200], ["Portland", 950], ["San Francisco", 1800]],
                headers=["Store", "Daily Sales ($)"],
                interactive=False
            )

if __name__ == "__main__":
    demo.launch()