Gradio + DBRX

TODO: link to github

app.py

import gradio as gr
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()

def respond(message, history):

    if len(message.strip()) == 0:
        return "ERROR the question should not be empty"

    try:
        response = w.serving_endpoints.query(
            name="databricks-dbrx-instruct", 
            messages=[ChatMessage(content=message, role=ChatMessageRole.USER)]
        )
    except Exception as error:
        response = f"ERROR status_code: {type(error).__name__}"


    return response.choices[0].message.content


demo = gr.ChatInterface(
    respond,
    chatbot=gr.Chatbot(show_label=False, container=False, show_copy_button=True, bubble_full_width=True),
    textbox=gr.Textbox(placeholder="Ask me a question",
                       container=False, scale=7),
    title="Databricks App demo - Chat with DBRX Databricks model serving endpoint",
    examples=[["What is DBRX?"],
              ["How can I start a Databricks cluster?"],
              ["What is a Databricks Cluster Policy?"],
              ["How can I track billing usage on my workspaces?"],],
    cache_examples=False,
    # theme=gr.themes.Monochrome(),
    retry_btn=None,
    undo_btn=None,
    clear_btn="Clear",
)
demo.launch()