"""
Databricks MCP Server - Streamlit Demo Chat Application
This demo application shows how to integrate the Databricks MCP Server
with Claude AI to create an intelligent Databricks workspace assistant.
"""
import os
import json
import asyncio
import streamlit as st
from anthropic import Anthropic
from dotenv import load_dotenv
import subprocess
import sys
# Load environment variables
load_dotenv()
# Page configuration
st.set_page_config(
page_title="Databricks AI Assistant",
page_icon="🤖",
layout="wide",
initial_sidebar_state="expanded"
)
# Custom CSS for better UI
st.markdown("""
<style>
.main-header {
font-size: 2.5rem;
color: #FF3621;
font-weight: bold;
margin-bottom: 0.5rem;
}
.sub-header {
font-size: 1.2rem;
color: #666;
margin-bottom: 2rem;
}
.chat-message {
padding: 1rem;
border-radius: 0.5rem;
margin-bottom: 1rem;
}
.user-message {
background-color: #E8F4F8;
}
.assistant-message {
background-color: #F5F5F5;
}
.tool-call {
background-color: #FFF8E1;
padding: 0.5rem;
border-radius: 0.3rem;
margin: 0.5rem 0;
border-left: 3px solid #FFC107;
}
.example-query {
background-color: #F0F7FF;
padding: 0.8rem;
border-radius: 0.3rem;
margin: 0.5rem 0;
cursor: pointer;
border: 1px solid #BBDEFB;
}
.example-query:hover {
background-color: #E3F2FD;
border-color: #2196F3;
}
</style>
""", unsafe_allow_html=True)
# Initialize session state
if "messages" not in st.session_state:
st.session_state.messages = []
if "anthropic_client" not in st.session_state:
api_key = os.getenv("ANTHROPIC_API_KEY")
if api_key:
st.session_state.anthropic_client = Anthropic(api_key=api_key)
else:
st.session_state.anthropic_client = None
# Sidebar
with st.sidebar:
st.markdown("## 🛠️ Configuration")
# API Key input
if not st.session_state.anthropic_client:
api_key = st.text_input("Anthropic API Key", type="password")
if api_key:
st.session_state.anthropic_client = Anthropic(api_key=api_key)
st.success("✅ API Key configured!")
else:
st.success("✅ API Key configured!")
st.markdown("---")
# Databricks configuration
st.markdown("### Databricks Settings")
databricks_host = os.getenv("DATABRICKS_HOST", "")
st.text_input("Workspace URL", value=databricks_host[:30] + "..." if len(databricks_host) > 30 else databricks_host, disabled=True)
if os.getenv("DATABRICKS_TOKEN"):
st.success("✅ Databricks token configured")
else:
st.error("❌ Databricks token not configured")
st.markdown("---")
# Example queries
st.markdown("### 💡 Example Queries")
examples = [
"📊 List all running clusters",
"💰 Show me last month's costs",
"📔 List my notebooks",
"⚙️ What jobs are scheduled?",
"🔍 Show optimization tips",
"📈 Analyze spending trends",
"🗄️ List Unity Catalog tables",
"👥 Who has admin access?"
]
for example in examples:
if st.button(example, key=f"example_{example}", use_container_width=True):
# Extract text without emoji
query_text = " ".join(example.split()[1:])
st.session_state.pending_query = query_text
st.markdown("---")
# Settings
st.markdown("### ⚙️ Settings")
temperature = st.slider("Temperature", 0.0, 1.0, 0.7, 0.1)
max_tokens = st.slider("Max Tokens", 1024, 4096, 2048, 128)
st.markdown("---")
# Clear chat button
if st.button("🗑️ Clear Chat", use_container_width=True):
st.session_state.messages = []
st.rerun()
st.markdown("---")
# Info
st.markdown("### ℹ️ About")
st.markdown("""
This demo uses:
- **Claude AI** via Anthropic API
- **Databricks MCP Server** for workspace operations
- **Streamlit** for the UI
The AI can manage clusters, jobs, notebooks, analyze costs, and more!
""")
# Main content
st.markdown('<div class="main-header">🤖 Databricks AI Assistant</div>', unsafe_allow_html=True)
st.markdown('<div class="sub-header">Powered by Claude AI and Model Context Protocol</div>', unsafe_allow_html=True)
# Check if configuration is complete
if not st.session_state.anthropic_client:
st.warning("⚠️ Please configure your Anthropic API key in the sidebar to get started.")
st.stop()
if not os.getenv("DATABRICKS_HOST") or not os.getenv("DATABRICKS_TOKEN"):
st.warning("⚠️ Please configure Databricks credentials in your .env file.")
st.info("""
Create a `.env` file with:
```
DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
DATABRICKS_TOKEN=your-databricks-token
```
""")
st.stop()
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Show tool calls if present
if "tool_calls" in message and message["tool_calls"]:
with st.expander("🔧 Tool Calls", expanded=False):
for tool_call in message["tool_calls"]:
st.markdown(f"""
<div class="tool-call">
<strong>🛠️ {tool_call['name']}</strong><br>
<small>{json.dumps(tool_call.get('input', {}), indent=2)}</small>
</div>
""", unsafe_allow_html=True)
# Handle pending query from example buttons
if hasattr(st.session_state, 'pending_query'):
user_message = st.session_state.pending_query
delattr(st.session_state, 'pending_query')
else:
# Chat input
user_message = st.chat_input("Ask me anything about your Databricks workspace...")
if user_message:
# Add user message to chat
st.session_state.messages.append({"role": "user", "content": user_message})
with st.chat_message("user"):
st.markdown(user_message)
# Show thinking indicator
with st.chat_message("assistant"):
with st.spinner("🤔 Thinking and querying Databricks..."):
try:
# Note: This is a simplified version. In a full implementation,
# you would integrate with the MCP server using the MCP SDK
# For this demo, we'll simulate the interaction
# Simulate calling Claude with MCP tools
response_text = """
I'll help you with that! Let me check your Databricks workspace.
**Note**: This is a demo interface. To fully integrate with the Databricks MCP Server:
1. Install the MCP server: `npm install -g databricks-mcp-server`
2. Configure it with your credentials
3. Use the MCP SDK to connect Claude to the server
The MCP server provides these capabilities:
- 🖥️ **Cluster Management**: Create, start, stop, resize clusters
- 📔 **Notebook Operations**: List, create, run notebooks
- ⚙️ **Job Management**: Schedule and monitor jobs
- 💰 **Cost Analytics**: Track spending and get optimization tips
- 🗄️ **Unity Catalog**: Manage data and permissions
- 👥 **User Management**: Control access and permissions
For your query: "{}"
In a full implementation, I would:
1. Connect to your Databricks workspace via MCP
2. Execute the appropriate API calls
3. Return structured results
4. Provide insights and recommendations
Try these example queries:
- "List all clusters and show which ones are idle"
- "What were my costs last month and how can I optimize?"
- "Create a new cluster for machine learning"
- "Show me all scheduled jobs and their statuses"
""".format(user_message)
# Add assistant response
st.session_state.messages.append({
"role": "assistant",
"content": response_text,
"tool_calls": []
})
st.markdown(response_text)
except Exception as e:
st.error(f"Error: {str(e)}")
st.info("Make sure the Databricks MCP Server is running and configured correctly.")
# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666; font-size: 0.9rem;'>
<p>Built with ❤️ using Streamlit, Claude AI, and Databricks MCP Server</p>
<p>
<a href='https://github.com/yourusername/databricks-mcp-server' target='_blank'>GitHub</a> •
<a href='https://docs.anthropic.com/claude/docs' target='_blank'>Claude Docs</a> •
<a href='https://modelcontextprotocol.io' target='_blank'>MCP Docs</a>
</p>
</div>
""", unsafe_allow_html=True)
