fix: GEPA current_generation, task_id mapping, Connect DB button, remove difficulty from header

backend/api/demo.py

@@ -27,7 +27,16 @@ from env.database import (
     get_schema_info,
     get_schema_graph,
     execute_query,
+    connect_external_db,
+    get_active_db_label,
 )
+
+# Map frontend difficulty names → backend task IDs
+_DIFFICULTY_MAP = {
+    "easy": "simple_queries",
+    "medium": "join_queries",
+    "hard": "complex_queries",
+}
 from env.tasks import TASKS, get_task
 from env.sql_env import SQLAgentEnv, Action, get_env, BASE_SYSTEM_PROMPT, _clean_sql
 from rl.environment import get_bandit_state
@@ -46,7 +55,22 @@ router = APIRouter()
 async def init_db():
     seeded = ensure_seeded()
     tables = get_table_stats()
-    return {"tables": tables, "seeded": seeded}
+    return {"tables": tables, "seeded": seeded, "dbLabel": get_active_db_label()}
+
+
+# ─── /api/connect-db ──────────────────────────────────────────────
+
+class ConnectDbRequest(BaseModel):
+    path: str  # SQLite file path or :memory:
+
+
+@router.post("/connect-db")
+async def connect_db(req: ConnectDbRequest):
+    success, message = connect_external_db(req.path)
+    if success:
+        tables = get_table_stats()
+        return {"success": True, "message": message, "tables": tables, "dbLabel": get_active_db_label()}
+    return {"success": False, "message": message, "tables": [], "dbLabel": get_active_db_label()}
 
 
 # ─── /api/execute-query ───────────────────────────────────────────
@@ -60,10 +84,12 @@ class ExecuteQueryRequest(BaseModel):
 async def execute_query_stream(req: ExecuteQueryRequest):
     async def event_generator() -> AsyncIterator[dict]:
         env = get_env()
-        obs = env.reset(req.task_id)
+        # Accept difficulty names ('easy'/'medium'/'hard') or direct task IDs
+        task_id = _DIFFICULTY_MAP.get(req.task_id, req.task_id)
+        obs = env.reset(task_id)
 
         # Pick first question of task matching question text, or default
-        task = get_task(req.task_id)
+        task = get_task(task_id)
         question_obj = task.questions[0]
         # Override question text
         env._episode.question = req.question  # type: ignore[union-attr]
@@ -155,7 +181,7 @@ async def execute_query_stream(req: ExecuteQueryRequest):
 
             from env.tasks import grade_response
             task_score = grade_response(
-                req.task_id, question_obj.id, generated_sql, rows, error, attempt
+                task_id, question_obj.id, generated_sql, rows, error, attempt
             )
             attempt_success = task_score >= 0.8
 
@@ -303,7 +329,8 @@ class BenchmarkRequest(BaseModel):
 @router.post("/benchmark")
 async def run_benchmark(req: BenchmarkRequest):
     async def event_generator() -> AsyncIterator[dict]:
-        task = get_task(req.task_id)
+        task_id = _DIFFICULTY_MAP.get(req.task_id, req.task_id)
+        task = get_task(task_id)
         scores: list[float] = []
 
         for question_obj in task.questions:
@@ -315,7 +342,7 @@ async def run_benchmark(req: BenchmarkRequest):
 
             # Run the question through the env
             env = SQLAgentEnv()
-            obs = env.reset_with_question(req.task_id, question_obj.id)
+            obs = env.reset_with_question(task_id, question_obj.id)
 
             attempt = 0
             sql = ""
@@ -373,7 +400,7 @@ async def run_benchmark(req: BenchmarkRequest):
                 rows, error = execute_query(sql)
                 from env.tasks import grade_response
                 task_score = grade_response(
-                    req.task_id, question_obj.id, sql, rows, error, attempt
+                    task_id, question_obj.id, sql, rows, error, attempt
                 )
                 success = task_score >= 0.8
 
@@ -428,7 +455,7 @@ async def run_benchmark(req: BenchmarkRequest):
         yield {"data": json.dumps({
             "type": "done",
             "overall_score": overall_score,
-            "task_id": req.task_id,
+            "task_id": task_id,
         })}
 
     return EventSourceResponse(event_generator())

backend/env/database.py

@@ -21,6 +21,34 @@ from typing import Any
 _DATA_DIR = Path(os.environ.get("DATA_DIR", Path(__file__).parent.parent / "data"))
 DB_PATH = _DATA_DIR / "benchmark.db"
 
+# Active DB path — can be overridden via connect_external_db()
+_active_db_path: str = str(DB_PATH)
+_active_db_label: str = "benchmark (built-in)"
+
+
+def connect_external_db(path: str) -> tuple[bool, str]:
+    """Switch the active SQLite database. Returns (success, message)."""
+    global _active_db_path, _active_db_label
+    try:
+        conn = sqlite3.connect(path)
+        tables = conn.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
+        ).fetchall()
+        conn.close()
+        _active_db_path = path
+        _active_db_label = Path(path).name
+        return True, f"Connected to {Path(path).name} ({len(tables)} tables)"
+    except Exception as e:
+        return False, str(e)
+
+
+def get_active_db_label() -> str:
+    return _active_db_label
+
+
+def _get_db_path() -> str:
+    return _active_db_path
+
 
 # ─── Schema ───────────────────────────────────────────────────────
 
@@ -350,10 +378,14 @@ def get_schema_info() -> str:
     """
     Return a concise textual schema summary for use in prompts.
     """
-    conn = sqlite3.connect(str(DB_PATH))
+    conn = sqlite3.connect(_get_db_path())
     try:
+        cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
+        tables_in_db = [r[0] for r in cur.fetchall()]
+        # Fall back to all tables if schema is unknown
+        tables = tables_in_db if tables_in_db else ["sellers", "users", "products", "orders", "reviews"]
         lines = []
-        for table in ["sellers", "users", "products", "orders", "reviews"]:
+        for table in tables:
             info = conn.execute(f"PRAGMA table_info({table})").fetchall()
             cols = ", ".join(
                 f"{col[1]} {col[2]}{'(PK)' if col[5] else ''}"
@@ -371,7 +403,7 @@ def execute_query(sql: str) -> tuple[list[dict], str | None]:
     Execute a SQL query and return (rows, error_message).
     rows is a list of dicts; error_message is None on success.
     """
-    conn = sqlite3.connect(str(DB_PATH))
+    conn = sqlite3.connect(_get_db_path())
     conn.row_factory = sqlite3.Row
     try:
         cursor = conn.execute(sql)
@@ -385,9 +417,10 @@ def execute_query(sql: str) -> tuple[list[dict], str | None]:
 
 def get_table_stats() -> list[dict]:
     """Return [{name, rows}, ...] for all tables."""
-    conn = sqlite3.connect(str(DB_PATH))
+    conn = sqlite3.connect(_get_db_path())
     try:
-        tables = ["sellers", "users", "products", "orders", "reviews"]
+        cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
+        tables = [r[0] for r in cur.fetchall()] or ["sellers", "users", "products", "orders", "reviews"]
         return [
             {
                 "name": t,

backend/gepa/optimizer.py

@@ -59,7 +59,7 @@ class Candidate(BaseModel):
 def _make_client() -> AsyncOpenAI:
     return AsyncOpenAI(
         api_key=os.environ.get("HF_TOKEN", ""),
-        base_url=os.environ.get("API_BASE_URL", "https://api.openai.com/v1"),
+        base_url=os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1"),
     )
 
 
@@ -158,6 +158,12 @@ class GEPAOptimizer:
             )
         self._save()
 
+    @property
+    def current_generation(self) -> int:
+        if not self._pareto_front:
+            return 0
+        return max(c.generation for c in self._pareto_front)
+
     def should_optimize(self) -> bool:
         return len(self._history) > 0 and len(self._history) % 4 == 0
 

frontend/src/App.tsx

@@ -9,6 +9,7 @@ import { BenchmarkPanel } from './components/BenchmarkPanel'
 import { ERDiagram } from './components/ERDiagram'
 import { RightSidebar } from './components/RightSidebar'
 import { DemoMode } from './components/DemoMode'
+import { ConnectDB } from './components/ConnectDB'
 import { useStore } from './store/useStore'
 import { fetchInit } from './lib/api'
 
@@ -25,8 +26,9 @@ export default function App() {
   const [leftOpen, setLeftOpen] = useState(false)
   const [rightOpen, setRightOpen] = useState(false)
   const [demoOpen, setDemoOpen] = useState(false)
+  const [connectDbOpen, setConnectDbOpen] = useState(false)
 
-  const { theme, setDbSeeded, setTables, setSchemaGraph } = useStore()
+  const { theme, setDbSeeded, setTables, setSchemaGraph, setDbLabel } = useStore()
 
   // Apply theme on mount / change
   useEffect(() => {
@@ -50,6 +52,7 @@ export default function App() {
       .then((d) => {
         setDbSeeded(true)
         setTables(d.tables)
+        if (d.dbLabel) setDbLabel(d.dbLabel)
         // Lazy-load schema graph
         fetch('/api/schema-graph')
           .then((r) => r.json())
@@ -74,11 +77,15 @@ export default function App() {
         onToggleLeft={() => { setLeftOpen((v) => !v); setRightOpen(false) }}
         onToggleRight={() => { setRightOpen((v) => !v); setLeftOpen(false) }}
         onDemo={() => setDemoOpen(true)}
+        onConnectDb={() => setConnectDbOpen(true)}
       />
 
       <AnimatePresence>
         {demoOpen && <DemoMode onClose={() => setDemoOpen(false)} />}
       </AnimatePresence>
+      <AnimatePresence>
+        {connectDbOpen && <ConnectDB onClose={() => setConnectDbOpen(false)} />}
+      </AnimatePresence>
 
       <div className="flex flex-1 overflow-hidden relative">
         {/* Overlay backdrop (mobile) */}

frontend/src/components/BenchmarkPanel.tsx

@@ -251,6 +251,23 @@ export function BenchmarkPanel() {
     <div className="flex flex-col h-full">
       {/* Header */}
       <div className="px-4 py-3 border-b border-white/[0.06] shrink-0">
+        {/* Difficulty tabs */}
+        <div className="flex items-center gap-1 mb-2.5 p-0.5 rounded-lg border border-white/[0.06] w-fit">
+          {DIFFICULTY_TABS.map((tab) => (
+            <button
+              key={tab.id}
+              onClick={() => setTaskDifficulty(tab.id)}
+              disabled={isBenchmarking}
+              className={`text-[10px] font-semibold px-2.5 py-1 rounded transition-all disabled:opacity-50 ${
+                taskDifficulty === tab.id
+                  ? 'bg-violet-600/25 text-violet-300 border border-violet-500/30'
+                  : 'text-gray-500 hover:text-gray-300 border border-transparent'
+              }`}
+            >
+              {tab.label}
+            </button>
+          ))}
+        </div>
         <div className="flex items-center justify-between mb-2">
           <div className="flex items-center gap-2">
             <Target size={14} className="text-violet-400" />

frontend/src/components/ConnectDB.tsx

@@ -0,0 +1,190 @@
+import { useState } from 'react'
+import { motion } from 'framer-motion'
+import { X, PlugZap, Database, CheckCircle2, XCircle, Loader2, RotateCcw } from 'lucide-react'
+import { useStore } from '../store/useStore'
+import { connectExternalDb } from '../lib/api'
+
+interface ConnectDBProps {
+  onClose: () => void
+}
+
+const EXAMPLES = [
+  { label: 'In-memory (blank)', value: ':memory:' },
+  { label: 'Custom path', value: '/path/to/your/database.db' },
+]
+
+export function ConnectDB({ onClose }: ConnectDBProps) {
+  const { dbLabel, setDbLabel, setTables, setDbSeeded } = useStore()
+  const [path, setPath] = useState('')
+  const [status, setStatus] = useState<'idle' | 'connecting' | 'success' | 'error'>('idle')
+  const [message, setMessage] = useState('')
+
+  const handleConnect = async () => {
+    if (!path.trim()) return
+    setStatus('connecting')
+    setMessage('')
+    try {
+      const res = await connectExternalDb(path.trim())
+      if (res.success) {
+        setDbLabel(res.dbLabel)
+        setTables(res.tables)
+        setDbSeeded(true)
+        setStatus('success')
+        setMessage(res.message)
+      } else {
+        setStatus('error')
+        setMessage(res.message)
+      }
+    } catch (e) {
+      setStatus('error')
+      setMessage(e instanceof Error ? e.message : 'Connection failed')
+    }
+  }
+
+  const handleReset = async () => {
+    setStatus('connecting')
+    try {
+      const res = await connectExternalDb('/app/backend/data/benchmark.db')
+      if (res.success) {
+        setDbLabel(res.dbLabel)
+        setTables(res.tables)
+        setDbSeeded(true)
+        setStatus('success')
+        setMessage('Reset to built-in benchmark database')
+      } else {
+        setStatus('error')
+        setMessage(res.message)
+      }
+    } catch (e) {
+      setStatus('error')
+      setMessage(e instanceof Error ? e.message : 'Reset failed')
+    }
+  }
+
+  return (
+    <motion.div
+      initial={{ opacity: 0 }}
+      animate={{ opacity: 1 }}
+      exit={{ opacity: 0 }}
+      className="fixed inset-0 z-[200] flex items-center justify-center p-4"
+      style={{ background: 'rgba(0,0,0,0.6)', backdropFilter: 'blur(4px)' }}
+      onClick={(e) => { if (e.target === e.currentTarget) onClose() }}
+    >
+      <motion.div
+        initial={{ scale: 0.95, opacity: 0, y: 8 }}
+        animate={{ scale: 1, opacity: 1, y: 0 }}
+        exit={{ scale: 0.95, opacity: 0 }}
+        transition={{ duration: 0.15 }}
+        className="w-full max-w-md rounded-2xl border shadow-2xl overflow-hidden"
+        style={{ background: 'var(--bg-secondary)', borderColor: 'var(--border-color)' }}
+      >
+        {/* Header */}
+        <div className="flex items-center justify-between px-5 py-4 border-b" style={{ borderColor: 'var(--border-color)' }}>
+          <div className="flex items-center gap-2.5">
+            <div className="w-7 h-7 rounded-lg flex items-center justify-center" style={{ background: 'linear-gradient(135deg,#1e3a5f,#2d1b69)' }}>
+              <PlugZap size={13} className="text-white" />
+            </div>
+            <div>
+              <h2 className="text-sm font-semibold theme-text-primary">Connect Database</h2>
+              <p className="text-[10px] text-gray-500">SQLite file path</p>
+            </div>
+          </div>
+          <button onClick={onClose} className="p-1.5 rounded-lg hover:bg-white/5 text-gray-500 hover:text-gray-300 transition-colors">
+            <X size={15} />
+          </button>
+        </div>
+
+        {/* Current DB */}
+        <div className="px-5 py-3 border-b flex items-center gap-2" style={{ borderColor: 'var(--border-color)', background: 'var(--bg-tertiary)' }}>
+          <Database size={11} className="text-violet-400 shrink-0" />
+          <span className="text-[11px] text-gray-500">Active:</span>
+          <span className="text-[11px] font-semibold text-violet-400 truncate">{dbLabel}</span>
+          <button
+            onClick={handleReset}
+            className="ml-auto flex items-center gap-1 text-[10px] text-gray-600 hover:text-gray-400 transition-colors"
+            title="Reset to built-in demo database"
+          >
+            <RotateCcw size={9} />
+            Reset to demo
+          </button>
+        </div>
+
+        {/* Body */}
+        <div className="px-5 py-4 flex flex-col gap-4">
+          <div>
+            <label className="text-[10px] font-semibold text-gray-500 uppercase tracking-wider block mb-1.5">
+              SQLite File Path
+            </label>
+            <input
+              type="text"
+              value={path}
+              onChange={(e) => { setPath(e.target.value); setStatus('idle') }}
+              onKeyDown={(e) => e.key === 'Enter' && void handleConnect()}
+              placeholder="/path/to/database.db"
+              className="w-full px-3 py-2.5 text-sm rounded-xl border focus:outline-none transition-all font-mono"
+              style={{
+                background: 'var(--bg-tertiary)',
+                borderColor: 'var(--border-color)',
+                color: 'var(--text-primary)',
+              }}
+              autoFocus
+            />
+          </div>
+
+          {/* Quick examples */}
+          <div className="flex flex-col gap-1.5">
+            <span className="text-[10px] text-gray-600 uppercase tracking-wider">Quick select</span>
+            <div className="flex flex-wrap gap-1.5">
+              {EXAMPLES.map((ex) => (
+                <button
+                  key={ex.value}
+                  onClick={() => { setPath(ex.value); setStatus('idle') }}
+                  className="text-[10px] px-2.5 py-1 rounded-full border transition-all text-gray-500 hover:text-gray-300"
+                  style={{ borderColor: 'var(--border-color)', background: 'var(--bg-tertiary)' }}
+                >
+                  {ex.label}
+                </button>
+              ))}
+            </div>
+          </div>
+
+          {/* Status message */}
+          {status !== 'idle' && (
+            <div className={`flex items-start gap-2 rounded-xl px-3 py-2.5 text-xs ${
+              status === 'success' ? 'bg-green-500/10 border border-green-500/20 text-green-400' :
+              status === 'error' ? 'bg-red-500/10 border border-red-500/20 text-red-400' :
+              'bg-violet-500/10 border border-violet-500/20 text-violet-400'
+            }`}>
+              {status === 'connecting' && <Loader2 size={12} className="animate-spin shrink-0 mt-0.5" />}
+              {status === 'success' && <CheckCircle2 size={12} className="shrink-0 mt-0.5" />}
+              {status === 'error' && <XCircle size={12} className="shrink-0 mt-0.5" />}
+              <span>{status === 'connecting' ? 'Connecting…' : message}</span>
+            </div>
+          )}
+        </div>
+
+        {/* Footer */}
+        <div className="px-5 pb-5 flex items-center justify-end gap-2">
+          <button
+            onClick={onClose}
+            className="px-4 py-2 rounded-xl text-xs font-medium text-gray-500 hover:text-gray-300 transition-colors"
+          >
+            {status === 'success' ? 'Close' : 'Cancel'}
+          </button>
+          <button
+            onClick={() => void handleConnect()}
+            disabled={!path.trim() || status === 'connecting'}
+            className="flex items-center gap-1.5 px-4 py-2 rounded-xl text-xs font-semibold text-white transition-all active:scale-95 disabled:opacity-40 disabled:cursor-not-allowed"
+            style={{ background: 'linear-gradient(135deg,#7c3aed,#2563eb)' }}
+          >
+            {status === 'connecting' ? (
+              <><Loader2 size={11} className="animate-spin" /> Connecting…</>
+            ) : (
+              <><PlugZap size={11} /> Connect</>
+            )}
+          </button>
+        </div>
+      </motion.div>
+    </motion.div>
+  )
+}

frontend/src/components/Header.tsx

@@ -1,21 +1,15 @@
-import { Database, Sun, Moon, PanelLeftOpen, PanelRightOpen, Cpu, Play } from 'lucide-react'
+import { Database, Sun, Moon, PanelLeftOpen, PanelRightOpen, Cpu, Play, PlugZap } from 'lucide-react'
 import { useStore } from '../store/useStore'
-import type { Difficulty } from '../lib/types'
 
 interface HeaderProps {
   onToggleLeft: () => void
   onToggleRight: () => void
   onDemo: () => void
+  onConnectDb: () => void
 }
 
-const DIFFICULTIES: { id: Difficulty; label: string; color: string }[] = [
-  { id: 'easy', label: 'Easy', color: 'text-green-400 border-green-500/30 bg-green-500/10' },
-  { id: 'medium', label: 'Medium', color: 'text-amber-400 border-amber-500/30 bg-amber-500/10' },
-  { id: 'hard', label: 'Hard', color: 'text-red-400 border-red-500/30 bg-red-500/10' },
-]
-
-export function Header({ onToggleLeft, onToggleRight, onDemo }: HeaderProps) {
-  const { theme, toggleTheme, dbSeeded, taskDifficulty, setTaskDifficulty } = useStore()
+export function Header({ onToggleLeft, onToggleRight, onDemo, onConnectDb }: HeaderProps) {
+  const { theme, toggleTheme, dbSeeded, dbLabel } = useStore()
 
   return (
     <header
@@ -71,22 +65,16 @@ export function Header({ onToggleLeft, onToggleRight, onDemo }: HeaderProps) {
           LinUCB Active
         </div>
 
-        {/* Difficulty selector */}
-        <div className="flex items-center gap-1 border border-white/[0.06] rounded-lg p-0.5">
-          {DIFFICULTIES.map((d) => (
-            <button
-              key={d.id}
-              onClick={() => setTaskDifficulty(d.id)}
-              className={`text-[10px] font-semibold px-2 py-1 rounded transition-all ${
-                taskDifficulty === d.id
-                  ? `${d.color} border`
-                  : 'text-gray-500 hover:text-gray-300 border border-transparent'
-              }`}
-            >
-              {d.label}
-            </button>
-          ))}
-        </div>
+        {/* Connect DB button */}
+        <button
+          onClick={onConnectDb}
+          className="hidden sm:flex items-center gap-1.5 px-2.5 py-1.5 rounded-lg text-[11px] font-medium transition-all hover:bg-white/5 theme-border border"
+          style={{ color: 'var(--text-muted)' }}
+          title={`Active: ${dbLabel}`}
+        >
+          <PlugZap size={11} />
+          <span className="hidden md:inline">Connect DB</span>
+        </button>
 
         {/* Demo button */}
         <button

frontend/src/lib/api.ts

@@ -95,3 +95,13 @@ export async function fetchPromptHistory() {
   if (!res.ok) throw new Error(`HTTP ${res.status}`)
   return res.json()
 }
+
+export async function connectExternalDb(path: string): Promise<{ success: boolean; message: string; tables: { name: string; rows: number }[]; dbLabel: string }> {
+  const res = await fetch(`${BASE_URL}/api/connect-db`, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ path }),
+  })
+  if (!res.ok) throw new Error(`HTTP ${res.status}`)
+  return res.json()
+}

frontend/src/store/useStore.ts

@@ -19,6 +19,9 @@ interface Store {
   taskDifficulty: Difficulty
   setTaskId: (id: string) => void
   setTaskDifficulty: (d: Difficulty) => void
+  // DB
+  dbLabel: string
+  setDbLabel: (label: string) => void
 
   // Init / DB
   dbSeeded: boolean
@@ -97,17 +100,22 @@ export const useStore = create<Store>((set) => ({
     }),
 
   // Task
-  taskId: 'easy',
+  taskId: 'simple_queries',
   taskDifficulty: 'easy',
   setTaskId: (id) => set({ taskId: id }),
-  setTaskDifficulty: (d) =>
+  setTaskDifficulty: (d) => {
+    const taskId = d === 'easy' ? 'simple_queries' : d === 'medium' ? 'join_queries' : 'complex_queries'
     set({
       taskDifficulty: d,
-      taskId: d,
+      taskId,
       benchmarkResults:
         d === 'easy' ? EASY_QUERIES : d === 'medium' ? MEDIUM_QUERIES : HARD_QUERIES,
       overallScore: null,
-    }),
+    })
+  },
+  // DB
+  dbLabel: 'benchmark (built-in)',
+  setDbLabel: (label) => set({ dbLabel: label }),
 
   // Init
   dbSeeded: false,