Gradio 챗봇 예제

databricks-sdk
gradio

command: ['python', 'app.py']
env:
  - name: SERVING_ENDPOINT
    valueFrom: serving_endpoint
resources:
  - name: serving_endpoint
    type: serving-endpoint

import os
import gradio as gr
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
endpoint_name = os.getenv("SERVING_ENDPOINT")
port = int(os.environ.get("DATABRICKS_APP_PORT", 7860))


def chat(message, history, request: gr.Request):
    """스트리밍 응답을 지원하는 챗 함수"""
    # 대화 이력을 API 형식으로 변환
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    # 스트리밍 응답
    response = w.serving_endpoints.query(
        name=endpoint_name,
        messages=messages,
        stream=True,
    )

    partial_message = ""
    for chunk in response:
        if chunk.choices and chunk.choices[0].delta.content:
            partial_message += chunk.choices[0].delta.content
            yield partial_message


# Gradio 챗봇 UI
demo = gr.ChatInterface(
    fn=chat,
    title="AI Assistant",
    description="Databricks Foundation Model API 기반 AI 챗봇",
    examples=["데이터 레이크하우스란 무엇인가요?", "Delta Lake의 장점을 설명해주세요"],
    theme=gr.themes.Soft(),
)

demo.launch(server_name="0.0.0.0", server_port=port)