Upload folder using huggingface_hub

.env.example

@@ -0,0 +1,6 @@
+LLM_PROVIDER=
+LLM_BASE_URL=
+LLM_MODEL=
+LLM_API_KEY=
+GOOGLE_API_KEY=
+ANTHROPIC_API_KEY=""

.github/workflows/update_space.yml

@@ -0,0 +1,28 @@
+name: Run Python script
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Install Gradio
+      run: python -m pip install gradio
+
+    - name: Log in to Hugging Face
+      run: python -c 'import huggingface_hub; huggingface_hub.login(token="${{ secrets.hf_token }}")'
+
+    - name: Deploy to Spaces
+      run: gradio deploy

.gitignore

@@ -0,0 +1,13 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# .env files
+.env

.python-version

@@ -0,0 +1 @@
+3.13

README.md

@@ -1,12 +1,277 @@
----
-title: DS STAR
-emoji: 🌍
-colorFrom: indigo
-colorTo: purple
-sdk: gradio
-sdk_version: 6.0.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: DS-STAR
+emoji: ✨
+colorFrom: purple
+colorTo: indigo
+sdk: gradio
+sdk_version: 6.0.1
+app_file: app.py
+pinned: false
+license: mit
+short_description: Multi-Agent AI System for Automated Data Science Tasks
+tags:
+  - mcp-in-action-track-consumer
+  - langgraph
+  - multi-agent
+  - data-science
+  - automation
+---
+
+<div align="center">
+
+# ✨ DS-STAR
+
+### **D**ata **S**cience - **S**tructured **T**ask **A**nalysis and **R**esolution
+
+[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Anurag-Deo/DS-STAR)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+[![LangGraph](https://img.shields.io/badge/Built%20with-LangGraph-orange)](https://langchain-ai.github.io/langgraph/)
+
+**A powerful multi-agent AI system that automates data science tasks through intelligent collaboration.**
+
+[🚀 Try Demo](https://huggingface.co/spaces/Anurag-Deo/DS-STAR) • [📖 Documentation](#-usage) • [🐛 Report Bug](https://github.com/Anurag-Deo/DS-STAR/issues)
+
+</div>
+
+---
+
+## 🎯 What is DS-STAR?
+
+DS-STAR is a **multi-agent AI system** built with LangGraph that takes your natural language questions about data and automatically:
+
+1. 📊 **Analyzes** your data files to understand their structure
+2. 📝 **Plans** a step-by-step approach to answer your question  
+3. 💻 **Generates** Python code to perform the analysis
+4. ✅ **Verifies** the solution meets your requirements
+5. 🔄 **Iterates** with smart backtracking if needed
+6. 🎯 **Delivers** polished, accurate results
+
+> **Built for the 🤗 Hugging Face MCP 1st Birthday Hackathon**
+
+---
+
+## ✨ Key Features
+
+| Feature | Description |
+|---------|-------------|
+| 🤖 **Multi-Agent Architecture** | Six specialized agents working in harmony |
+| 🔄 **Iterative Refinement** | Automatically improves solutions through multiple cycles |
+| 🔙 **Smart Backtracking** | Intelligently reverts failed approaches |
+| 📊 **Auto Data Analysis** | Understands your data structure automatically |
+| 💻 **Code Generation** | Produces clean, executable Python code |
+| 🌐 **Multi-Provider Support** | Works with Google, OpenAI, Anthropic, or custom APIs |
+| 🎨 **Modern UI** | Beautiful dark-themed Gradio interface |
+
+---
+
+## 🏗️ Architecture
+
+DS-STAR uses a sophisticated multi-agent workflow powered by LangGraph:
+
+```
+┌─────────────┐     ┌─────────────┐     ┌─────────────┐
+│   Analyzer  │────▶│   Planner   │────▶│    Coder    │
+│  📊 Analyze │     │  📝 Plan    │     │  💻 Code    │
+└─────────────┘     └─────────────┘     └─────────────┘
+                                               │
+                                               ▼
+┌─────────────┐     ┌─────────────┐     ┌─────────────┐
+│  Finalyzer  │◀────│   Router    │◀────│  Verifier   │
+│  🎯 Polish  │     │  🔀 Route   │     │  ✅ Verify  │
+└─────────────┘     └─────────────┘     └─────────────┘
+                           │
+                           ▼
+                    ┌─────────────┐
+                    │  Backtrack  │
+                    │  ↩️ Retry   │
+                    └─────────────┘
+```
+
+### Agent Roles
+
+| Agent | Role | Description |
+|-------|------|-------------|
+| **Analyzer** | 📊 | Examines all data files and creates detailed descriptions |
+| **Planner** | 📝 | Generates the next logical step in the solution |
+| **Coder** | 💻 | Implements the plan as executable Python code |
+| **Verifier** | ✅ | Validates if the solution answers the query |
+| **Router** | 🔀 | Decides to continue, add steps, or backtrack |
+| **Finalyzer** | 🎯 | Polishes and formats the final output |
+
+---
+
+## 🚀 Quick Start
+
+### Online Demo
+
+Try DS-STAR instantly on Hugging Face Spaces:
+
+👉 **[Launch DS-STAR Demo](https://huggingface.co/spaces/Anurag-Deo/DS-STAR)**
+
+### Local Installation
+
+```bash
+# Clone the repository
+git clone https://github.com/Anurag-Deo/DS-STAR.git
+cd DS-STAR
+
+# Create virtual environment
+python -m venv .venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Run the application
+python app.py
+```
+
+Then open http://localhost:7860 in your browser.
+
+---
+
+## 💡 Usage
+
+### Web Interface
+
+1. **Select Provider** — Choose Google, OpenAI, Anthropic, or Custom
+2. **Enter API Key** — Or set via environment variable
+3. **Upload Data** — Drop your CSV, JSON, Excel, or Parquet files
+4. **Ask Questions** — Type your data science question
+5. **Run Analysis** — Click "Run Analysis" and watch the magic!
+
+### Example Queries
+
+```
+📊 "What percentage of transactions use credit cards?"
+📈 "Show me the distribution of transaction amounts"
+🏆 "Which category has the highest total sales?"
+🔗 "Find correlations between numeric columns"
+📋 "Create a summary statistics report"
+```
+
+### Python API
+
+```python
+from src.graph import run_ds_star
+from src.config import get_llm
+
+# Initialize LLM
+llm = get_llm(provider="google", model="gemini-2.0-flash")
+
+# Run DS-STAR
+result = run_ds_star(
+    query="What is the average transaction amount?",
+    llm=llm,
+    max_iterations=20
+)
+```
+
+---
+
+## 🔌 Supported Providers
+
+| Provider | Models | Environment Variable |
+|----------|--------|---------------------|
+| **Google** | Gemini 2.0, 1.5 Pro, 1.5 Flash | `GOOGLE_API_KEY` |
+| **OpenAI** | GPT-4o, GPT-4, GPT-3.5 | `OPENAI_API_KEY` |
+| **Anthropic** | Claude 3.5, Claude 3 | `ANTHROPIC_API_KEY` |
+| **Custom** | Any OpenAI-compatible API | Custom Base URL |
+
+---
+
+## 📁 Project Structure
+
+```
+DS-STAR/
+├── 📱 app.py                 # Gradio web application
+├── 📜 main.py                # CLI entry point
+├── 📋 requirements.txt       # Dependencies
+├── 📂 src/
+│   ├── 🤖 agents/            # Agent implementations
+│   │   ├── analyzer_agent.py
+│   │   ├── planner_agent.py
+│   │   ├── coder_agent.py
+│   │   ├── verifier_agent.py
+│   │   ├── router_agent.py
+│   │   └── finalyzer_agent.py
+│   ├── 🔧 utils/             # Shared utilities
+│   │   ├── state.py          # State schema
+│   │   ├── formatters.py     # Text formatting
+│   │   └── code_execution.py # Safe code execution
+│   ├── ⚙️ config/            # Configuration
+│   │   └── llm_config.py     # LLM setup
+│   └── 🔄 graph.py           # LangGraph workflow
+├── 🧪 tests/                 # Test suite
+└── 📊 data/                  # Sample data files
+```
+
+---
+
+## 🧪 Testing
+
+```bash
+# Run complete workflow test
+python tests/test_complete_workflow.py
+
+# Test individual agents
+python -c "from src.agents import test_analyzer; test_analyzer(llm)"
+```
+
+---
+
+## 🛠️ Configuration
+
+### Environment Variables
+
+```bash
+# Set your API keys
+export GOOGLE_API_KEY="your-google-api-key"
+export OPENAI_API_KEY="your-openai-api-key"
+export ANTHROPIC_API_KEY="your-anthropic-api-key"
+```
+
+### Advanced Settings
+
+| Setting | Default | Description |
+|---------|---------|-------------|
+| Max Iterations | 20 | Maximum refinement cycles |
+| Temperature | 0.0 | LLM temperature (0 = deterministic) |
+
+---
+
+## 🤝 Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+
+1. Fork the repository
+2. Create your feature branch (`git checkout -b feature/AmazingFeature`)
+3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)
+4. Push to the branch (`git push origin feature/AmazingFeature`)
+5. Open a Pull Request
+
+---
+
+## 📄 License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+---
+
+## 🙏 Acknowledgments
+
+- Built with [LangGraph](https://langchain-ai.github.io/langgraph/) by LangChain
+- UI powered by [Gradio](https://gradio.app/)
+- Created for the [🤗 Hugging Face MCP 1st Birthday Hackathon](https://huggingface.co/)
+
+---
+
+<div align="center">
+
+**Made with ❤️ by [Anurag Deo](https://github.com/Anurag-Deo)**
+
+⭐ Star this repo if you find it helpful!
+
+</div>
+

app.py

@@ -0,0 +1,1406 @@
+"""
+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("""
+            <div class="header-section">
+                <div class="header-content">
+                    <h1 class="header-title">
+                        <span class="star-icon">✨</span>
+                        DS-STAR
+                    </h1>
+                    <p class="header-subtitle">Multi-Agent System for Automated Data Science Tasks</p>
+                    <div class="header-badges">
+                        <span class="header-badge">🔗 LangGraph</span>
+                        <span class="header-badge">🤗 HuggingFace MCP</span>
+                        <span class="header-badge">🤖 Multi-Agent</span>
+                    </div>
+                </div>
+            </div>
+        """)
+
+        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("""
+                    <div class="section-header">
+                        <div class="section-icon">💬</div>
+                        <h3 class="section-title">Ask Your Question</h3>
+                    </div>
+                """)
+
+                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("""
+                    <div style="margin-top: 16px; padding-top: 12px; border-top: 1px solid var(--border-color);">
+                        <span style="font-size: 0.8rem; font-weight: 600; color: var(--text-muted); margin-bottom: 8px; display: block;">💡 Quick Examples</span>
+                    </div>
+                """)
+                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("""
+                    <div class="section-header" style="margin-top: 20px;">
+                        <div class="section-icon">📊</div>
+                        <h3 class="section-title">Status</h3>
+                    </div>
+                """)
+
+                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("""
+                            <div class="workflow-container">
+                                <div class="workflow-step">
+                                    <div class="step-number">1</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Analyzer</div>
+                                        <div class="step-desc">Examines data structure</div>
+                                    </div>
+                                </div>
+                                <div class="workflow-step">
+                                    <div class="step-number">2</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Planner</div>
+                                        <div class="step-desc">Creates execution plan</div>
+                                    </div>
+                                </div>
+                                <div class="workflow-step">
+                                    <div class="step-number">3</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Coder</div>
+                                        <div class="step-desc">Generates Python code</div>
+                                    </div>
+                                </div>
+                                <div class="workflow-step">
+                                    <div class="step-number">4</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Verifier</div>
+                                        <div class="step-desc">Validates solution</div>
+                                    </div>
+                                </div>
+                                <div class="workflow-step">
+                                    <div class="step-number">5</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Router</div>
+                                        <div class="step-desc">Decides next step</div>
+                                    </div>
+                                </div>
+                                <div class="workflow-step">
+                                    <div class="step-number">6</div>
+                                    <div class="step-content">
+                                        <div class="step-title">Finalyzer</div>
+                                        <div class="step-desc">Delivers final result</div>
+                                    </div>
+                                </div>
+                            </div>
+                        """)
+
+                    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)

main.py

@@ -0,0 +1,69 @@
+"""
+DS-STAR: Multi-Agent System for Data Science Tasks
+
+This is the main entry point for the refactored DS-STAR system.
+All agents are modularized and can be tested independently.
+
+Usage:
+    python main_new.py
+
+Or customize:
+    from src.graph import run_ds_star
+    from src.config import get_llm
+
+    llm = get_llm(provider="google", model="gemini-1.5-flash")
+    result = run_ds_star("Your question here", llm, max_iterations=20)
+"""
+
+import sys
+
+from src.config import DEFAULT_CONFIG, get_llm
+from src.graph import run_ds_star
+
+
+def main():
+    """
+    Main execution function for DS-STAR.
+    """
+    # Configuration
+    query = "What percentage of transactions use credit cards?"
+    max_iterations = DEFAULT_CONFIG["max_iterations"]
+    provider = DEFAULT_CONFIG["provider"]
+    model = DEFAULT_CONFIG["model"]
+
+    print("Initializing DS-STAR Multi-Agent System...")
+    print(f"Provider: {provider}")
+    print(f"Model: {model}")
+    print()
+
+    try:
+        # Initialize LLM
+        llm = get_llm(
+            provider=provider, model=model, temperature=DEFAULT_CONFIG["temperature"]
+        )
+
+        # Run DS-STAR workflow
+        final_state = run_ds_star(
+            query=query,
+            llm=llm,
+            max_iterations=max_iterations,
+            thread_id="ds-star-main-session",
+        )
+
+        if final_state:
+            print("\n✅ Workflow completed successfully!")
+            return 0
+        else:
+            print("\n❌ Workflow failed!")
+            return 1
+
+    except Exception as e:
+        print(f"\n❌ Fatal error: {str(e)}")
+        import traceback
+
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "ds-star"
+version = "0.1.0"
+description = "DS-STAR: Multi-Agent System for Automated Data Science Tasks"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "python-dotenv>=1.0.0",
+    "langchain>=0.3.0",
+    "langchain-anthropic>=0.3.0",
+    "langchain-core>=0.3.0",
+    "langchain-google-genai>=2.0.0",
+    "langchain-openai>=0.3.0",
+    "langgraph>=0.2.0",
+    "pandas>=2.0.0",
+    "gradio>=5.0.0",
+]
+
+[project.scripts]
+ds-star = "app:create_gradio_app"
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "black>=24.0.0",
+    "ruff>=0.5.0",
+]

requirements.txt

@@ -0,0 +1,18 @@
+# DS-STAR Dependencies
+# For Hugging Face Spaces deployment
+
+# Core dependencies
+python-dotenv>=1.0.0
+langchain>=0.3.0
+langchain-anthropic>=0.3.0
+langchain-core>=0.3.0
+langchain-google-genai>=2.0.0
+langchain-openai>=0.3.0
+langgraph>=0.2.0
+pandas>=2.0.0
+
+# Web interface
+gradio>=5.0.0
+
+# Additional utilities
+numpy>=1.20.0

src/__init__.py

@@ -0,0 +1 @@
+"""DS-STAR multi-agent system for data science tasks."""

src/agents/__init__.py

@@ -0,0 +1,24 @@
+"""Agent modules for DS-STAR system."""
+
+from .analyzer_agent import analyzer_node, test_analyzer
+from .coder_agent import coder_node, test_coder
+from .finalyzer_agent import finalyzer_node, test_finalyzer
+from .planner_agent import planner_node, test_planner
+from .router_agent import backtrack_node, router_node, test_router
+from .verifier_agent import test_verifier, verifier_node
+
+__all__ = [
+    "analyzer_node",
+    "planner_node",
+    "coder_node",
+    "verifier_node",
+    "router_node",
+    "backtrack_node",
+    "finalyzer_node",
+    "test_analyzer",
+    "test_planner",
+    "test_coder",
+    "test_verifier",
+    "test_router",
+    "test_finalyzer",
+]

src/agents/analyzer_agent.py

@@ -0,0 +1,167 @@
+"""
+Analyzer Agent: Analyzes data files and generates descriptions.
+
+This agent runs once at the beginning to understand available data.
+"""
+
+import os
+from pathlib import Path
+
+from langchain_core.messages import AIMessage
+
+from ..utils.code_execution import execute_with_debug
+from ..utils.formatters import extract_code, gemini_text
+from ..utils.state import DSStarState
+
+
+def analyzer_node(state: DSStarState) -> dict:
+    """
+    Analyzer Agent Node: Analyzes all data files in the data/ directory.
+
+    For each file, generates and executes Python code to:
+    - Load the file
+    - Print structure, types, and sample data
+    - Capture essential information
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - data_descriptions: Dict mapping filename to analysis result
+        - messages: Agent communication messages
+        - next: Next node to visit ("planner" or "__end__")
+    """
+    print("=" * 60)
+    print("DATA ANALYZER AGENT STARTING...")
+    print("=" * 60)
+
+    data_dir = "data/"
+    descriptions = {}
+
+    # Check if data directory exists
+    if not os.path.exists(data_dir):
+        print(f"Error: {data_dir} directory not found")
+        return {
+            "data_descriptions": {"error": "Data directory not found"},
+            "messages": [AIMessage(content="Error: data/ directory not found")],
+            "next": "__end__",
+        }
+
+    # Get list of files
+    files = [
+        f for f in os.listdir(data_dir) if os.path.isfile(os.path.join(data_dir, f))
+    ]
+
+    if not files:
+        print(f"Error: No files found in {data_dir}")
+        return {
+            "data_descriptions": {"error": "No data files found"},
+            "messages": [AIMessage(content="Error: No files in data/ directory")],
+            "next": "__end__",
+        }
+
+    print(f"Found {len(files)} files to analyze")
+
+    # Analyze each file
+    for filename in files:
+        filepath = os.path.join(data_dir, filename)
+        file_ext = Path(filepath).suffix.lower()
+
+        print(f"\nAnalyzing: {filename}")
+
+        # Generate analysis script
+        analysis_prompt = f"""Generate a Python script to analyze the file: {filepath}
+
+File type: {file_ext}
+
+Requirements:
+- Load the file using appropriate method for {file_ext} format
+- Print essential information:
+  * Data structure and types
+  * Column names (for structured data like CSV, Excel)
+  * First 3-5 rows/examples
+  * Shape/size information
+- Handle common formats: CSV, JSON, Excel, TXT, MD
+- Use pandas for structured data
+- No try-except blocks
+- All files are in 'data/' directory
+- Print output clearly
+
+Provide ONLY the Python code in a markdown code block."""
+
+        try:
+            # Get LLM response
+            response = state["llm"].invoke(analysis_prompt)
+
+            # Handle different response formats (Gemini vs OpenAI)
+            if hasattr(response, "content") and isinstance(response.content, list):
+                # Gemini format
+                response_text = gemini_text(response)
+            elif hasattr(response, "content"):
+                response_text = response.content
+            else:
+                response_text = str(response)
+
+            code = extract_code(response_text)
+
+            # Execute with debugging
+            result = execute_with_debug(code, state["llm"], is_analysis=True)
+
+            descriptions[filename] = result
+            print(f"✓ Successfully analyzed {filename}")
+
+        except Exception as e:
+            descriptions[filename] = f"Error analyzing file: {str(e)}"
+            print(f"✗ Failed to analyze {filename}: {str(e)}")
+
+    print("\n" + "=" * 60)
+    print(f"ANALYSIS COMPLETE: {len(files)} files processed")
+    print("=" * 60)
+
+    return {
+        "data_descriptions": descriptions,
+        "messages": [AIMessage(content=f"Analyzed {len(files)} data files")],
+        "next": "planner",
+    }
+
+
+# Standalone test function
+def test_analyzer(llm, data_dir: str = "data/"):
+    """
+    Test the analyzer agent independently.
+
+    Args:
+        llm: LLM instance
+        data_dir: Directory containing data files
+
+    Returns:
+        Dictionary with analysis results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": "Test query",
+        "data_descriptions": {},
+        "plan": [],
+        "current_code": "",
+        "execution_result": "",
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "analyzer",
+    }
+
+    result = analyzer_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("ANALYZER TEST RESULTS")
+    print("=" * 60)
+    for filename, description in result["data_descriptions"].items():
+        print(f"\n{filename}:")
+        print("-" * 60)
+        print(description)
+
+    return result

src/agents/coder_agent.py

@@ -0,0 +1,174 @@
+"""
+Coder Agent: Implements the plan as executable Python code.
+
+This agent generates Python code that implements all steps in the current plan.
+"""
+
+from langchain_core.messages import AIMessage
+
+from ..utils.code_execution import execute_with_debug
+from ..utils.formatters import extract_code, format_data_descriptions, format_plan
+from ..utils.state import DSStarState
+
+
+def coder_node(state: DSStarState) -> dict:
+    """
+    Coder Agent Node: Generates and executes Python code for the plan.
+
+    On first call: Generates code implementing all plan steps
+    On subsequent calls: Updates code to include new plan steps
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - current_code: Generated Python code
+        - execution_result: Output from code execution
+        - messages: Agent communication messages
+        - next: Next node to visit ("verifier")
+    """
+    print("=" * 60)
+    print("CODER AGENT STARTING...")
+    print("=" * 60)
+
+    data_context = format_data_descriptions(state["data_descriptions"])
+    plan_text = format_plan(state["plan"])
+
+    is_initial = state["current_code"] == ""
+
+    if is_initial:
+        print("Generating INITIAL code implementation...")
+        prompt = f"""You are an expert Python developer for data science.
+
+Available Data Files:
+{data_context}
+
+Plan to Implement:
+{plan_text}
+
+Task: Write a Python script that implements ALL steps in the plan.
+
+Requirements:
+- Use pandas for data manipulation
+- All files are in 'data/' directory
+- Print intermediate results for each step
+- No try-except blocks
+- Clean, readable code
+
+Provide ONLY the Python code in a markdown code block."""
+    else:
+        print(f"Updating code to implement {len(state['plan'])} steps...")
+        prompt = f"""You are an expert Python developer for data science.
+
+Available Data Files:
+{data_context}
+
+Complete Plan:
+{plan_text}
+
+Previous Code:
+{state["current_code"]}
+
+Task: Update the code to implement the COMPLETE current plan.
+Build upon the previous code, extending it to include all plan steps.
+
+Requirements:
+- Use pandas for data manipulation
+- All files are in 'data/' directory
+- Print intermediate and final results
+- No try-except blocks
+- Clean, readable code
+
+Provide ONLY the updated Python code in a markdown code block."""
+
+    try:
+        # Get LLM response
+        response = state["llm"].invoke(prompt)
+
+        # Handle different response formats
+        if hasattr(response, "content") and isinstance(response.content, list):
+            from ..utils.formatters import gemini_text
+
+            response_text = gemini_text(response)
+        elif hasattr(response, "content"):
+            response_text = response.content
+        else:
+            response_text = str(response)
+
+        code = extract_code(response_text)
+
+        print("\nGenerated Code:")
+        print("-" * 60)
+        print(code[:200] + "..." if len(code) > 200 else code)
+        print("-" * 60)
+
+        print("\nExecuting code...")
+
+        # Execute with debugging
+        result = execute_with_debug(
+            code, state["llm"], is_analysis=False, data_context=data_context
+        )
+
+        print("\nExecution Result:")
+        print("-" * 60)
+        print(result[:200] + "..." if len(result) > 200 else result)
+        print("-" * 60)
+        print("=" * 60)
+
+        return {
+            "current_code": code,
+            "execution_result": result,
+            "messages": [AIMessage(content="Code executed")],
+            "next": "verifier",
+        }
+
+    except Exception as e:
+        print(f"\n✗ Coder error: {str(e)}")
+        return {
+            "messages": [AIMessage(content=f"Coder error: {str(e)}")],
+            "next": "__end__",
+        }
+
+
+# Standalone test function
+def test_coder(llm, query: str, data_descriptions: dict, plan: list):
+    """
+    Test the coder agent independently.
+
+    Args:
+        llm: LLM instance
+        query: User query
+        data_descriptions: Dict of filename -> description
+        plan: List of plan steps to implement
+
+    Returns:
+        Dictionary with coder results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": query,
+        "data_descriptions": data_descriptions,
+        "plan": plan,
+        "current_code": "",
+        "execution_result": "",
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "coder",
+    }
+
+    result = coder_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("CODER TEST RESULTS")
+    print("=" * 60)
+    print("Code:")
+    print(result.get("current_code", "No code generated"))
+    print("\nExecution Result:")
+    print(result.get("execution_result", "No result"))
+
+    return result

src/agents/finalyzer_agent.py

@@ -0,0 +1,162 @@
+"""
+Finalyzer Agent: Creates final polished solution with clear output.
+
+This agent runs when the verifier confirms the plan is sufficient.
+It generates a final version of the code with improved formatting and output.
+"""
+
+from langchain_core.messages import AIMessage
+
+from ..utils.code_execution import execute_code_safely
+from ..utils.formatters import extract_code, format_data_descriptions, gemini_text
+from ..utils.state import DSStarState
+
+
+def finalyzer_node(state: DSStarState) -> dict:
+    """
+    Finalyzer Agent Node: Creates final polished solution.
+
+    Takes the working code and creates a final version with:
+    - Clear answer to the original question
+    - Proper output formatting
+    - Self-contained executable code
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - current_code: Final polished code
+        - execution_result: Final execution output
+        - messages: Agent communication messages
+        - next: "__end__" (workflow complete)
+    """
+    print("=" * 60)
+    print("FINALYZER AGENT STARTING...")
+    print("=" * 60)
+
+    data_context = format_data_descriptions(state["data_descriptions"])
+
+    prompt = f"""You are an expert data analyst creating final solutions.
+
+Original Question: {state["query"]}
+
+Available Data:
+{data_context}
+
+Working Code:
+{state["current_code"]}
+
+Execution Result:
+{state["execution_result"]}
+
+Task: Create a final version of the code that:
+1. Clearly prints the answer to the question
+2. Includes proper formatting of the output
+3. Is self-contained and executable
+4. Has clear print statements
+
+Provide ONLY the final Python code in a markdown code block."""
+
+    try:
+        # Get LLM response
+        response = state["llm"].invoke(prompt)
+
+        # Handle different response formats
+        if hasattr(response, "content") and isinstance(response.content, list):
+            response_text = gemini_text(response)
+        elif hasattr(response, "content"):
+            response_text = response.content
+        else:
+            response_text = str(response)
+
+        final_code = extract_code(response_text)
+
+        print("\nFinal Code Generated:")
+        print("-" * 60)
+        print(final_code[:300] + "..." if len(final_code) > 300 else final_code)
+        print("-" * 60)
+
+        # Execute final code
+        print("\nExecuting final code...")
+        success, stdout, stderr = execute_code_safely(final_code)
+
+        if success:
+            final_result = stdout
+            print("\n✓ Final execution successful")
+        else:
+            # If final execution fails, use previous result
+            print("\n⚠ Final execution failed, using previous result")
+            final_result = state["execution_result"]
+
+        print("\nFinal Result:")
+        print("-" * 60)
+        print(final_result[:300] + "..." if len(final_result) > 300 else final_result)
+        print("-" * 60)
+        print("=" * 60)
+        print("SOLUTION COMPLETE ✓")
+        print("=" * 60)
+
+        return {
+            "current_code": final_code,
+            "execution_result": final_result,
+            "messages": [AIMessage(content="Solution finalized")],
+            "next": "__end__",
+        }
+
+    except Exception as e:
+        # On error, return current state
+        print(f"\n✗ Finalyzer error: {str(e)}")
+        print("Using current solution as final")
+        return {
+            "messages": [
+                AIMessage(content=f"Finalyzer error: {str(e)}, using current solution")
+            ],
+            "next": "__end__",
+        }
+
+
+# Standalone test function
+def test_finalyzer(
+    llm, query: str, data_descriptions: dict, current_code: str, execution_result: str
+):
+    """
+    Test the finalyzer agent independently.
+
+    Args:
+        llm: LLM instance
+        query: User query
+        data_descriptions: Dict of filename -> description
+        current_code: Working code to finalize
+        execution_result: Current execution result
+
+    Returns:
+        Dictionary with finalyzer results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": query,
+        "data_descriptions": data_descriptions,
+        "plan": [],
+        "current_code": current_code,
+        "execution_result": execution_result,
+        "is_sufficient": True,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "finalyzer",
+    }
+
+    result = finalyzer_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("FINALYZER TEST RESULTS")
+    print("=" * 60)
+    print("Final Code:")
+    print(result.get("current_code", "No code"))
+    print("\nFinal Result:")
+    print(result.get("execution_result", "No result"))
+
+    return result

src/agents/planner_agent.py

@@ -0,0 +1,155 @@
+"""
+Planner Agent: Generates next plan step to answer the query.
+
+This agent generates ONE step at a time based on:
+- The original query
+- Available data files
+- Previously completed steps (if any)
+"""
+
+from langchain_core.messages import AIMessage
+
+from ..utils.formatters import format_data_descriptions, format_plan, gemini_text
+from ..utils.state import DSStarState, PlanStep
+
+
+def planner_node(state: DSStarState) -> dict:
+    """
+    Planner Agent Node: Generates the next step in the plan.
+
+    On first call: Generates initial step to start answering the query
+    On subsequent calls: Generates next step based on progress so far
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - plan: Updated plan with new step appended
+        - messages: Agent communication messages
+        - next: Next node to visit ("coder")
+    """
+    print("=" * 60)
+    print("PLANNER AGENT STARTING...")
+    print("=" * 60)
+
+    is_initial = len(state["plan"]) == 0
+    data_context = format_data_descriptions(state["data_descriptions"])
+
+    if is_initial:
+        print("Generating INITIAL plan step...")
+        prompt = f"""You are an expert data analyst.
+
+Question to answer: {state["query"]}
+
+Available Data Files:
+{data_context}
+
+Task: Generate list of simple, executable steps to start answering this question.
+Examples of good steps:
+- "Load the transactions.csv file"
+- "Read and explore the sales data"
+
+Provide ONLY the step description (one sentence) in one line in bullet points, no explanation."""
+    else:
+        print(f"Generating NEXT step (current plan has {len(state['plan'])} steps)...")
+        plan_text = format_plan(state["plan"])
+
+        prompt = f"""You are an expert data analyst.
+
+Question to answer: {state["query"]}
+
+Available Data Files:
+{data_context}
+
+Current Plan (completed steps):
+{plan_text}
+
+Last Execution Result:
+{state["execution_result"][:500]}...
+
+Task: Suggest the NEXT step to progress toward answering the question.
+Make it simple and executable (one clear action).
+
+Provide ONLY the next step description (one sentence), no explanation."""
+
+    try:
+        # Get LLM response
+        response = state["llm"].invoke(prompt)
+
+        # Handle different response formats
+        if hasattr(response, "content") and isinstance(response.content, list):
+            response_text = gemini_text(response)
+        elif hasattr(response, "content"):
+            response_text = response.content
+        else:
+            response_text = str(response)
+
+        # Create new step
+        new_step = PlanStep(
+            step_number=len(state["plan"]), description=response_text.strip()
+        )
+
+        # Add new step to existing plan
+        updated_plan = state["plan"] + [new_step]
+
+        print(
+            f"\n✓ Generated step {new_step['step_number'] + 1}: {new_step['description']}"
+        )
+        print("=" * 60)
+
+        return {
+            "plan": updated_plan,
+            "messages": [
+                AIMessage(content=f"Added step {new_step['step_number'] + 1}")
+            ],
+            "next": "coder",
+        }
+
+    except Exception as e:
+        print(f"✗ Planner error: {str(e)}")
+        return {
+            "messages": [AIMessage(content=f"Planner error: {str(e)}")],
+            "next": "__end__",
+        }
+
+
+# Standalone test function
+def test_planner(llm, query: str, data_descriptions: dict, existing_plan: list = None):
+    """
+    Test the planner agent independently.
+
+    Args:
+        llm: LLM instance
+        query: User query
+        data_descriptions: Dict of filename -> description
+        existing_plan: Optional existing plan steps
+
+    Returns:
+        Dictionary with planner results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": query,
+        "data_descriptions": data_descriptions,
+        "plan": existing_plan or [],
+        "current_code": "",
+        "execution_result": "",
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "planner",
+    }
+
+    result = planner_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("PLANNER TEST RESULTS")
+    print("=" * 60)
+    print(f"Updated Plan ({len(result.get('plan', []))} steps):")
+    print(format_plan(result.get("plan", [])))
+
+    return result

src/agents/router_agent.py

@@ -0,0 +1,220 @@
+"""
+Router Agent: Decides how to improve an insufficient plan.
+
+When the verifier determines the plan is insufficient, the router decides:
+- "Add Step": Add a new step to the plan
+- "Step N": Backtrack to step N and fix it
+"""
+
+import re
+
+from langchain_core.messages import AIMessage
+
+from ..utils.formatters import format_data_descriptions, format_plan, gemini_text
+from ..utils.state import DSStarState
+
+
+def router_node(state: DSStarState) -> dict:
+    """
+    Router Agent Node: Decides how to improve the plan.
+
+    Analyzes the current situation and determines whether to:
+    1. Add a new step to the plan
+    2. Backtrack and fix an existing step
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - router_decision: "Add Step" or "Step N"
+        - iteration: Incremented iteration count
+        - messages: Agent communication messages
+        - next: "planner" (add step) or "backtrack" (fix step)
+    """
+    print("=" * 60)
+    print("ROUTER AGENT STARTING...")
+    print("=" * 60)
+
+    data_context = format_data_descriptions(state["data_descriptions"])
+    plan_text = format_plan(state["plan"])
+
+    prompt = f"""You are an expert data analyst router.
+
+The current plan is INSUFFICIENT to answer the question.
+
+Original Question: {state["query"]}
+
+Available Data:
+{data_context}
+
+Current Plan:
+{plan_text}
+
+Execution Result:
+{state["execution_result"][:500]}
+
+Task: Decide how to improve the plan:
+1. If a current step is WRONG or needs fixing: Answer "Step N" (where N is the step number, e.g., "Step 2")
+2. If we need to ADD a NEW step: Answer "Add Step"
+
+Answer with ONLY: "Step 1", "Step 2", etc. OR "Add Step"
+No explanation needed."""
+
+    try:
+        # Get LLM response
+        response = state["llm"].invoke(prompt)
+
+        # Handle different response formats
+        if hasattr(response, "content") and isinstance(response.content, list):
+            response_text = gemini_text(response)
+        elif hasattr(response, "content"):
+            response_text = response.content
+        else:
+            response_text = str(response)
+
+        # Parse decision
+        response_lower = response_text.strip().lower()
+        if "add step" in response_lower:
+            decision = "Add Step"
+            next_node = "planner"
+        else:
+            # Try to extract step number
+            match = re.search(r"step\s+(\d+)", response_lower)
+            if match:
+                decision = f"Step {match.group(1)}"
+                next_node = "backtrack"
+            else:
+                # Default to adding new step
+                decision = "Add Step"
+                next_node = "planner"
+
+        print(f"\nRouter Decision: {decision}")
+        print(
+            f"Next Action: {'Backtrack' if next_node == 'backtrack' else 'Add New Step'}"
+        )
+        print("=" * 60)
+
+        return {
+            "router_decision": decision,
+            "messages": [AIMessage(content=f"Router: {decision}")],
+            "iteration": state["iteration"] + 1,
+            "next": next_node,
+        }
+
+    except Exception as e:
+        # On error, default to adding new step
+        print(f"\n✗ Router error: {str(e)}")
+        print("Defaulting to 'Add Step'")
+        return {
+            "router_decision": "Add Step",
+            "messages": [AIMessage(content=f"Router error, adding step: {str(e)}")],
+            "iteration": state["iteration"] + 1,
+            "next": "planner",
+        }
+
+
+def backtrack_node(state: DSStarState) -> dict:
+    """
+    Backtrack Node: Truncates plan to remove incorrect steps.
+
+    When router identifies a wrong step, this node:
+    1. Parses the step number from router_decision
+    2. Truncates the plan to remove that step and all subsequent steps
+    3. Routes back to planner to regenerate from that point
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - plan: Truncated plan
+        - messages: Agent communication messages
+        - next: "planner" to regenerate from truncation point
+    """
+    print("=" * 60)
+    print("BACKTRACK NODE ACTIVATING...")
+    print("=" * 60)
+
+    try:
+        # Extract step number from router decision
+        match = re.search(r"step\s+(\d+)", state["router_decision"].lower())
+        if match:
+            step_num = int(match.group(1))
+        else:
+            # If parsing fails, just add new step
+            print("Failed to parse step number, adding new step instead")
+            return {
+                "messages": [
+                    AIMessage(content="Backtrack parsing failed, adding new step")
+                ],
+                "next": "planner",
+            }
+
+        # Truncate plan to steps before the wrong one
+        # Keep steps 0 to (step_num - 2), which are steps 1 to (step_num - 1) in human counting
+        truncated_plan = state["plan"][: step_num - 1] if step_num > 1 else []
+
+        print(
+            f"Truncating plan from {len(state['plan'])} to {len(truncated_plan)} steps"
+        )
+        print(f"Removed step {step_num} and beyond")
+        print("=" * 60)
+
+        # Return the truncated plan (replaces entire plan, not appends)
+        return {
+            "plan": truncated_plan,
+            "messages": [AIMessage(content=f"Backtracked to step {step_num - 1}")],
+            "next": "planner",
+        }
+
+    except Exception as e:
+        print(f"✗ Backtrack error: {str(e)}")
+        return {
+            "messages": [AIMessage(content=f"Backtrack error: {str(e)}, continuing")],
+            "next": "planner",
+        }
+
+
+# Standalone test function
+def test_router(
+    llm, query: str, data_descriptions: dict, plan: list, execution_result: str
+):
+    """
+    Test the router agent independently.
+
+    Args:
+        llm: LLM instance
+        query: User query
+        data_descriptions: Dict of filename -> description
+        plan: Current plan steps
+        execution_result: Result from code execution
+
+    Returns:
+        Dictionary with router results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": query,
+        "data_descriptions": data_descriptions,
+        "plan": plan,
+        "current_code": "",
+        "execution_result": execution_result,
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "router",
+    }
+
+    result = router_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("ROUTER TEST RESULTS")
+    print("=" * 60)
+    print(f"Decision: {result.get('router_decision', 'unknown')}")
+    print(f"Next Node: {result.get('next', 'unknown')}")
+
+    return result

src/agents/verifier_agent.py

@@ -0,0 +1,143 @@
+"""
+Verifier Agent: Checks if the current plan and implementation sufficiently answer the query.
+
+This agent evaluates whether the work done so far is enough to answer the original question.
+"""
+
+from langchain_core.messages import AIMessage
+
+from ..utils.formatters import format_plan, gemini_text
+from ..utils.state import DSStarState
+
+
+def verifier_node(state: DSStarState) -> dict:
+    """
+    Verifier Agent Node: Determines if plan sufficiently answers the query.
+
+    Analyzes:
+    - Original query
+    - Current plan
+    - Code implementation
+    - Execution results
+
+    Args:
+        state: Current DSStarState
+
+    Returns:
+        Dictionary with updated state fields:
+        - is_sufficient: Boolean indicating if work is complete
+        - messages: Agent communication messages
+        - next: "finalyzer" if sufficient, "router" if not
+    """
+    print("=" * 60)
+    print("VERIFIER AGENT STARTING...")
+    print("=" * 60)
+
+    plan_text = format_plan(state["plan"])
+
+    prompt = f"""You are an expert data analyst verifier.
+
+Original Question: {state["query"]}
+
+Current Plan:
+{plan_text}
+
+Implementation Code:
+{state["current_code"]}
+
+Execution Result:
+{state["execution_result"][:1000]}
+
+Task: Verify if this plan and implementation are SUFFICIENT to fully answer the question.
+
+Consider:
+- Does the plan address all aspects of the question?
+- Does the execution result contain the answer?
+- Is any additional analysis needed?
+
+Answer with ONLY one word: "Yes" or "No"
+- "Yes" if sufficient to answer the question
+- "No" if more analysis is needed"""
+
+    try:
+        # Get LLM response
+        response = state["llm"].invoke(prompt)
+
+        # Handle different response formats
+        if hasattr(response, "content") and isinstance(response.content, list):
+            response_text = gemini_text(response)
+        elif hasattr(response, "content"):
+            response_text = response.content
+        else:
+            response_text = str(response)
+
+        response_lower = response_text.strip().lower()
+        is_sufficient = "yes" in response_lower
+
+        status = "SUFFICIENT ✓" if is_sufficient else "INSUFFICIENT ✗"
+        print(f"\nVerification Result: {status}")
+        print("=" * 60)
+
+        next_node = "finalyzer" if is_sufficient else "router"
+
+        return {
+            "is_sufficient": is_sufficient,
+            "messages": [
+                AIMessage(
+                    content=f"Verification: {'Sufficient' if is_sufficient else 'Insufficient'}"
+                )
+            ],
+            "next": next_node,
+        }
+
+    except Exception as e:
+        # On error, assume insufficient and continue
+        print(f"\n✗ Verifier error: {str(e)}")
+        print("Defaulting to insufficient, continuing...")
+        return {
+            "is_sufficient": False,
+            "messages": [AIMessage(content=f"Verifier error: {str(e)}, continuing...")],
+            "next": "router",
+        }
+
+
+# Standalone test function
+def test_verifier(llm, query: str, plan: list, code: str, execution_result: str):
+    """
+    Test the verifier agent independently.
+
+    Args:
+        llm: LLM instance
+        query: User query
+        plan: List of plan steps
+        code: Generated code
+        execution_result: Result from code execution
+
+    Returns:
+        Dictionary with verifier results
+    """
+    # Create minimal test state
+    test_state = {
+        "llm": llm,
+        "query": query,
+        "data_descriptions": {},
+        "plan": plan,
+        "current_code": code,
+        "execution_result": execution_result,
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": 20,
+        "messages": [],
+        "next": "verifier",
+    }
+
+    result = verifier_node(test_state)
+
+    print("\n" + "=" * 60)
+    print("VERIFIER TEST RESULTS")
+    print("=" * 60)
+    print(f"Is Sufficient: {result.get('is_sufficient', False)}")
+    print(f"Next Node: {result.get('next', 'unknown')}")
+
+    return result

src/config/__init__.py

@@ -0,0 +1,5 @@
+"""Configuration package."""
+
+from .llm_config import DEFAULT_CONFIG, get_llm
+
+__all__ = ["get_llm", "DEFAULT_CONFIG"]

src/config/llm_config.py

@@ -0,0 +1,90 @@
+"""
+Configuration for DS-STAR system.
+Centralizes LLM setup and system parameters.
+"""
+
+import os
+from typing import Optional
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def get_llm(
+    provider: str = "google",
+    model: Optional[str] = None,
+    api_key: Optional[str] = None,
+    temperature: float = 0,
+    base_url: Optional[str] = None,
+):
+    """
+    Get configured LLM instance.
+
+    Args:
+        provider: LLM provider ("google", "openai", "anthropic")
+        model: Model name (uses default if not specified)
+        api_key: API key (uses environment variable if not specified)
+        temperature: Temperature for generation
+        base_url: Custom base URL for OpenAI-compatible APIs
+
+    Returns:
+        Configured LLM instance
+    """
+    if provider == "google":
+        from langchain_google_genai import ChatGoogleGenerativeAI
+
+        default_model = "gemini-flash-latest"
+        api_key = api_key or os.getenv("GOOGLE_API_KEY", "")
+
+        return ChatGoogleGenerativeAI(
+            model=model or default_model,
+            temperature=temperature,
+            google_api_key=api_key,
+        )
+
+    elif provider == "openai":
+        from langchain_openai import ChatOpenAI
+
+        default_model = "gpt-4"
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+
+        # Use provided base_url, then env var, then default
+        effective_base_url = base_url or os.getenv(
+            "LLM_BASE_URL",
+            "https://api.openai.com/v1",
+        )
+
+        return ChatOpenAI(
+            model=model or default_model,
+            temperature=temperature,
+            api_key=api_key,
+            base_url=effective_base_url,
+        )
+
+    elif provider == "anthropic":
+        from langchain_anthropic import ChatAnthropic
+
+        default_model = "claude-3-5-sonnet-20241022"
+        api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
+
+        return ChatAnthropic(
+            model=model or default_model,
+            temperature=temperature,
+            api_key=api_key,
+        )
+
+    else:
+        raise ValueError(
+            f"Unknown provider: {provider}. Choose from: google, openai, anthropic"
+        )
+
+
+# Default configuration
+DEFAULT_CONFIG = {
+    "max_iterations": 20,
+    "provider": "openai",
+    "model": "google/gemini-2.5-flash",
+    "temperature": 0,
+    "data_dir": "data/",
+}

src/graph.py

@@ -0,0 +1,232 @@
+"""
+DS-STAR Graph: Connects all agents into a workflow.
+
+This module builds the LangGraph StateGraph that orchestrates the multi-agent system.
+"""
+
+from typing import Literal
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, StateGraph
+
+from .agents.analyzer_agent import analyzer_node
+from .agents.coder_agent import coder_node
+from .agents.finalyzer_agent import finalyzer_node
+from .agents.planner_agent import planner_node
+from .agents.router_agent import backtrack_node, router_node
+from .agents.verifier_agent import verifier_node
+from .utils.state import DSStarState
+
+# ==================== CONDITIONAL ROUTING FUNCTIONS ====================
+
+
+def route_after_analyzer(state: DSStarState) -> Literal["planner", "__end__"]:
+    """
+    Route after analyzer based on success.
+
+    If analyzer found errors, end workflow.
+    Otherwise, proceed to planner.
+    """
+    if "error" in state.get("data_descriptions", {}):
+        return "__end__"
+    return state.get("next", "planner")
+
+
+def route_after_planner(state: DSStarState) -> Literal["coder", "__end__"]:
+    """Route after planner to coder."""
+    return state.get("next", "coder")
+
+
+def route_after_coder(state: DSStarState) -> Literal["verifier", "__end__"]:
+    """Route after coder to verifier."""
+    return state.get("next", "verifier")
+
+
+def route_after_verifier(
+    state: DSStarState,
+) -> Literal["router", "finalyzer", "__end__"]:
+    """
+    Route after verifier based on sufficiency and iteration count.
+
+    If max iterations reached, go to finalyzer.
+    If sufficient, go to finalyzer.
+    Otherwise, go to router to decide next action.
+    """
+    # Check max iterations
+    if state["iteration"] >= state["max_iterations"]:
+        print(f"\n⚠ Max iterations ({state['max_iterations']}) reached, finalizing...")
+        return "finalyzer"
+
+    return state.get("next", "router")
+
+
+def route_after_router(state: DSStarState) -> Literal["planner", "backtrack"]:
+    """
+    Route after router based on decision.
+
+    If router says "Add Step", go to planner.
+    If router says "Step N", go to backtrack.
+    """
+    return state.get("next", "planner")
+
+
+# ==================== GRAPH BUILDER ====================
+
+
+def build_ds_star_graph(llm, max_iterations: int = 20):
+    """
+    Constructs the LangGraph workflow for DS-STAR.
+
+    The workflow follows this pattern:
+    1. Analyzer: Analyze data files (runs once)
+    2. Planner: Generate next plan step
+    3. Coder: Implement plan as code
+    4. Verifier: Check if sufficient
+    5. If insufficient:
+       a. Router: Decide to add step or backtrack
+       b. Backtrack (optional): Remove wrong steps
+       c. Go back to Planner
+    6. If sufficient:
+       Finalyzer: Create polished final solution
+
+    Args:
+        llm: LLM instance (e.g., ChatOpenAI, ChatGoogleGenerativeAI)
+        max_iterations: Maximum refinement iterations (default: 20)
+
+    Returns:
+        Compiled LangGraph application with checkpointing
+    """
+    # Initialize graph with state schema
+    workflow = StateGraph(DSStarState)
+
+    # Add all agent nodes
+    workflow.add_node("analyzer", analyzer_node)
+    workflow.add_node("planner", planner_node)
+    workflow.add_node("coder", coder_node)
+    workflow.add_node("verifier", verifier_node)
+    workflow.add_node("router", router_node)
+    workflow.add_node("backtrack", backtrack_node)
+    workflow.add_node("finalyzer", finalyzer_node)
+
+    # Set entry point
+    workflow.set_entry_point("analyzer")
+
+    # Add conditional edges with proper routing
+    workflow.add_conditional_edges(
+        "analyzer", route_after_analyzer, {"planner": "planner", "__end__": END}
+    )
+
+    workflow.add_conditional_edges(
+        "planner", route_after_planner, {"coder": "coder", "__end__": END}
+    )
+
+    workflow.add_conditional_edges(
+        "coder", route_after_coder, {"verifier": "verifier", "__end__": END}
+    )
+
+    workflow.add_conditional_edges(
+        "verifier",
+        route_after_verifier,
+        {"router": "router", "finalyzer": "finalyzer", "__end__": END},
+    )
+
+    workflow.add_conditional_edges(
+        "router", route_after_router, {"planner": "planner", "backtrack": "backtrack"}
+    )
+
+    workflow.add_edge("backtrack", "planner")
+    workflow.add_edge("finalyzer", END)
+
+    # Add memory/checkpointing
+    memory = MemorySaver()
+
+    # Compile graph
+    app = workflow.compile(checkpointer=memory)
+
+    return app
+
+
+def create_initial_state(query: str, llm, max_iterations: int = 20) -> DSStarState:
+    """
+    Create initial state for the DS-STAR workflow.
+
+    Args:
+        query: User's question to answer
+        llm: LLM instance
+        max_iterations: Maximum refinement iterations
+
+    Returns:
+        Initial DSStarState dictionary
+    """
+    return {
+        "query": query,
+        "data_descriptions": {},
+        "plan": [],
+        "current_code": "",
+        "execution_result": "",
+        "is_sufficient": False,
+        "router_decision": "",
+        "iteration": 0,
+        "max_iterations": max_iterations,
+        "messages": [],
+        "next": "analyzer",
+        "llm": llm,
+    }
+
+
+def run_ds_star(
+    query: str, llm, max_iterations: int = 20, thread_id: str = "ds-star-1"
+):
+    """
+    Run the complete DS-STAR workflow.
+
+    Args:
+        query: User's question to answer
+        llm: LLM instance
+        max_iterations: Maximum refinement iterations
+        thread_id: Unique thread ID for checkpointing
+
+    Returns:
+        Final state after workflow completion
+    """
+    print("=" * 60)
+    print("DS-STAR MULTI-AGENT SYSTEM")
+    print("=" * 60)
+    print(f"Query: {query}")
+    print(f"Max Iterations: {max_iterations}")
+    print("=" * 60)
+
+    # Build graph
+    app = build_ds_star_graph(llm, max_iterations)
+
+    # Create initial state
+    initial_state = create_initial_state(query, llm, max_iterations)
+
+    # Run with checkpointing
+    config = {"configurable": {"thread_id": thread_id}}
+
+    try:
+        # Execute the workflow
+        final_state = app.invoke(initial_state, config)
+
+        # Display results
+        print("\n" + "=" * 60)
+        print("FINAL SOLUTION")
+        print("=" * 60)
+        print("\nGenerated Code:")
+        print("-" * 60)
+        print(final_state["current_code"])
+        print("\n" + "-" * 60)
+        print("Execution Result:")
+        print("-" * 60)
+        print(final_state["execution_result"])
+        print("=" * 60)
+
+        return final_state
+
+    except Exception as e:
+        print(f"\n✗ Error during execution: {str(e)}")
+        import traceback
+
+        traceback.print_exc()
+        return None

src/utils/__init__.py

@@ -0,0 +1,16 @@
+"""Utility modules for DS-STAR system."""
+
+from .code_execution import execute_code_safely, execute_with_debug
+from .formatters import extract_code, format_data_descriptions, format_plan, gemini_text
+from .state import DSStarState, PlanStep
+
+__all__ = [
+    "DSStarState",
+    "PlanStep",
+    "extract_code",
+    "format_data_descriptions",
+    "format_plan",
+    "gemini_text",
+    "execute_code_safely",
+    "execute_with_debug",
+]

src/utils/code_execution.py

@@ -0,0 +1,129 @@
+"""
+Code execution utilities with debugging capabilities.
+"""
+
+import os
+import subprocess
+import sys
+from typing import Tuple
+
+
+def execute_code_safely(code: str, timeout: int = 30) -> Tuple[bool, str, str]:
+    """
+    Execute Python code safely in a subprocess and capture output.
+
+    Args:
+        code: Python code to execute
+        timeout: Maximum execution time in seconds
+
+    Returns:
+        Tuple of (success: bool, stdout: str, stderr: str)
+    """
+    temp_file = "temp_script.py"
+
+    try:
+        # Write code to temporary file
+        with open(temp_file, "w", encoding="utf-8") as f:
+            f.write(code)
+
+        # Execute with subprocess
+        result = subprocess.run(
+            [sys.executable, temp_file], capture_output=True, text=True, timeout=timeout
+        )
+
+        success = result.returncode == 0
+        return success, result.stdout, result.stderr
+
+    except subprocess.TimeoutExpired:
+        return False, "", f"Execution timed out after {timeout} seconds"
+    except Exception as e:
+        return False, "", f"Execution error: {str(e)}"
+    finally:
+        # Clean up temp file
+        if os.path.exists(temp_file):
+            try:
+                os.remove(temp_file)
+            except Exception as e:
+                print(f"Warning: Failed to remove temp file: {str(e)}")
+
+
+def execute_with_debug(
+    code: str, llm, is_analysis: bool, data_context: str = "", max_attempts: int = 3
+) -> str:
+    """
+    Execute code with automatic debugging via LLM.
+
+    If execution fails, the LLM is asked to fix the error.
+    This repeats for up to max_attempts.
+
+    Args:
+        code: Python code to execute
+        llm: LLM instance for debugging
+        is_analysis: Whether this is data analysis stage (simpler prompts)
+        data_context: Context about available data files
+        max_attempts: Maximum debugging attempts
+
+    Returns:
+        Execution output or error message
+    """
+    from .formatters import extract_code
+
+    for attempt in range(max_attempts):
+        success, stdout, stderr = execute_code_safely(code)
+
+        if success:
+            return stdout if stdout else "Code executed successfully (no output)"
+
+        # Debug the error
+        print(f"  Debug attempt {attempt + 1}/{max_attempts}")
+
+        if is_analysis:
+            debug_prompt = f"""Fix this Python code error:
+{code}
+
+Error:
+{stderr}
+
+Requirements:
+- Fix the error
+- Keep the same functionality
+- No try-except blocks
+- All files are in 'data/' directory
+- Provide ONLY the corrected code in a markdown code block"""
+        else:
+            debug_prompt = f"""Fix this Python code error:
+
+Available Data:
+{data_context}
+
+Code with error:
+{code}
+
+Error:
+{stderr}
+
+Requirements:
+- Fix the error using data context
+- Keep the same functionality
+- No try-except blocks
+- Provide ONLY the corrected code in a markdown code block"""
+
+        try:
+            response = llm.invoke(debug_prompt)
+
+            # Handle Gemini response format
+            if hasattr(response, "content") and isinstance(response.content, list):
+                # Gemini returns list of dicts
+                from .formatters import gemini_text
+
+                response_text = gemini_text(response)
+            elif hasattr(response, "content"):
+                response_text = response.content
+            else:
+                response_text = str(response)
+
+            code = extract_code(response_text)
+        except Exception as e:
+            return f"Debugging failed: {str(e)}"
+
+    return f"Failed after {max_attempts} attempts. Last error:\n{stderr}"

src/utils/formatters.py

@@ -0,0 +1,89 @@
+"""
+Utility functions for formatting and extracting data.
+"""
+
+
+def gemini_text(res) -> str:
+    """
+    Extract text from Gemini API response.
+
+    Args:
+        res: Gemini API response object
+
+    Returns:
+        Concatenated text from all text parts
+    """
+    return "".join(part["text"] for part in res.content if part.get("type") == "text")
+
+
+def extract_code(response: str) -> str:
+    """
+    Extract code from markdown code blocks.
+    Handles ```python and ``` formats.
+
+    Args:
+        response: Text response potentially containing code blocks
+
+    Returns:
+        Extracted code string
+    """
+    # Try to find python code block first
+    if "```python" in response:
+        parts = response.split("```python", 1)
+        if len(parts) > 1:
+            code_part = parts[1].split("```", 1)
+            if len(code_part) > 0:
+                return code_part[0].strip()
+
+    # Try generic code block
+    elif "```" in response:
+        parts = response.split("```", 1)
+        if len(parts) > 1:
+            code_part = parts[1].split("```", 1)
+            if len(code_part) > 0:
+                # Remove language identifier if present
+                code = code_part[0].strip()
+                # Remove first line if it's a language identifier
+                lines = code.split("\n")
+                if lines and lines[0].strip() in ["python", "py", "python3"]:
+                    return "\n".join(lines[1:]).strip()
+                return code
+
+    # If no code blocks found, return as is
+    return response.strip()
+
+
+def format_data_descriptions(descriptions: dict) -> str:
+    """
+    Format data descriptions dictionary into readable string.
+
+    Args:
+        descriptions: Dict mapping filename to description
+
+    Returns:
+        Formatted string with file descriptions
+    """
+    if not descriptions:
+        return "No data files analyzed yet."
+
+    formatted_parts = []
+    for filename, description in descriptions.items():
+        formatted_parts.append(f"## File: {filename}\n{description}\n")
+
+    return "\n".join(formatted_parts)
+
+
+def format_plan(plan: list) -> str:
+    """
+    Format plan steps into readable string.
+
+    Args:
+        plan: List of PlanStep dictionaries
+
+    Returns:
+        Formatted plan string
+    """
+    if not plan:
+        return "No plan steps yet."
+
+    return "\n".join([f"{i + 1}. {step['description']}" for i, step in enumerate(plan)])

src/utils/state.py

@@ -0,0 +1,56 @@
+"""
+State schema for DS-STAR multi-agent system.
+Defines the centralized state structure shared across all agents.
+"""
+
+import operator
+from typing import Annotated, List, TypedDict
+
+from langchain_core.messages import AIMessage, HumanMessage
+
+
+class PlanStep(TypedDict):
+    """Individual plan step with number and description"""
+
+    step_number: int
+    description: str
+
+
+class DSStarState(TypedDict):
+    """
+    Centralized state for DS-STAR pipeline.
+
+    This state is passed between all agents in the graph.
+    Uses reducer pattern for message accumulation.
+    """
+
+    # User input
+    query: str
+
+    # Data file descriptions (from Stage 1 - Analyzer)
+    data_descriptions: dict[str, str]
+
+    # Current plan - list of completed steps
+    # NO REDUCER - we need full control for backtracking
+    plan: List[PlanStep]
+
+    # Code and execution
+    current_code: str
+    execution_result: str
+
+    # Verification and routing
+    is_sufficient: bool
+    router_decision: str  # "Add Step" or "Step N"
+
+    # Iteration tracking
+    iteration: int
+    max_iterations: int
+
+    # Messages for agent communication (accumulated with reducer)
+    messages: Annotated[List[HumanMessage | AIMessage], operator.add]
+
+    # Next node routing (for internal control flow)
+    next: str
+
+    # LLM instance (shared across all agents)
+    llm: object

tests/test_complete_workflow.py

@@ -0,0 +1,112 @@
+"""
+Comprehensive test for DS-STAR workflow.
+
+This test runs the complete multi-agent system to verify:
+1. All agents are properly connected
+2. The graph routing works correctly
+3. The workflow can complete successfully
+"""
+
+import os
+import sys
+
+from dotenv import load_dotenv
+
+load_dotenv()
+LLM_MODEL = os.getenv("LLM_MODEL", "google/gemini-2.5-flash")
+LLM_API_KEY = os.getenv("LLM_API_KEY", "")
+
+# Add parent directory to path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from src.config import get_llm
+from src.graph import run_ds_star
+
+
+def test_complete_workflow():
+    """Test the complete DS-STAR workflow."""
+    print("=" * 60)
+    print("COMPREHENSIVE DS-STAR WORKFLOW TEST")
+    print("=" * 60)
+
+    # Configuration
+    query = "What percentage of transactions use credit cards?"
+    max_iterations = 10  # Reduced for testing
+
+    print(f"\nTest Query: {query}")
+    print(f"Max Iterations: {max_iterations}")
+    print()
+
+    try:
+        # Initialize LLM
+        print("Initializing LLM (Gemini 1.5 Flash)...")
+        llm = get_llm(
+            provider="openai",
+            model=LLM_MODEL,
+            temperature=0,
+            api_key=LLM_API_KEY,
+        )
+
+        print("✓ LLM initialized")
+        print()
+
+        # Run workflow
+        print("Starting DS-STAR workflow...")
+        print("=" * 60)
+
+        final_state = run_ds_star(
+            query=query,
+            llm=llm,
+            max_iterations=max_iterations,
+            thread_id="test-session",
+        )
+
+        # Verify results
+        print("\n" + "=" * 60)
+        print("TEST VERIFICATION")
+        print("=" * 60)
+
+        if final_state is None:
+            print("❌ FAILED: Workflow returned None")
+            return False
+
+        # Check that we got results
+        checks = [
+            ("Data descriptions", len(final_state.get("data_descriptions", {})) > 0),
+            ("Plan generated", len(final_state.get("plan", [])) > 0),
+            ("Code generated", len(final_state.get("current_code", "")) > 0),
+            ("Execution result", len(final_state.get("execution_result", "")) > 0),
+        ]
+
+        all_passed = True
+        for check_name, passed in checks:
+            status = "✓" if passed else "✗"
+            print(f"{status} {check_name}: {'PASS' if passed else 'FAIL'}")
+            all_passed = all_passed and passed
+
+        print("\n" + "=" * 60)
+        if all_passed:
+            print("✅ ALL TESTS PASSED")
+            print("=" * 60)
+            return True
+        else:
+            print("❌ SOME TESTS FAILED")
+            print("=" * 60)
+            return False
+
+    except Exception as e:
+        print(f"\n❌ TEST FAILED WITH EXCEPTION: {str(e)}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+
+def main():
+    """Run the test."""
+    success = test_complete_workflow()
+    return 0 if success else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())