RonyForAI/Mirage_DB_RL
0
1import streamlit as st2import time3import sys4import os5 6# Add project root to path so Mirage_RL package is found7sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))8 9from Mirage_RL.client import QueryClient10from Mirage_RL.models import QueryAction11from Mirage_RL.training.agent import Agent12 13# ─── Page Config ──────────────────────────────────────────────────────────────14st.set_page_config(15 page_title="Mirage RL — DBMS Query Optimizer",16 page_icon="🧠",17 layout="wide",18 initial_sidebar_state="expanded",19)20 21# ─── CSS ──────────────────────────────────────────────────────────────────────22st.markdown("""23<style>24@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');25 26html, body, [class*="css"] { font-family: 'Inter', sans-serif; }27 28.metric-card {29 background: linear-gradient(135deg, #1e1e2e 0%, #2a2a3e 100%);30 border: 1px solid #3a3a5e;31 border-radius: 12px;32 padding: 20px;33 text-align: center;34 box-shadow: 0 4px 20px rgba(0,0,0,0.3);35}36.metric-card h3 { color: #a0aec0; font-size: 0.85rem; font-weight: 400; margin: 0 0 8px 0; letter-spacing: 0.08em; text-transform: uppercase; }37.metric-card .value { font-size: 2rem; font-weight: 700; color: #7c3aed; }38.metric-card .sub { font-size: 0.75rem; color: #718096; margin-top: 4px; }39 40.step-card {41 background: linear-gradient(135deg, #0f3460 0%, #16213e 100%);42 border-left: 4px solid #7c3aed;43 border-radius: 8px;44 padding: 16px 20px;45 margin: 8px 0;46 animation: fadeIn 0.4s ease;47}48.step-card .step-title { color: #e2e8f0; font-weight: 600; font-size: 1rem; }49.step-card .step-detail { color: #94a3b8; font-size: 0.85rem; margin-top: 6px; }50.step-card .reward-pos { color: #10b981; font-weight: 600; }51.step-card .reward-neg { color: #f87171; font-weight: 600; }52 53.table-chip {54 display: inline-block;55 background: #7c3aed22;56 border: 1px solid #7c3aed66;57 color: #a78bfa;58 border-radius: 20px;59 padding: 3px 12px;60 font-size: 0.8rem;61 margin: 2px;62 font-weight: 600;63}64.table-chip.chosen {65 background: #10b98122;66 border-color: #10b98166;67 color: #34d399;68}69 70.badge-hash { background: #3b82f622; border: 1px solid #3b82f666; color: #60a5fa; border-radius: 6px; padding: 2px 8px; font-size: 0.75rem; }71.badge-nested { background: #f59e0b22; border: 1px solid #f59e0b66; color: #fbbf24; border-radius: 6px; padding: 2px 8px; font-size: 0.75rem; }72.badge-merge { background: #10b98122; border: 1px solid #10b98166; color: #34d399; border-radius: 6px; padding: 2px 8px; font-size: 0.75rem; }73 74.hero-banner {75 background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);76 border: 1px solid #3a3a5e;77 border-radius: 16px;78 padding: 32px;79 margin-bottom: 24px;80 text-align: center;81}82.hero-banner h1 { color: #e2e8f0; font-size: 2.2rem; font-weight: 700; margin: 0; }83.hero-banner p { color: #94a3b8; font-size: 1rem; margin: 10px 0 0 0; }84 85.pill-green { background: #10b98122; border: 1px solid #10b981; color: #34d399; border-radius: 20px; padding: 4px 14px; font-size: 0.8rem; font-weight: 600; display: inline-block; }86.pill-red { background: #ef444422; border: 1px solid #ef4444; color: #f87171; border-radius: 20px; padding: 4px 14px; font-size: 0.8rem; font-weight: 600; display: inline-block; }87 88@keyframes fadeIn { from { opacity: 0; transform: translateY(8px); } to { opacity: 1; transform: translateY(0); } }89</style>90""", unsafe_allow_html=True)91 92# ─── Hero ─────────────────────────────────────────────────────────────────────93st.markdown("""94<div class="hero-banner">95 <h1>🧠 Mirage RL — DBMS Query Optimizer</h1>96 <p>Reinforcement Learning agent learns the optimal join order, join strategy, and index usage for SQL queries</p>97</div>98""", unsafe_allow_html=True)99 100# ─── Sidebar ──────────────────────────────────────────────────────────────────101with st.sidebar:102 st.markdown("## ⚙️ Configuration")103 server_url = st.text_input("Server URL", value="http://localhost:8000")104 mode = st.radio("🎮 Mode", ["AI Agent (Trained)", "Manual Control"])105 st.markdown("---")106 st.markdown("### 📖 Join Types")107 st.markdown("**0 — Hash Join** · Fast for large, unordered sets")108 st.markdown("**1 — Nested Loop** · Good for small tables or indexed lookups")109 st.markdown("**2 — Merge Join** · Efficient for sorted/pre-sorted data")110 st.markdown("---")111 st.markdown("### 🗄️ Tables in Query")112 st.markdown("| Table | Rows | Selectivity | Index |")113 st.markdown("|---|---|---|---|")114 st.markdown("| A | 1,000 | 10% | ✅ |")115 st.markdown("| B | 5,000 | 50% | ❌ |")116 st.markdown("| C | 200 | 5% | ✅ |")117 118# ─── Connection check ─────────────────────────────────────────────────────────119import requests120col_status, col_btn = st.columns([3, 1])121with col_status:122 try:123 r = requests.get(f"{server_url}/health", timeout=2)124 st.markdown('<span class="pill-green">🟢 Server Connected</span>', unsafe_allow_html=True)125 server_ok = True126 except Exception:127 try:128 r = requests.get(f"{server_url}/docs", timeout=2)129 st.markdown('<span class="pill-green">🟢 Server Connected</span>', unsafe_allow_html=True)130 server_ok = True131 except Exception:132 st.markdown('<span class="pill-red">🔴 Server Offline — open new terminal and run: uv run server</span>', unsafe_allow_html=True)133 server_ok = False134 135agent = Agent(num_tables=3)136 137# ─── Run button ───────────────────────────────────────────────────────────────138st.markdown("<br>", unsafe_allow_html=True)139run_col, _ = st.columns([1, 3])140with run_col:141 run_clicked = st.button("▶ Run Simulation", type="primary", use_container_width=True, disabled=not server_ok)142 143if run_clicked:144 env = QueryClient(base_url=server_url).sync().__enter__()145 result = env.reset()146 obs = result.observation147 148 # ── Top metrics row ──149 st.markdown("<br>", unsafe_allow_html=True)150 st.markdown("### 📊 Query Environment")151 m1, m2, m3, m4 = st.columns(4)152 m1.markdown(f'<div class="metric-card"><h3>Tables</h3><div class="value">{len(obs.tables)}</div><div class="sub">A, B, C</div></div>', unsafe_allow_html=True)153 m2.markdown(f'<div class="metric-card"><h3>Total Rows</h3><div class="value">{sum(obs.table_rows):,}</div><div class="sub">across all tables</div></div>', unsafe_allow_html=True)154 m3.markdown(f'<div class="metric-card"><h3>Indexed Tables</h3><div class="value">{sum(obs.has_index)}/{len(obs.tables)}</div><div class="sub">A and C</div></div>', unsafe_allow_html=True)155 m4.markdown(f'<div class="metric-card"><h3>Mode</h3><div class="value">{"🤖" if "AI" in mode else "🕹️"}</div><div class="sub">{mode.split("(")[0].strip()}</div>', unsafe_allow_html=True)156 157 st.markdown("<br>", unsafe_allow_html=True)158 159 # ── Table details ──160 st.markdown("### 🗂️ Table Details")161 table_cols = st.columns(len(obs.tables))162 table_names = obs.tables163 for i, col in enumerate(table_cols):164 with col:165 st.markdown(f"""166 <div class="metric-card">167 <h3>Table {table_names[i]}</h3>168 <div class="value" style="font-size:1.4rem">{obs.table_rows[i]:,}</div>169 <div class="sub">rows</div>170 <div class="sub" style="margin-top:8px">Selectivity: <b style="color:#a78bfa">{obs.selectivities[i]*100:.0f}%</b></div>171 <div class="sub">Index: <b style="color:{'#34d399' if obs.has_index[i] else '#f87171'}">{'✅ Yes' if obs.has_index[i] else '❌ No'}</b></div>172 </div>173 """, unsafe_allow_html=True)174 175 # ── Simulation loop ──176 st.markdown("<br>", unsafe_allow_html=True)177 st.markdown("### 🔄 Step-by-Step Execution")178 179 done = False180 step = 0181 total_reward = 0.0182 step_log = []183 184 cost_chart_placeholder = st.empty()185 steps_placeholder = st.empty()186 187 join_names = {0: "Hash Join", 1: "Nested Loop", 2: "Merge Join"}188 join_badges = {0: "badge-hash", 1: "badge-nested", 2: "badge-merge"}189 190 while not done:191 step += 1192 193 if "Manual" in mode:194 st.markdown(f"**Step {step} — Choose action for remaining tables: {[table_names[t] for t in obs.remaining_tables]}**")195 col_a, col_b, col_c = st.columns(3)196 with col_a:197 table_choice = st.selectbox("Next Table", obs.remaining_tables,198 format_func=lambda x: f"Table {table_names[x]} ({obs.table_rows[x]:,} rows)",199 key=f"t{step}")200 with col_b:201 join_choice = st.selectbox("Join Type", [0, 1, 2],202 format_func=lambda x: join_names[x], key=f"j{step}")203 with col_c:204 index_choice = st.selectbox("Use Index", [0, 1],205 format_func=lambda x: "Yes ✅" if x else "No ❌", key=f"i{step}")206 table, join, index = table_choice, join_choice, index_choice207 else:208 (table, join, index), _ = agent.select_action(obs)209 210 action = QueryAction(next_table=table, join_type=join, use_index=index)211 result = env.step(action)212 obs = result.observation213 done = result.done214 reward = result.reward215 total_reward += reward216 217 step_log.append({218 "step": step,219 "table": table_names[table],220 "join": join_names[join],221 "index": bool(index),222 "cost": obs.current_cost,223 "reward": reward,224 })225 226 # Rebuild steps display227 steps_html = ""228 for s in step_log:229 r_class = "reward-pos" if s["reward"] >= 0 else "reward-neg"230 r_sign = "+" if s["reward"] >= 0 else ""231 badge_cls = join_badges[list(join_names.values()).index(s["join"])]232 idx_label = "Index ✅" if s["index"] else "No Index ❌"233 steps_html += f"""234 <div class="step-card">235 <div class="step-title">Step {s['step']} → Join Table <span class="table-chip chosen">{s['table']}</span></div>236 <div class="step-detail">237 <span class="{badge_cls}">{s['join']}</span> 238 {idx_label} 239 Cumulative Cost: <b style="color:#e2e8f0">{s['cost']:.1f}</b>240 Reward: <span class="{r_class}">{r_sign}{s['reward']:.1f}</span>241 </div>242 </div>"""243 244 steps_placeholder.markdown(steps_html, unsafe_allow_html=True)245 246 if "AI" in mode:247 time.sleep(0.6)248 249 # ── Final Results ──250 st.markdown("<br>", unsafe_allow_html=True)251 st.markdown("### 🏆 Final Results")252 253 final_cost = obs.current_cost254 chosen_order = [table_names[i] for i in obs.chosen_order]255 256 r1, r2, r3 = st.columns(3)257 r1.markdown(f'<div class="metric-card"><h3>Final Query Cost</h3><div class="value" style="color:#f87171">{final_cost:.1f}</div><div class="sub">lower is better</div></div>', unsafe_allow_html=True)258 r2.markdown(f'<div class="metric-card"><h3>Total Reward</h3><div class="value" style="color:#{"10b981" if total_reward >= 0 else "f87171"}">{total_reward:.1f}</div><div class="sub">agent performance</div></div>', unsafe_allow_html=True)259 r3.markdown(f'<div class="metric-card"><h3>Join Order</h3><div class="value" style="font-size:1.4rem">{"→".join(chosen_order)}</div><div class="sub">chosen sequence</div></div>', unsafe_allow_html=True)260 261 st.markdown("<br>", unsafe_allow_html=True)262 263 # ── Cost breakdown chart ──264 st.markdown("### 📈 Cost Accumulation Per Step")265 import pandas as pd266 df = pd.DataFrame(step_log)267 st.line_chart(df.set_index("step")["cost"], use_container_width=True)268 269 st.markdown("### 📋 Full Step Log")270 st.dataframe(df.rename(columns={271 "step": "Step", "table": "Table Joined", "join": "Join Type",272 "index": "Used Index", "cost": "Cumulative Cost", "reward": "Reward"273 }), use_container_width=True)274 275 if total_reward > -500:276 st.success(f"✅ Simulation complete! Final cost: **{final_cost:.1f}** | Join order: **{'→'.join(chosen_order)}**")277 else:278 st.warning(f"⚠️ High cost path taken. Final cost: **{final_cost:.1f}**. Try AI Agent mode for better results.")