Agent UI 예제

databricks-sdk
streamlit

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

import os
import streamlit as st
from databricks.sdk import WorkspaceClient

# Databricks 클라이언트 초기화 (인증 자동 처리)
w = WorkspaceClient()
endpoint_name = os.getenv("SERVING_ENDPOINT")

st.title("AI Assistant")
st.caption("Databricks Model Serving 기반 AI 챗봇")

# 세션 상태에 대화 이력 저장
if "messages" not in st.session_state:
    st.session_state.messages = []

# 기존 대화 이력 표시
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 사용자 입력 처리
if prompt := st.chat_input("질문을 입력하세요..."):
    # 사용자 메시지 추가
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # AI 응답 생성
    with st.chat_message("assistant"):
        with st.spinner("생각 중..."):
            try:
                response = w.serving_endpoints.query(
                    name=endpoint_name,
                    messages=[
                        {"role": m["role"], "content": m["content"]}
                        for m in st.session_state.messages
                    ],
                )
                assistant_message = response.choices[0].message.content
                st.markdown(assistant_message)
                st.session_state.messages.append(
                    {"role": "assistant", "content": assistant_message}
                )
            except Exception as e:
                st.error(f"오류 발생: {str(e)}")