"""
DS-STAR Gradio Application
A modern web interface for the DS-STAR Multi-Agent Data Science System.
Created for the Hugging Face MCP 1st Birthday Hackathon.
"""
import os
import shutil
from typing import Generator
import gradio as gr
import httpx
from src.config import get_llm
from src.graph import build_ds_star_graph, create_initial_state
# ==================== MODEL FETCHING ====================
# Store fetch status for UI feedback
_last_fetch_status = {"success": False, "message": "", "from_api": False}
def fetch_google_models(api_key: str | None = None) -> tuple[list[str], str]:
"""Fetch available models from Google Gemini API. Returns (models, status_message)."""
api_key = api_key or os.getenv("GOOGLE_API_KEY", "")
fallback = [
"gemini-2.0-flash",
"gemini-1.5-pro",
"gemini-1.5-flash",
"gemini-1.0-pro",
]
if not api_key:
return fallback, "⚠️ No API key - showing default models"
try:
url = f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}"
response = httpx.get(url, timeout=15)
if response.status_code == 200:
data = response.json()
models = []
for model in data.get("models", []):
name = model.get("name", "").replace("models/", "")
# Filter for chat/generate models
if "generateContent" in model.get("supportedGenerationMethods", []):
models.append(name)
if models:
return sorted(
models, reverse=True
), f"✅ Fetched {len(models)} models from API"
return fallback, "⚠️ No compatible models found - showing defaults"
elif response.status_code == 400:
return fallback, "❌ Invalid API key format"
elif response.status_code == 403:
return fallback, "❌ API key invalid or expired"
else:
return fallback, f"❌ API error: {response.status_code}"
except httpx.TimeoutException:
return fallback, "❌ Request timed out"
except httpx.ConnectError:
return fallback, "❌ Connection failed - check internet"
except Exception as e:
return fallback, f"❌ Error: {str(e)[:50]}"
def fetch_openai_models(
api_key: str | None = None, base_url: str | None = None
) -> tuple[list[str], str]:
"""Fetch available models from OpenAI API or compatible endpoint. Returns (models, status_message)."""
api_key = api_key or os.getenv("OPENAI_API_KEY", "")
base_url = base_url or "https://api.openai.com/v1"
fallback = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
if not api_key:
return fallback, "⚠️ No API key - showing default models"
try:
headers = {"Authorization": f"Bearer {api_key}"}
endpoint = f"{base_url.rstrip('/')}/models"
response = httpx.get(endpoint, headers=headers, timeout=15)
if response.status_code == 200:
data = response.json()
models = []
for model in data.get("data", []):
model_id = model.get("id", "")
# For OpenAI, filter chat models; for custom endpoints, include all
if base_url != "https://api.openai.com/v1":
models.append(model_id)
elif model_id.startswith(("gpt-", "o1", "o3", "o4", "chatgpt")):
models.append(model_id)
if models:
return sorted(
models, reverse=True
), f"✅ Fetched {len(models)} models from API"
return fallback, "⚠️ No chat models found - showing defaults"
elif response.status_code == 401:
return fallback, "❌ Invalid API key"
elif response.status_code == 403:
return fallback, "❌ Access denied"
else:
return fallback, f"❌ API error: {response.status_code}"
except httpx.TimeoutException:
return fallback, "❌ Request timed out"
except httpx.ConnectError:
return fallback, "❌ Connection failed - check URL/internet"
except Exception as e:
return fallback, f"❌ Error: {str(e)[:50]}"
def fetch_anthropic_models(api_key: str | None = None) -> tuple[list[str], str]:
"""Fetch available models from Anthropic API. Returns (models, status_message)."""
api_key = api_key or os.getenv("ANTHROPIC_API_KEY", "")
fallback = [
"claude-sonnet-4-20250514",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022",
"claude-3-opus-20240229",
]
if not api_key:
return fallback, "⚠️ No API key - showing default models"
try:
headers = {"x-api-key": api_key, "anthropic-version": "2023-06-01"}
response = httpx.get(
"https://api.anthropic.com/v1/models", headers=headers, timeout=15
)
if response.status_code == 200:
data = response.json()
models = [
model.get("id", "") for model in data.get("data", []) if model.get("id")
]
if models:
return sorted(
models, reverse=True
), f"✅ Fetched {len(models)} models from API"
return fallback, "⚠️ No models found - showing defaults"
elif response.status_code == 401:
return fallback, "❌ Invalid API key"
elif response.status_code == 403:
return fallback, "❌ Access denied"
else:
# Anthropic may not have a public models endpoint, use fallback
return fallback, "ℹ️ Using known Anthropic models"
except httpx.TimeoutException:
return fallback, "❌ Request timed out"
except httpx.ConnectError:
return fallback, "❌ Connection failed"
except Exception as e:
return fallback, f"❌ Error: {str(e)[:50]}"
def fetch_models_for_provider(
provider: str, api_key: str | None = None, base_url: str | None = None
) -> tuple[list[str], str]:
"""Fetch models for the given provider. Returns (models, status_message)."""
if provider == "google":
return fetch_google_models(api_key)
elif provider == "openai":
return fetch_openai_models(api_key)
elif provider == "anthropic":
return fetch_anthropic_models(api_key)
elif provider == "custom":
if not base_url:
return [], "❌ Please enter a Base URL for custom provider"
return fetch_openai_models(api_key, base_url)
return [], "❌ Unknown provider"
# ==================== CUSTOM THEME ====================
def create_ds_star_theme():
"""Create a modern dark theme for DS-STAR."""
return gr.themes.Base(
primary_hue=gr.themes.colors.violet,
secondary_hue=gr.themes.colors.purple,
neutral_hue=gr.themes.colors.slate,
font=[
gr.themes.GoogleFont("Inter"),
"ui-sans-serif",
"system-ui",
"sans-serif",
],
font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
# Body - Dark background
body_background_fill="#0a0a0f",
body_background_fill_dark="#0a0a0f",
body_text_color="#e4e4e7",
body_text_color_dark="#e4e4e7",
# Buttons
button_primary_background_fill="linear-gradient(135deg, #7c3aed 0%, #8b5cf6 100%)",
button_primary_background_fill_hover="linear-gradient(135deg, #6d28d9 0%, #7c3aed 100%)",
button_primary_text_color="white",
button_primary_border_color="transparent",
button_secondary_background_fill="transparent",
button_secondary_background_fill_hover="rgba(139, 92, 246, 0.15)",
button_secondary_border_color="#7c3aed",
button_secondary_text_color="#a78bfa",
# Blocks
block_background_fill="#18181b",
block_background_fill_dark="#18181b",
block_border_width="1px",
block_border_color="#27272a",
block_border_color_dark="#27272a",
block_shadow="none",
block_title_text_weight="600",
block_title_text_size="*text_md",
block_label_text_weight="500",
block_label_text_size="*text_sm",
block_radius="12px",
block_padding="16px",
# Inputs
input_background_fill="#27272a",
input_background_fill_dark="#27272a",
input_border_color="#3f3f46",
input_border_color_dark="#3f3f46",
input_border_width="1px",
input_shadow="none",
input_radius="8px",
# Panels
panel_background_fill="#18181b",
panel_background_fill_dark="#18181b",
panel_border_width="0px",
# Spacing
layout_gap="16px",
# Shadows
shadow_drop="none",
shadow_drop_lg="none",
# Checkbox
checkbox_background_color="#27272a",
checkbox_background_color_dark="#27272a",
checkbox_border_color="#3f3f46",
checkbox_border_color_dark="#3f3f46",
checkbox_label_text_color="#a1a1aa",
checkbox_label_text_color_dark="#a1a1aa",
# Slider
slider_color="#7c3aed",
slider_color_dark="#8b5cf6",
)
# ==================== CSS STYLING ====================
CUSTOM_CSS = """
/* Modern Font Import */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600&display=swap');
/* Root variables */
:root {
--bg-primary: #0a0a0f;
--bg-secondary: #18181b;
--bg-tertiary: #27272a;
--border-color: #3f3f46;
--text-primary: #fafafa;
--text-secondary: #a1a1aa;
--text-muted: #71717a;
--accent-primary: #8b5cf6;
--accent-secondary: #7c3aed;
--accent-glow: rgba(139, 92, 246, 0.3);
--success: #22c55e;
--error: #ef4444;
}
/* Main container - Dark background */
.gradio-container {
max-width: 1400px !important;
margin: 0 auto !important;
padding: 32px !important;
font-family: 'Inter', sans-serif !important;
background: var(--bg-primary) !important;
min-height: 100vh;
}
/* Remove all default shadows and borders for cleaner look */
.gradio-container * {
box-shadow: none !important;
}
/* ===== HEADER ===== */
.header-section {
background: linear-gradient(135deg, #1e1b4b 0%, #312e81 50%, #3730a3 100%);
border-radius: 16px;
padding: 40px 32px;
margin-bottom: 24px;
border: 1px solid #4338ca;
position: relative;
overflow: hidden;
}
.header-section::before {
content: '';
position: absolute;
top: 0;
right: 0;
width: 40%;
height: 100%;
background: radial-gradient(ellipse at top right, rgba(139, 92, 246, 0.3), transparent 70%);
}
.header-content {
position: relative;
z-index: 1;
text-align: center;
}
.header-title {
font-size: 2.75rem;
font-weight: 800;
color: #fff;
margin: 0 0 8px 0;
letter-spacing: -0.02em;
display: flex;
align-items: center;
justify-content: center;
gap: 12px;
}
.header-title .star-icon {
font-size: 2.2rem;
}
.header-subtitle {
font-size: 1.1rem;
color: rgba(255, 255, 255, 0.7);
margin: 0;
font-weight: 400;
}
.header-badges {
display: flex;
justify-content: center;
gap: 10px;
margin-top: 20px;
flex-wrap: wrap;
}
.header-badge {
padding: 6px 14px;
background: rgba(255, 255, 255, 0.1);
border: 1px solid rgba(255, 255, 255, 0.15);
border-radius: 20px;
font-size: 0.8rem;
font-weight: 500;
color: rgba(255, 255, 255, 0.9);
}
/* ===== CARDS & GROUPS ===== */
.dark-card {
background: var(--bg-secondary) !important;
border: 1px solid var(--border-color) !important;
border-radius: 12px !important;
padding: 20px !important;
margin-bottom: 16px !important;
}
/* Group styling */
.gr-group {
background: var(--bg-secondary) !important;
border: 1px solid var(--border-color) !important;
border-radius: 12px !important;
padding: 20px !important;
}
/* ===== SECTION HEADERS ===== */
.section-header {
display: flex;
align-items: center;
gap: 10px;
margin-bottom: 16px;
padding-bottom: 12px;
border-bottom: 1px solid var(--border-color);
}
.section-icon {
width: 32px;
height: 32px;
display: flex;
align-items: center;
justify-content: center;
background: linear-gradient(135deg, var(--accent-secondary), var(--accent-primary));
border-radius: 8px;
font-size: 1rem;
}
.section-title {
font-size: 1.1rem;
font-weight: 600;
color: var(--text-primary);
margin: 0;
}
/* ===== ACCORDIONS ===== */
.gr-accordion {
border: 1px solid var(--border-color) !important;
border-radius: 10px !important;
overflow: hidden !important;
margin-bottom: 12px !important;
background: var(--bg-secondary) !important;
}
.gr-accordion > .label-wrap {
background: var(--bg-tertiary) !important;
padding: 12px 16px !important;
cursor: pointer !important;
border: none !important;
}
.gr-accordion > .label-wrap:hover {
background: #323238 !important;
}
.gr-accordion > .label-wrap span {
font-weight: 600 !important;
font-size: 0.95rem !important;
color: var(--text-primary) !important;
}
.gr-accordion > div:last-child {
padding: 16px !important;
background: var(--bg-secondary) !important;
}
/* ===== BUTTONS ===== */
button.primary, .gr-button-primary {
background: linear-gradient(135deg, var(--accent-secondary), var(--accent-primary)) !important;
border: none !important;
border-radius: 8px !important;
padding: 10px 20px !important;
font-weight: 600 !important;
font-size: 0.9rem !important;
color: white !important;
cursor: pointer !important;
transition: all 0.2s ease !important;
}
button.primary:hover, .gr-button-primary:hover {
opacity: 0.9 !important;
transform: translateY(-1px) !important;
}
button.secondary, .gr-button-secondary {
background: transparent !important;
border: 1px solid var(--accent-primary) !important;
border-radius: 8px !important;
color: var(--accent-primary) !important;
font-weight: 500 !important;
padding: 10px 20px !important;
transition: all 0.2s ease !important;
}
button.secondary:hover, .gr-button-secondary:hover {
background: rgba(139, 92, 246, 0.1) !important;
}
/* Refresh button */
.refresh-btn {
min-width: 40px !important;
width: 40px !important;
height: 40px !important;
padding: 0 !important;
display: flex !important;
align-items: center !important;
justify-content: center !important;
border-radius: 8px !important;
background: var(--bg-tertiary) !important;
border: 1px solid var(--border-color) !important;
font-size: 1rem !important;
color: var(--text-secondary) !important;
transition: all 0.2s ease !important;
}
.refresh-btn:hover {
background: var(--accent-primary) !important;
border-color: var(--accent-primary) !important;
color: white !important;
}
/* ===== INPUTS ===== */
input, textarea, select, .gr-input, .gr-text-input {
background: var(--bg-tertiary) !important;
border: 1px solid var(--border-color) !important;
border-radius: 8px !important;
color: var(--text-primary) !important;
padding: 10px 14px !important;
font-size: 0.9rem !important;
transition: border-color 0.2s ease !important;
}
input:focus, textarea:focus, select:focus {
border-color: var(--accent-primary) !important;
outline: none !important;
}
input::placeholder, textarea::placeholder {
color: var(--text-muted) !important;
}
/* Dropdown */
.gr-dropdown {
background: var(--bg-tertiary) !important;
}
/* ===== TABS ===== */
.gr-tabs {
background: transparent !important;
}
.gr-tab-nav {
background: var(--bg-secondary) !important;
border-radius: 10px !important;
padding: 4px !important;
margin-bottom: 16px !important;
border: 1px solid var(--border-color) !important;
gap: 4px !important;
display: flex !important;
overflow-x: auto !important;
}
.gr-tab-nav button {
border-radius: 8px !important;
padding: 10px 16px !important;
font-weight: 500 !important;
font-size: 0.85rem !important;
background: transparent !important;
border: none !important;
color: var(--text-secondary) !important;
white-space: nowrap !important;
transition: all 0.2s ease !important;
}
.gr-tab-nav button:hover {
background: var(--bg-tertiary) !important;
color: var(--text-primary) !important;
}
.gr-tab-nav button.selected {
background: var(--accent-primary) !important;
color: white !important;
}
/* ===== CODE OUTPUT ===== */
.code-wrap, .gr-code {
border-radius: 10px !important;
overflow: hidden !important;
border: 1px solid var(--border-color) !important;
}
.code-wrap pre, .gr-code pre {
background: #0d0d12 !important;
padding: 16px !important;
margin: 0 !important;
font-family: 'JetBrains Mono', monospace !important;
font-size: 0.85rem !important;
line-height: 1.5 !important;
max-height: 400px !important;
overflow: auto !important;
}
/* ===== FILE UPLOAD ===== */
.file-upload, .gr-file {
border: 2px dashed var(--border-color) !important;
border-radius: 10px !important;
background: var(--bg-tertiary) !important;
padding: 24px !important;
transition: all 0.2s ease !important;
}
.file-upload:hover, .gr-file:hover {
border-color: var(--accent-primary) !important;
}
/* ===== STATUS DISPLAYS ===== */
.status-box textarea {
font-weight: 600 !important;
color: var(--accent-primary) !important;
background: rgba(139, 92, 246, 0.1) !important;
border: 1px solid rgba(139, 92, 246, 0.3) !important;
border-radius: 8px !important;
}
.step-box textarea {
font-family: 'JetBrains Mono', monospace !important;
font-size: 0.85rem !important;
color: var(--text-secondary) !important;
background: var(--bg-tertiary) !important;
border-radius: 8px !important;
}
/* ===== EXAMPLES - Redesigned as chips ===== */
.gr-examples {
margin-top: 16px !important;
padding-top: 16px !important;
border-top: 1px solid var(--border-color) !important;
}
.gr-examples .label {
font-size: 0.85rem !important;
font-weight: 600 !important;
color: var(--text-secondary) !important;
margin-bottom: 12px !important;
}
.gr-examples-table {
display: flex !important;
flex-wrap: wrap !important;
gap: 8px !important;
}
.gr-examples-table tbody {
display: flex !important;
flex-wrap: wrap !important;
gap: 8px !important;
}
.gr-examples-table tr {
display: contents !important;
}
.gr-examples-table td {
display: block !important;
padding: 0 !important;
}
.gr-examples-table button, .gr-samples-table button {
background: var(--bg-tertiary) !important;
border: 1px solid var(--border-color) !important;
border-radius: 20px !important;
padding: 8px 16px !important;
font-size: 0.8rem !important;
color: var(--text-secondary) !important;
cursor: pointer !important;
transition: all 0.2s ease !important;
white-space: nowrap !important;
max-width: none !important;
}
.gr-examples-table button:hover, .gr-samples-table button:hover {
background: var(--accent-primary) !important;
border-color: var(--accent-primary) !important;
color: white !important;
}
/* ===== FILE LIST ===== */
.file-list {
background: var(--bg-tertiary) !important;
border-radius: 8px !important;
padding: 12px !important;
font-size: 0.85rem !important;
color: var(--text-secondary) !important;
max-height: 120px !important;
overflow-y: auto !important;
}
/* ===== WORKFLOW STEPS ===== */
.workflow-container {
display: grid;
grid-template-columns: repeat(2, 1fr);
gap: 12px;
padding: 16px;
}
.workflow-step {
display: flex;
align-items: center;
gap: 12px;
padding: 12px 14px;
background: var(--bg-tertiary);
border: 1px solid var(--border-color);
border-radius: 10px;
transition: all 0.2s ease;
}
.workflow-step:hover {
border-color: var(--accent-primary);
}
.step-number {
width: 28px;
height: 28px;
display: flex;
align-items: center;
justify-content: center;
background: linear-gradient(135deg, var(--accent-secondary), var(--accent-primary));
border-radius: 6px;
color: white;
font-weight: 700;
font-size: 0.8rem;
flex-shrink: 0;
}
.step-content {
flex: 1;
min-width: 0;
}
.step-title {
font-weight: 600;
color: var(--text-primary);
font-size: 0.85rem;
}
.step-desc {
font-size: 0.75rem;
color: var(--text-muted);
margin-top: 2px;
white-space: nowrap;
overflow: hidden;
text-overflow: ellipsis;
}
/* ===== SCROLLBAR ===== */
::-webkit-scrollbar {
width: 6px;
height: 6px;
}
::-webkit-scrollbar-track {
background: var(--bg-tertiary);
border-radius: 3px;
}
::-webkit-scrollbar-thumb {
background: var(--border-color);
border-radius: 3px;
}
::-webkit-scrollbar-thumb:hover {
background: #52525b;
}
/* ===== LAYOUT FIXES ===== */
.gr-row {
gap: 12px !important;
}
.gr-column {
gap: 12px !important;
}
/* Remove excess padding/margins */
.gr-form {
gap: 12px !important;
}
.gr-block {
padding: 0 !important;
}
/* Label styling */
label, .gr-input-label {
font-size: 0.85rem !important;
font-weight: 500 !important;
color: var(--text-secondary) !important;
margin-bottom: 6px !important;
}
/* Info text */
.gr-info {
font-size: 0.75rem !important;
color: var(--text-muted) !important;
}
/* Fix button alignment in rows */
.gr-button {
height: 40px !important;
}
/* Slider styling */
input[type="range"] {
accent-color: var(--accent-primary) !important;
}
.gr-slider input {
background: var(--bg-tertiary) !important;
}
/* ===== RESPONSIVE ===== */
@media (max-width: 768px) {
.header-title {
font-size: 1.75rem;
}
.gradio-container {
padding: 16px !important;
}
.workflow-container {
grid-template-columns: 1fr;
}
}
"""
# ==================== HELPER FUNCTIONS ====================
def validate_api_key(
provider: str, api_key: str, base_url: str = ""
) -> tuple[bool, str]:
"""Validate that the API key is provided for the selected provider."""
if provider == "custom":
if not api_key or api_key.strip() == "":
return False, "❌ Please provide an API key for custom provider"
if not base_url or base_url.strip() == "":
return False, "❌ Please provide a Base URL for custom provider"
return True, "✅ Custom provider configured"
if not api_key or api_key.strip() == "":
env_var = {
"google": "GOOGLE_API_KEY",
"openai": "OPENAI_API_KEY",
"anthropic": "ANTHROPIC_API_KEY",
}.get(provider, "")
# Check environment variable
if os.getenv(env_var):
return True, f"✅ Using API key from environment variable ({env_var})"
return (
False,
f"❌ Please provide an API key or set the {env_var} environment variable",
)
return True, "✅ API key provided"
def get_model_choices(
provider: str, api_key: str | None = None, base_url: str | None = None
) -> list[str]:
"""Get available models for each provider (fallback list)."""
fallback_models = {
"google": [
"gemini-2.0-flash",
"gemini-1.5-pro",
"gemini-1.5-flash",
],
"openai": [
"gpt-4o",
"gpt-4o-mini",
"gpt-4-turbo",
"gpt-4",
"gpt-3.5-turbo",
],
"anthropic": [
"claude-sonnet-4-20250514",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022",
"claude-3-opus-20240229",
],
"custom": [],
}
return fallback_models.get(provider, [])
def update_model_dropdown(provider: str, api_key: str = "", base_url: str = ""):
"""Update model dropdown when provider changes."""
# Fetch models with status
models, status = fetch_models_for_provider(
provider,
api_key.strip() if api_key and api_key.strip() else None,
base_url.strip() if base_url and base_url.strip() else None,
)
if not models:
models = get_model_choices(provider)
return gr.update(choices=models, value=models[0] if models else None)
def update_base_url_visibility(provider: str):
"""Show/hide base URL field based on provider."""
return gr.update(visible=(provider == "custom"))
def refresh_models(provider: str, api_key: str, base_url: str):
"""Refresh the model list from the API."""
models, status = fetch_models_for_provider(
provider,
api_key.strip() if api_key and api_key.strip() else None,
base_url.strip() if base_url and base_url.strip() else None,
)
if models:
return gr.update(choices=models, value=models[0]), status
# Use fallback if no models returned
fallback = get_model_choices(provider)
if fallback:
return gr.update(
choices=fallback, value=fallback[0]
), status or "ℹ️ Using default models"
return gr.update(), status or "❌ No models available"
def copy_uploaded_files(files: list) -> str:
"""Copy uploaded files to the data directory."""
data_dir = os.path.join(os.path.dirname(__file__), "data")
# Clear existing files in data directory (except .gitkeep)
if os.path.exists(data_dir):
for f in os.listdir(data_dir):
if f != ".gitkeep":
file_path = os.path.join(data_dir, f)
if os.path.isfile(file_path):
os.remove(file_path)
else:
os.makedirs(data_dir)
# Copy new files
copied_files = []
if files:
for file_path in files:
if file_path:
filename = os.path.basename(file_path)
dest_path = os.path.join(data_dir, filename)
shutil.copy2(file_path, dest_path)
copied_files.append(filename)
if copied_files:
return f"✅ Uploaded {len(copied_files)} file(s): {', '.join(copied_files)}"
return "ℹ️ No files uploaded. Using existing files in data/ directory."
def list_data_files() -> str:
"""List files currently in the data directory."""
data_dir = os.path.join(os.path.dirname(__file__), "data")
if not os.path.exists(data_dir):
return "No data directory found."
files = [
f
for f in os.listdir(data_dir)
if f != ".gitkeep" and os.path.isfile(os.path.join(data_dir, f))
]
if files:
file_list = "\n".join([f" 📄 {f}" for f in files])
return f"**Files in data/ directory:**\n{file_list}"
return "No data files found. Please upload some files."
# ==================== MAIN WORKFLOW ====================
def run_ds_star_workflow(
query: str,
provider: str,
model: str,
api_key: str,
base_url: str,
max_iterations: int,
temperature: float,
progress=gr.Progress(),
) -> Generator[tuple[str, str, str, str], None, None]:
"""
Run the DS-STAR workflow with streaming updates.
Yields: (status, current_step, code_output, execution_result)
"""
# Validate inputs
if not query or query.strip() == "":
yield "❌ Error", "Please enter a query", "", ""
return
is_valid, message = validate_api_key(provider, api_key, base_url)
if not is_valid:
yield "❌ Configuration Error", message, "", ""
return
# Initialize LLM
yield "🔄 Initializing...", "Setting up LLM connection", "", ""
# For custom provider, use openai with custom base_url
actual_provider = "openai" if provider == "custom" else provider
try:
llm = get_llm(
provider=actual_provider,
model=model,
api_key=api_key if api_key.strip() else None,
temperature=temperature,
base_url=base_url if provider == "custom" and base_url.strip() else None,
)
except Exception as e:
yield "❌ LLM Error", f"Failed to initialize LLM: {str(e)}", "", ""
return
# Build graph
yield "🔄 Building Graph...", "Constructing multi-agent workflow", "", ""
try:
app = build_ds_star_graph(llm, max_iterations)
except Exception as e:
yield "❌ Graph Error", f"Failed to build graph: {str(e)}", "", ""
return
# Create initial state
initial_state = create_initial_state(query, llm, max_iterations)
config = {"configurable": {"thread_id": f"gradio-session-{os.urandom(4).hex()}"}}
# Run workflow with progress updates
step_descriptions = {
"analyzer": "📊 Analyzing data files...",
"planner": "📝 Creating execution plan...",
"coder": "💻 Generating code...",
"verifier": "✅ Verifying solution...",
"router": "🔀 Routing to next step...",
"backtrack": "↩️ Backtracking...",
"finalyzer": "🎯 Finalizing solution...",
}
yield "🚀 Running DS-STAR...", "Starting multi-agent workflow", "", ""
try:
# Stream through the workflow
current_code = ""
current_result = ""
iteration = 0
for event in app.stream(initial_state, config, stream_mode="values"):
# Update progress based on current state
next_node = event.get("next", "")
iteration = event.get("iteration", 0)
step_desc = step_descriptions.get(next_node, f"Processing: {next_node}")
progress_msg = (
f"Iteration {iteration}/{max_iterations}"
if iteration > 0
else "Starting..."
)
current_code = event.get("current_code", current_code) or ""
current_result = event.get("execution_result", current_result) or ""
yield f"🔄 {progress_msg}", step_desc, current_code, current_result
# Final state
final_code = event.get("current_code", "") or ""
final_result = event.get("execution_result", "") or ""
yield "✅ Complete!", "Workflow finished successfully", final_code, final_result
except Exception as e:
import traceback
error_trace = traceback.format_exc()
yield (
"❌ Execution Error",
f"Error: {str(e)}\n\n{error_trace}",
current_code,
current_result,
)
# ==================== GRADIO INTERFACE ====================
def create_gradio_app():
"""Create and configure the Gradio application."""
with gr.Blocks(title="DS-STAR | Multi-Agent Data Science") as demo:
# Header Section
gr.HTML("""
""")
with gr.Row():
# Left Column - Configuration (narrower)
with gr.Column(scale=1, min_width=300):
# LLM Configuration - Accordion
with gr.Accordion("🔑 LLM Configuration", open=True):
provider = gr.Dropdown(
choices=["google", "openai", "anthropic", "custom"],
value="google",
label="Provider",
info="Select your LLM provider",
)
base_url = gr.Textbox(
label="Base URL",
placeholder="https://api.together.xyz/v1",
info="OpenAI-compatible API endpoint",
visible=False,
)
with gr.Row():
model = gr.Dropdown(
choices=get_model_choices("google"),
value="gemini-2.0-flash",
label="Model",
scale=5,
)
refresh_models_btn = gr.Button(
"🔄",
variant="secondary",
size="sm",
scale=1,
min_width=40,
elem_classes="refresh-btn",
)
api_key = gr.Textbox(
label="API Key",
type="password",
placeholder="Enter API key or use env variable",
)
api_status = gr.Markdown(
"💡 *Enter API key or set environment variable*"
)
# Advanced Settings - Accordion (closed by default)
with gr.Accordion("⚙️ Advanced Settings", open=False):
max_iterations = gr.Slider(
minimum=1,
maximum=50,
value=20,
step=1,
label="Max Iterations",
info="Maximum refinement cycles",
)
temperature = gr.Slider(
minimum=0.0,
maximum=1.0,
value=0.0,
step=0.1,
label="Temperature",
info="Controls response creativity (0 = deterministic)",
)
# Data Files - Accordion
with gr.Accordion("📁 Data Files", open=True):
file_upload = gr.File(
label="Upload Files",
file_count="multiple",
file_types=[".csv", ".json", ".xlsx", ".parquet", ".txt"],
type="filepath",
)
upload_status = gr.Markdown(
list_data_files(), elem_classes="file-list"
)
refresh_btn = gr.Button(
"🔄 Refresh Files", variant="secondary", size="sm"
)
# Right Column - Main Interface (wider)
with gr.Column(scale=2, min_width=500):
# Query Input Section
gr.HTML("""
""")
query_input = gr.Textbox(
label="",
placeholder="e.g., What percentage of transactions use credit cards? Show the distribution by category.",
lines=2,
max_lines=4,
show_label=False,
)
# Buttons Row - properly aligned
with gr.Row():
run_btn = gr.Button(
"🚀 Run Analysis",
variant="primary",
size="lg",
scale=3,
)
clear_btn = gr.Button(
"🗑️ Clear",
variant="secondary",
size="lg",
scale=1,
)
# Example Queries - as clickable chips
gr.HTML("""
💡 Quick Examples
""")
gr.Examples(
examples=[
["Show distribution of transaction amounts"],
["Which category has highest sales?"],
["Find correlations between columns"],
["Create summary statistics report"],
],
inputs=query_input,
label="",
examples_per_page=4,
)
# Status Section
gr.HTML("""
""")
with gr.Row():
status_display = gr.Textbox(
label="Status",
value="⏳ Ready",
interactive=False,
scale=1,
elem_classes="status-box",
lines=1,
max_lines=1,
)
current_step = gr.Textbox(
label="Current Step",
value="Waiting for query...",
interactive=False,
scale=2,
elem_classes="step-box",
lines=1,
max_lines=1,
)
# Results Tabs
with gr.Tabs():
with gr.TabItem("💻 Code", id=0):
code_output = gr.Code(
label="",
language="python",
lines=16,
interactive=False,
show_label=False,
)
with gr.TabItem("📊 Output", id=1):
result_output = gr.Textbox(
label="",
lines=14,
interactive=False,
show_label=False,
)
with gr.TabItem("🔄 Workflow", id=2):
gr.HTML("""
1
Analyzer
Examines data structure
2
Planner
Creates execution plan
3
Coder
Generates Python code
4
Verifier
Validates solution
6
Finalyzer
Delivers final result
""")
with gr.TabItem("ℹ️ About", id=3):
gr.Markdown("""
## About DS-STAR
**DS-STAR** (Data Science - Structured Task Analysis and Resolution) is a multi-agent system for automating data science tasks.
### ✨ Features
- 🤖 **Multi-Agent** — Specialized agents for analysis, planning, coding & verification
- 🔄 **Iterative** — Automatically refines solutions
- 🔙 **Backtracking** — Smart rollback when needed
- 💻 **Code Gen** — Produces clean Python code
### 🔌 Providers
**Google Gemini** • **OpenAI GPT** • **Anthropic Claude** • **Custom API**
---
Built for the **HuggingFace MCP Hackathon** • [GitHub](https://github.com/Anurag-Deo/DS-STAR)
""")
# Event Handlers
provider.change(
fn=update_model_dropdown,
inputs=[provider, api_key, base_url],
outputs=[model],
)
provider.change(
fn=update_base_url_visibility,
inputs=[provider],
outputs=[base_url],
)
provider.change(
fn=lambda p, k, b: validate_api_key(p, k, b)[1],
inputs=[provider, api_key, base_url],
outputs=[api_status],
)
api_key.change(
fn=lambda p, k, b: validate_api_key(p, k, b)[1],
inputs=[provider, api_key, base_url],
outputs=[api_status],
)
base_url.change(
fn=lambda p, k, b: validate_api_key(p, k, b)[1],
inputs=[provider, api_key, base_url],
outputs=[api_status],
)
refresh_models_btn.click(
fn=refresh_models,
inputs=[provider, api_key, base_url],
outputs=[model, api_status],
)
file_upload.change(
fn=copy_uploaded_files, inputs=[file_upload], outputs=[upload_status]
)
refresh_btn.click(fn=list_data_files, outputs=[upload_status])
run_btn.click(
fn=run_ds_star_workflow,
inputs=[
query_input,
provider,
model,
api_key,
base_url,
max_iterations,
temperature,
],
outputs=[status_display, current_step, code_output, result_output],
)
clear_btn.click(
fn=lambda: ("⏳ Ready", "Waiting for query...", "", ""),
outputs=[status_display, current_step, code_output, result_output],
)
return demo
# ==================== MAIN ====================
if __name__ == "__main__":
demo = create_gradio_app()
theme = create_ds_star_theme()
demo.launch(share=False, show_error=True, theme=theme, css=CUSTOM_CSS)