feat: admin user list — sort by last login desc, convert to EST

chat_interface.py

@@ -4603,7 +4603,7 @@ def create_chat_interface():
                             with gr.Column(scale=2):
                                 gr.Markdown("#### Current Users")
                                 user_list_display = gr.Dataframe(
-                                    headers=["Email", "Display Name", "Admin", "Active", "Last Login"],
+                                    headers=["Email", "Display Name", "Admin", "Active", "Last Login EST"],
                                     datatype=["str", "str", "bool", "bool", "str"],
                                     interactive=False,
                                     label="Users"
@@ -4633,16 +4633,29 @@ def create_chat_interface():
                             """Load user list from Supabase."""
                             try:
                                 from supabase_client import UserManager
+                                from zoneinfo import ZoneInfo
                                 um = UserManager()
                                 users = um.list_users()
+                                users = sorted(users, key=lambda u: u.get('last_login') or '', reverse=True)
+                                eastern = ZoneInfo('America/New_York')
                                 rows = []
                                 for u in users:
+                                    raw_login = u.get('last_login')
+                                    if raw_login:
+                                        try:
+                                            from datetime import datetime
+                                            dt = datetime.fromisoformat(str(raw_login).replace('Z', '+00:00'))
+                                            last_login_str = dt.astimezone(eastern).strftime('%Y-%m-%d %H:%M')
+                                        except Exception:
+                                            last_login_str = str(raw_login)[:19]
+                                    else:
+                                        last_login_str = 'Never'
                                     rows.append([
                                         u.get('email', ''),
                                         u.get('display_name', ''),
                                         u.get('is_admin', False),
                                         u.get('is_active', True),
-                                        str(u.get('last_login', 'Never'))[:19] if u.get('last_login') else 'Never'
+                                        last_login_str
                                     ])
                                 return rows
                             except Exception as e:

tests/e2e_quality.py

@@ -1,24 +1,19 @@
 """
 Quality regression test suite for DemoPrep.
 
-Runs 6 pipeline tests against the live HF Space:
-  - 2 fixed         (same company/use case every run — regression baseline)
-  - 2 random        (drawn from the pool in quality_config.yaml)
-  - 2 AI-generated  (Claude picks a novel company + use case each run)
+Runs 6 pipeline tests against the live HF Space using the form UI:
+  - 2 fixed         (same every run — regression baselines)
+  - 2 random        (use case picked from pool, AI selects matching company)
+  - 2 AI-generated  (AI picks vertical, line, function, and company)
 
-Each test submits a prompt via the Chat tab, monitors pipeline progress,
-then grades the result:
+Scoring (100 pts total):
   - Stage completion  → up to 25 pts  (research 5, ddl 7, data 8, thoughtspot 5)
   - Data quality      → up to 50 pts  (LLM grades model TML + Snowflake sample, 0-100 scaled)
   - Liveboard quality → up to 25 pts  (LLM grades liveboard TML, 0-100 scaled)
-  Total: 100 pts, A-F grade
 
 Usage:
     source demoprep/bin/activate
     python tests/e2e_quality.py
-
-Or via pytest:
-    pytest tests/e2e_quality.py -v -s
 """
 
 import json
@@ -37,7 +32,6 @@ import yaml
 from dotenv import load_dotenv
 from playwright.sync_api import Page, sync_playwright
 
-# Make project modules importable (snowflake_auth, main_research, etc.)
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
 # ---------------------------------------------------------------------------
@@ -49,8 +43,8 @@ BASE_URL      = "https://thoughtspot-dp-demoprep.hf.space"
 TEST_USER     = os.getenv("TEST_USER")
 TEST_PASSWORD = os.getenv("TEST_PASSWORD")
 
-CONFIG_FILE  = Path(__file__).parent / "quality_config.yaml"
-RESULTS_DIR  = Path(__file__).parent / "quality_results"
+CONFIG_FILE = Path(__file__).parent / "quality_config.yaml"
+RESULTS_DIR = Path(__file__).parent / "quality_results"
 RESULTS_DIR.mkdir(exist_ok=True)
 
 STAGE_LABELS = {
@@ -71,154 +65,187 @@ def load_config() -> dict:
 
 
 # ---------------------------------------------------------------------------
-# Test case generation
+# LLM helper (uses app's configured LLM)
 # ---------------------------------------------------------------------------
 def _get_researcher():
-    """Return a MultiLLMResearcher using the app's configured default LLM."""
     from main_research import MultiLLMResearcher
     from llm_config import DEFAULT_LLM_MODEL, map_llm_display_to_provider
     provider, model = map_llm_display_to_provider(DEFAULT_LLM_MODEL)
     return MultiLLMResearcher(provider=provider, model=model)
 
 
-def generate_ai_test_case(config: dict) -> dict:
-    """Ask the configured LLM to pick a novel company + use case."""
+def _llm(prompt: str, max_tokens: int = 300) -> str:
     researcher = _get_researcher()
-    prompt = config["ai_generated"]["generation_prompt"]
-    raw = researcher.make_request(
+    return (researcher.make_request(
         [{"role": "user", "content": prompt}],
-        max_tokens=200,
+        max_tokens=max_tokens,
         stream=False,
-    )
-    if not raw:
-        raise ValueError("LLM returned empty response for test case generation")
-    match = re.search(r'\{.*\}', raw.strip(), re.DOTALL)
+    ) or "").strip()
+
+
+def _parse_json(text: str) -> dict:
+    match = re.search(r'\{.*\}', text, re.DOTALL)
     if not match:
-        raise ValueError(f"LLM did not return valid JSON: {raw}")
-    data = json.loads(match.group())
+        raise ValueError(f"No JSON found in: {text[:200]}")
+    return json.loads(match.group())
+
+
+# ---------------------------------------------------------------------------
+# Test case generation
+# ---------------------------------------------------------------------------
+def generate_ai_test_case(config: dict) -> dict:
+    """AI picks vertical, line, function, and a matching company."""
+    prompt = config["ai_generated"]["generation_prompt"]
+    data = _parse_json(_llm(prompt))
     return {
-        "name": f"AI: {data['company']} — {data['use_case']}",
-        "type": "ai_generated",
-        "prompt": data["prompt"],
-        "company": data["company"],
-        "use_case": data["use_case"],
+        "name":        f"AI: {data['company']} — {data['vertical']} / {data['line']} / {data['function']}",
+        "type":        "ai_generated",
+        "company":     data["company"],
+        "company_url": data["company_url"],
+        "vertical":    data["vertical"],
+        "line":        data["line"],
+        "function":    data["function"],
     }
 
 
-def pick_random_test_case(config: dict, used_combos: set) -> dict:
-    """
-    Pick a use case at random, then ask the LLM to select a matching company.
-    Avoids use cases already used this run.
-    """
-    pool = config["random_pool"]
-    use_cases = pool["use_cases"]
-    company_prompt_template = pool["company_prompt"]
+def pick_random_test_case(config: dict, used_labels: set) -> dict:
+    """Pick a use case from the pool, ask AI for a matching company."""
+    pool = config["random_pool"]["use_cases"]
+    template = config["random_pool"]["company_prompt"]
+    fixed_companies = [t["company"] for t in config.get("fixed_tests", [])]
 
-    # Pick a use case not yet used this run
-    available = [uc for uc in use_cases if uc not in used_combos]
+    available = [uc for uc in pool if uc["label"] not in used_labels]
     if not available:
-        available = use_cases  # all used — allow repeats
-    use_case = random.choice(available)
-    used_combos.add(use_case)
-
-    # Fixed test companies to exclude so the LLM doesn't repeat them
-    fixed_companies = [tc["company"] for tc in config.get("fixed_tests", [])]
-    exclude_list = ", ".join(fixed_companies)
-
-    prompt_text = company_prompt_template.format(
-        use_case=use_case,
-        exclude_list=exclude_list,
+        available = pool
+    uc = random.choice(available)
+    used_labels.add(uc["label"])
+
+    exclude = ", ".join(fixed_companies)
+    prompt = template.format(
+        label=uc["label"],
+        vertical=uc["vertical"],
+        line=uc["line"],
+        function=uc["function"],
+        exclude_list=exclude,
     )
 
     try:
-        researcher = _get_researcher()
-        raw = researcher.make_request(
-            [{"role": "user", "content": prompt_text}],
-            max_tokens=200,
-            stream=False,
-        ) or ""
-        match = re.search(r'\{.*\}', raw.strip(), re.DOTALL)
-        if not match:
-            raise ValueError(f"No JSON in response: {raw}")
-        data = json.loads(match.group())
-        return {
-            "name": f"Random: {data['company']} — {use_case}",
-            "type": "random",
-            "prompt": data["prompt"],
-            "company": data["company"],
-            "use_case": use_case,
-        }
+        data = _parse_json(_llm(prompt))
+        company     = data["company"]
+        company_url = data["company_url"]
     except Exception as e:
-        # Hard fallback — shouldn't happen often
-        print(f"  ⚠️  Random company selection failed ({e}), using generic fallback")
-        fallbacks = {
-            "Retail Sales": ("Target", "target.com"),
-            "Retail Supply Chain": ("Walmart", "walmart.com"),
-            "Banking Sales": ("JPMorgan Chase", "jpmorganchase.com"),
-            "Banking Supply Chain": ("Bank of America", "bankofamerica.com"),
-            "Software Sales": ("Salesforce", "salesforce.com"),
-            "Software Supply Chain": ("Microsoft", "microsoft.com"),
-        }
-        company, url = fallbacks.get(use_case, ("Amazon", "amazon.com"))
-        return {
-            "name": f"Random: {company} — {use_case}",
-            "type": "random",
-            "prompt": f"{url}, {use_case}",
-            "company": company,
-            "use_case": use_case,
-        }
+        print(f"  ⚠️  Company selection failed ({e}), using fallback")
+        company     = uc["label"].split()[0]
+        company_url = "example.com"
+
+    return {
+        "name":        f"Random: {company} — {uc['label']}",
+        "type":        "random",
+        "company":     company,
+        "company_url": company_url,
+        "vertical":    uc["vertical"],
+        "line":        uc["line"],
+        "function":    uc["function"],
+    }
 
 
 def build_test_suite(config: dict) -> list:
     suite = []
-    used_combos = set()
+    used_labels = set()
 
     for tc in config["fixed_tests"]:
         suite.append({**tc, "type": "fixed"})
 
     for _ in range(2):
-        suite.append(pick_random_test_case(config, used_combos))
+        suite.append(pick_random_test_case(config, used_labels))
 
     ai_count = config["ai_generated"].get("count", 2)
     for i in range(ai_count):
         try:
             suite.append(generate_ai_test_case(config))
         except Exception as e:
-            print(f"  ⚠️  AI test case {i+1} generation failed ({e}), substituting random")
-            suite.append(pick_random_test_case(config, used_combos))
+            print(f"  ⚠️  AI test case {i+1} failed ({e}), substituting random")
+            suite.append(pick_random_test_case(config, used_labels))
 
     random.shuffle(suite)
+    if max_tests:
+        suite = suite[:max_tests]
     return suite
 
 
+# ---------------------------------------------------------------------------
+# Form interaction helpers
+# ---------------------------------------------------------------------------
+def select_gradio_dropdown(page: Page, label: str, value: str):
+    """Select a value from a Gradio dropdown by its label text."""
+    # Find the input inside the dropdown container that follows this label
+    container = page.locator(f'label:has-text("{label}")').locator('xpath=ancestor::div[contains(@class,"form") or contains(@class,"block")][1]')
+    inp = container.locator('input').first
+    inp.click()
+    inp.fill("")
+    inp.type(value, delay=50)
+    page.wait_for_timeout(400)
+    # Click the matching option in the listbox
+    option = page.get_by_role('option', name=value, exact=True)
+    option.wait_for(timeout=5000)
+    option.click()
+    page.wait_for_timeout(300)
+
+
+def submit_job(page: Page, test_case: dict):
+    """Fill the form and click GO."""
+    # Navigate to App tab, then App sub-tab
+    page.goto(BASE_URL, timeout=90000)
+    page.wait_for_selector('button[role=tab]', timeout=60000)
+    page.wait_for_timeout(2000)
+
+    page.click('button[role=tab]:has-text("📱 App")', timeout=10000)
+    page.wait_for_timeout(500)
+    page.get_by_role('tab', name='App', exact=True).click(timeout=10000)
+    page.wait_for_timeout(1000)
+
+    # Select dropdowns
+    select_gradio_dropdown(page, "Vertical",  test_case["vertical"])
+    select_gradio_dropdown(page, "Line",      test_case["line"])
+    select_gradio_dropdown(page, "Function",  test_case["function"])
+
+    # Fill company URL
+    url_input = page.locator('input[placeholder*="company"]').first
+    if not url_input.is_visible():
+        url_input = page.locator('label:has-text("Company URL")').locator('xpath=ancestor::div[1]//input')
+    url_input.click()
+    url_input.fill(test_case["company_url"])
+    page.wait_for_timeout(300)
+
+    # Click GO
+    page.click('button:has-text("→ GO")', timeout=10000)
+    print(f"  ✅ Form submitted: {test_case['vertical']} / {test_case['line']} / {test_case['function']} — {test_case['company_url']}")
+
+
 # ---------------------------------------------------------------------------
 # Pipeline stage detection
 # ---------------------------------------------------------------------------
 def read_progress(page: Page) -> dict:
     """
-    Read the pipeline progress from the page.
-    Returns stage_key -> 'complete' | 'running' | 'failed' | 'not_started' | 'unknown'
+    Read pipeline progress from the right-side progress panel.
+    Returns stage_key -> 'complete' | 'running' | 'not_started' | 'unknown'
     """
+    # Stay on App tab — progress panel is on the right side
     try:
         page.click('button[role=tab]:has-text("📱 App")', timeout=5000)
         page.wait_for_timeout(500)
+        page.get_by_role('tab', name='App', exact=True).click(timeout=3000)
+        page.wait_for_timeout(300)
     except Exception:
         pass
 
-    progress_text = page.evaluate('''() => {
-        const el = document.getElementById("component-39");
-        return el ? el.innerText : "";
-    }''')
-    if not progress_text:
-        progress_text = page.inner_text('body')
+    progress_text = page.inner_text('body')
 
     stages = {}
     for key, label in STAGE_LABELS.items():
-        if f"✅ {label}" in progress_text or f"✓ {label}" in progress_text:
+        if f"✓ {label}" in progress_text or f"✅ {label}" in progress_text:
             stages[key] = "complete"
-        elif f"❌ {label}" in progress_text or f"✗ {label}" in progress_text:
-            stages[key] = "failed"
-        elif f"▶ {label}" in progress_text or f"⏳ {label}" in progress_text:
+        elif f"▶ {label}" in progress_text:
             stages[key] = "running"
         elif f"○ {label}" in progress_text:
             stages[key] = "not_started"
@@ -228,33 +255,26 @@ def read_progress(page: Page) -> dict:
 
 
 def pipeline_finished(stages: dict) -> bool:
-    """True when Complete shows ✓. Timeout handles stuck pipelines."""
     return stages.get("complete") == "complete"
 
 
 # ---------------------------------------------------------------------------
-# Post-run context extraction (GUIDs + TS environment)
+# Post-run GUID extraction
 # ---------------------------------------------------------------------------
 def extract_run_context(page: Page) -> dict:
     """
-    After pipeline completion, scrape the page for model and liveboard URLs.
-    Returns: { ts_base_url, model_guid, liveboard_guid }
+    After completion, find model and liveboard URLs in the page.
+    Checks both the progress panel area and any result display.
     """
     body = page.inner_text('body')
 
-    # Model URL: https://ts.domain/#/data/tables/{guid}
-    model_match = re.search(
-        r'(https://[^\s/#]+)/#/data/tables/([a-f0-9-]{36})', body
-    )
-    # Liveboard URL: https://ts.domain/#/pinboard/{guid}
-    lb_match = re.search(
-        r'(https://[^\s/#]+)/#/pinboard/([a-f0-9-]{36})', body
-    )
+    model_match = re.search(r'(https://[^\s/#]+)/#/data/tables/([a-f0-9-]{36})', body)
+    lb_match    = re.search(r'(https://[^\s/#]+)/#/pinboard/([a-f0-9-]{36})',    body)
 
     return {
         "ts_base_url":    model_match.group(1) if model_match else None,
         "model_guid":     model_match.group(2) if model_match else None,
-        "liveboard_guid": lb_match.group(2) if lb_match else None,
+        "liveboard_guid": lb_match.group(2)    if lb_match    else None,
     }
 
 
@@ -262,32 +282,24 @@ def extract_run_context(page: Page) -> dict:
 # ThoughtSpot API helpers
 # ---------------------------------------------------------------------------
 def _find_ts_key_for_url(ts_base_url: str) -> str:
-    """Match a TS base URL to its trusted-auth key from .env TS_ENV_* vars."""
     target = (ts_base_url or "").rstrip("/")
     for i in range(1, 10):
         url = os.getenv(f"TS_ENV_{i}_URL", "").rstrip("/")
         key = os.getenv(f"TS_ENV_{i}_KEY_VAR", "")
         if url and key and url == target:
             return key
-    # Fallback to ENV_1
     return os.getenv("TS_ENV_1_KEY_VAR", "")
 
 
 def ts_authenticate(ts_base_url: str) -> requests.Session:
-    """Authenticate with ThoughtSpot trusted auth, return session with bearer token."""
     secret_key = _find_ts_key_for_url(ts_base_url)
     if not secret_key:
         raise RuntimeError(f"No trusted auth key found for {ts_base_url}")
-
     session = requests.Session()
     session.headers["Accept"] = "application/json"
     resp = session.post(
         f"{ts_base_url}/api/rest/2.0/auth/token/full",
-        json={
-            "username": TEST_USER,
-            "secret_key": secret_key,
-            "validity_time_in_sec": 3600,
-        },
+        json={"username": TEST_USER, "secret_key": secret_key, "validity_time_in_sec": 3600},
         timeout=30,
     )
     resp.raise_for_status()
@@ -298,14 +310,9 @@ def ts_authenticate(ts_base_url: str) -> requests.Session:
 
 
 def export_tml(ts_base_url: str, session: requests.Session, guid: str) -> str:
-    """Export TML for a model or liveboard GUID, return raw YAML string."""
     resp = session.post(
         f"{ts_base_url}/api/rest/2.0/metadata/tml/export",
-        json={
-            "metadata": [{"identifier": guid}],
-            "export_associated": False,
-            "export_fqn": True,
-        },
+        json={"metadata": [{"identifier": guid}], "export_associated": False, "export_fqn": True},
         timeout=30,
     )
     resp.raise_for_status()
@@ -314,7 +321,6 @@ def export_tml(ts_base_url: str, session: requests.Session, guid: str) -> str:
 
 
 def extract_db_schema(model_tml_str: str) -> tuple:
-    """Pull database + schema from model TML (needed for Snowflake query)."""
     try:
         tml = yaml.safe_load(model_tml_str)
         tables = tml.get("model", {}).get("tables", [])
@@ -327,21 +333,14 @@ def extract_db_schema(model_tml_str: str) -> tuple:
 
 
 def get_snowflake_sample(db: str, schema: str) -> str:
-    """
-    Connect to Snowflake and pull sample rows from each table in the schema.
-    Returns a formatted string suitable for pasting into the grader prompt.
-    """
     try:
         from snowflake_auth import get_snowflake_connection
-        conn = get_snowflake_connection()
+        conn   = get_snowflake_connection()
         cursor = conn.cursor()
-
         cursor.execute(f'SHOW TABLES IN SCHEMA "{db}"."{schema}"')
-        # Column index 1 = table name in SHOW TABLES output
         tables = [row[1] for row in cursor.fetchall()]
-
-        parts = []
-        for table in tables[:6]:  # cap at 6 tables
+        parts  = []
+        for table in tables[:6]:
             try:
                 cursor.execute(f'SELECT * FROM "{db}"."{schema}"."{table}" LIMIT 30')
                 cols = [d[0] for d in cursor.description]
@@ -351,250 +350,167 @@ def get_snowflake_sample(db: str, schema: str) -> str:
                 for row in rows[:15]:
                     parts.append("  " + str(dict(zip(cols, row))))
             except Exception as e:
-                parts.append(f"\nTable: {table}  — error fetching sample: {e}")
-
+                parts.append(f"\nTable: {table} — error: {e}")
         cursor.close()
         conn.close()
         return "\n".join(parts) if parts else "No tables found"
-
     except Exception as e:
         return f"Snowflake connection failed: {e}"
 
 
 # ---------------------------------------------------------------------------
-# LLM grading calls
+# AI quality grading
 # ---------------------------------------------------------------------------
 def _call_grader(prompt: str) -> dict:
-    """Send a grading prompt to the app's configured LLM. Returns parsed JSON dict."""
-    researcher = _get_researcher()
-    raw = researcher.make_request(
-        [{"role": "user", "content": prompt}],
-        max_tokens=1000,
-        stream=False,
-    ) or ""
-    raw = raw.strip()
+    raw = _llm(prompt, max_tokens=1000)
     match = re.search(r'\{.*\}', raw, re.DOTALL)
     if match:
         try:
             return json.loads(match.group())
         except json.JSONDecodeError:
             pass
-    return {
-        "score": 0,
-        "reasoning": "Could not parse grader response",
-        "strengths": [],
-        "weaknesses": [raw[:300]],
-    }
+    return {"score": 0, "reasoning": "Could not parse response", "strengths": [], "weaknesses": [raw[:200]]}
 
 
-def grade_data_quality(company: str, use_case: str, model_tml: str, sample_data: str) -> dict:
-    """
-    Grade the data model and Snowflake sample on a 0-100 scale.
-    Rubric: realism(20), story potential(30), time coverage(20),
-            schema fitness(15), completeness(15).
-    """
-    prompt = f"""You are grading a ThoughtSpot demo dataset for {company} ({use_case}).
+def grade_data_quality(company: str, vertical: str, line: str, function: str,
+                       model_tml: str, sample_data: str) -> dict:
+    prompt = f"""You are grading a ThoughtSpot demo dataset.
 
-The purpose of this demo is to tell a compelling business story for a persona —
-a line of business leader or data analyst — using realistic data with outliers
-and trends that drive a narrative conversation.
+Company: {company}
+Vertical: {vertical} / {line}
+Analytics function: {function}
 
-MODEL TML (schema and structure):
+The goal is a compelling demo with realistic data, outliers that drive a narrative,
+and a schema that supports the key KPIs for this use case.
+
+MODEL TML:
 {model_tml[:3000]}
 
 SNOWFLAKE SAMPLE DATA:
 {sample_data[:3000]}
 
-Grade this dataset 0–100 using this rubric:
-
-1. REALISM (20 pts): Do values look like real {company} data? Are measures in believable
-   ranges for this company's scale? Do dimension members reflect real-world entities
-   (actual product categories, realistic region names, plausible customer segments)?
-
-2. STORY POTENTIAL (30 pts): Is there a "so what"? Are there outliers — products,
-   regions, or time periods that clearly stand out and would anchor a demo conversation?
-   Do trends exist with directionality, or is the data flat/random noise?
-
-3. TIME COVERAGE (20 pts): Does the data span 12–24 months of history? Is there enough
-   range for meaningful trend analysis? Are seasonal patterns present where expected for
-   this industry?
-
-4. SCHEMA FITNESS (15 pts): Does the star schema match the {use_case} use case?
-   Are the right KPIs computable from this schema (revenue, units, margins, etc.)?
-   Are dimension tables appropriate for the vertical?
+Grade 0–100:
+1. REALISM (20 pts): Values look like real {company} data at scale.
+2. STORY POTENTIAL (30 pts): Outliers and trends exist that would anchor a demo conversation.
+3. TIME COVERAGE (20 pts): 12–24 months of history with meaningful trends.
+4. SCHEMA FITNESS (15 pts): Star schema supports the right KPIs for {line} {function}.
+5. COMPLETENESS (15 pts): Tables fully populated, dimensions have 20+ members.
 
-5. COMPLETENESS (15 pts): Are key measure columns populated with no suspicious nulls?
-   Do dimension tables have enough distinct members (20+) to support meaningful slicing?
-
-Return ONLY valid JSON — no other text:
+Return ONLY valid JSON:
 {{"score": 0, "reasoning": "...", "strengths": ["..."], "weaknesses": ["..."]}}"""
-
     return _call_grader(prompt)
 
 
-def grade_liveboard_quality(company: str, use_case: str, liveboard_tml: str) -> dict:
-    """
-    Grade the liveboard on a 0-100 scale.
-    Rubric: data coverage(25), trend coherence(20), story structure(25),
-            viz variety(15), use case alignment(15).
-    """
-    prompt = f"""You are grading a ThoughtSpot liveboard for {company} ({use_case}).
-
-A great demo liveboard tells a complete business story: it opens with top-line KPIs,
-shows trends over time with clear directionality, and breaks down performance by
-key dimensions so a presenter can walk a prospect through a compelling narrative.
-
-LIVEBOARD TML (structure, visualization titles, and question text):
-{liveboard_tml[:4000]}
-
-Grade this liveboard 0–100 using this rubric:
-
-1. DATA COVERAGE (25 pts): Do all visualizations appear to have backing data?
-   Are questions written against specific column names (not generic placeholders)?
-   Are there any obviously broken or empty chart patterns?
+def grade_liveboard_quality(company: str, vertical: str, line: str, function: str,
+                            liveboard_tml: str) -> dict:
+    prompt = f"""You are grading a ThoughtSpot liveboard.
 
-2. TREND COHERENCE (20 pts): Are line/trend charts asking questions that would produce
-   coherent time series — not random noise? Do KPI questions include a time grain
-   (weekly, monthly, or quarterly) as required by ThoughtSpot?
+Company: {company}
+Vertical: {vertical} / {line}
+Analytics function: {function}
 
-3. STORY STRUCTURE (25 pts): Does the liveboard flow logically — KPIs at the top,
-   trend charts in the middle, dimensional breakdowns to support the narrative?
-   Could a sales rep walk through this board and tell a story without confusion?
+A great liveboard opens with KPIs, shows trends with clear directionality,
+and breaks down performance by dimensions — telling a story a presenter can walk through.
 
-4. VISUALIZATION VARIETY (15 pts): Is there a healthy mix of KPIs, line/trend charts,
-   bar/categorical charts? Or is it monotonously the same chart type repeated?
+LIVEBOARD TML:
+{liveboard_tml[:4000]}
 
-5. USE CASE ALIGNMENT (15 pts): Do visualization titles and question text clearly match
-   the {use_case} use case at {company}? Are the most important KPIs for this
-   vertical present and prominent?
+Grade 0–100:
+1. DATA COVERAGE (25 pts): All vizzes have backing data, questions use real column names.
+2. TREND COHERENCE (20 pts): Line charts produce coherent time series; KPIs have time grains.
+3. STORY STRUCTURE (25 pts): Flows KPIs → trends → breakdowns; walkable in a demo.
+4. VISUALIZATION VARIETY (15 pts): Mix of KPIs, line charts, bar charts.
+5. USE CASE ALIGNMENT (15 pts): Titles and questions match {line} {function} at {company}.
 
-Return ONLY valid JSON — no other text:
+Return ONLY valid JSON:
 {{"score": 0, "reasoning": "...", "strengths": ["..."], "weaknesses": ["..."]}}"""
-
     return _call_grader(prompt)
 
 
-# ---------------------------------------------------------------------------
-# AI quality grading orchestration
-# ---------------------------------------------------------------------------
-def run_ai_grading(run_context: dict, company: str, use_case: str) -> dict:
-    """
-    Fetch TML from ThoughtSpot, sample Snowflake, call LLM graders.
-    Returns a dict with scores, points, reasoning, and any errors.
-    """
+def run_ai_grading(run_context: dict, company: str, vertical: str, line: str, function: str) -> dict:
     result = {
-        "data_score":          None,
-        "data_points":         0.0,
-        "data_reasoning":      "Not graded",
-        "data_strengths":      [],
-        "data_weaknesses":     [],
-        "liveboard_score":     None,
-        "liveboard_points":    0.0,
-        "liveboard_reasoning": "Not graded",
-        "liveboard_strengths": [],
-        "liveboard_weaknesses": [],
-        "grading_errors":      [],
+        "data_score": None, "data_points": 0.0,
+        "data_reasoning": "Not graded", "data_strengths": [], "data_weaknesses": [],
+        "liveboard_score": None, "liveboard_points": 0.0,
+        "liveboard_reasoning": "Not graded", "liveboard_strengths": [], "liveboard_weaknesses": [],
+        "grading_errors": [],
     }
 
-    ts_base     = run_context.get("ts_base_url")
-    model_guid  = run_context.get("model_guid")
-    lb_guid     = run_context.get("liveboard_guid")
+    ts_base    = run_context.get("ts_base_url")
+    model_guid = run_context.get("model_guid")
+    lb_guid    = run_context.get("liveboard_guid")
 
     if not ts_base or not model_guid:
-        result["grading_errors"].append(
-            "No model URL found in page output — pipeline may not have completed"
-        )
+        result["grading_errors"].append("No model URL found — pipeline may not have completed")
         return result
 
-    # Authenticate once, reuse session for both exports
     try:
         session = ts_authenticate(ts_base)
     except Exception as e:
         result["grading_errors"].append(f"ThoughtSpot auth failed: {e}")
         return result
 
-    # --- Data quality ---
+    # Data quality
     try:
         print("    🔍 Exporting model TML...")
         model_tml = export_tml(ts_base, session, model_guid)
         db, schema = extract_db_schema(model_tml)
-
-        if db and schema:
-            print(f"    🗄️  Sampling Snowflake {db}.{schema}...")
-            sample = get_snowflake_sample(db, schema)
-        else:
-            sample = "Could not determine database/schema from model TML"
-            result["grading_errors"].append("db/schema not found in model TML")
-
+        sample = get_snowflake_sample(db, schema) if db and schema else "Could not determine db/schema"
         print("    🤖 Grading data quality...")
-        dg = grade_data_quality(company, use_case, model_tml, sample)
+        dg    = grade_data_quality(company, vertical, line, function, model_tml, sample)
         score = max(0, min(100, int(dg.get("score", 0))))
-        result["data_score"]      = score
-        result["data_points"]     = round(score * 0.50, 1)
-        result["data_reasoning"]  = dg.get("reasoning", "")
-        result["data_strengths"]  = dg.get("strengths", [])
-        result["data_weaknesses"] = dg.get("weaknesses", [])
-        print(f"    📊 Data score: {score}/100 → {result['data_points']} pts")
-
+        result.update({
+            "data_score": score, "data_points": round(score * 0.50, 1),
+            "data_reasoning": dg.get("reasoning", ""),
+            "data_strengths": dg.get("strengths", []),
+            "data_weaknesses": dg.get("weaknesses", []),
+        })
+        print(f"    📊 Data: {score}/100 → {result['data_points']} pts")
     except Exception as e:
         result["grading_errors"].append(f"Data grading failed: {e}")
 
-    # --- Liveboard quality ---
+    # Liveboard quality
     if lb_guid:
         try:
             print("    🔍 Exporting liveboard TML...")
             lb_tml = export_tml(ts_base, session, lb_guid)
             print("    🤖 Grading liveboard quality...")
-            lg = grade_liveboard_quality(company, use_case, lb_tml)
+            lg      = grade_liveboard_quality(company, vertical, line, function, lb_tml)
             lb_score = max(0, min(100, int(lg.get("score", 0))))
-            result["liveboard_score"]      = lb_score
-            result["liveboard_points"]     = round(lb_score * 0.25, 1)
-            result["liveboard_reasoning"]  = lg.get("reasoning", "")
-            result["liveboard_strengths"]  = lg.get("strengths", [])
-            result["liveboard_weaknesses"] = lg.get("weaknesses", [])
-            print(f"    📊 Liveboard score: {lb_score}/100 → {result['liveboard_points']} pts")
+            result.update({
+                "liveboard_score": lb_score, "liveboard_points": round(lb_score * 0.25, 1),
+                "liveboard_reasoning": lg.get("reasoning", ""),
+                "liveboard_strengths": lg.get("strengths", []),
+                "liveboard_weaknesses": lg.get("weaknesses", []),
+            })
+            print(f"    📊 Liveboard: {lb_score}/100 → {result['liveboard_points']} pts")
         except Exception as e:
             result["grading_errors"].append(f"Liveboard grading failed: {e}")
     else:
-        result["grading_errors"].append(
-            "No liveboard GUID found — liveboard may not have been created"
-        )
+        result["grading_errors"].append("No liveboard GUID — liveboard may not have been created")
 
     return result
 
 
 # ---------------------------------------------------------------------------
-# Stage-based grading
+# Stage grading
 # ---------------------------------------------------------------------------
 def grade_stages(stages: dict, config: dict) -> dict:
-    """
-    Score stage completions (up to 25 pts).
-    Returns: { stage_scores, stage_total, breakdown }
-    """
     weights   = config["grading"]["stages"]
     breakdown = {}
     total     = 0
-
     for key, weight in weights.items():
         status = stages.get(key, "unknown")
-        if status == "complete":
-            earned = weight
-        elif status == "running":
-            earned = weight // 2
-        else:
-            earned = 0
+        earned = weight if status == "complete" else (weight // 2 if status == "running" else 0)
         breakdown[key] = {"weight": weight, "earned": earned, "status": status}
         total += earned
-
     return {"stage_total": total, "breakdown": breakdown}
 
 
-def compute_grade(total_score: float, config: dict) -> str:
-    pct = total_score  # already out of 100
-    thresholds = config["grading"]["thresholds"]
+def compute_grade(score: float, config: dict) -> str:
     grade = "F"
-    for letter, threshold in sorted(thresholds.items(), key=lambda x: -x[1]):
-        if pct >= threshold:
+    for letter, threshold in sorted(config["grading"]["thresholds"].items(), key=lambda x: -x[1]):
+        if score >= threshold:
             grade = letter
             break
     return grade
@@ -604,47 +520,21 @@ def compute_grade(total_score: float, config: dict) -> str:
 # Single test runner
 # ---------------------------------------------------------------------------
 def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
-    """
-    Submit one prompt, monitor until completion or timeout, then grade.
-    Returns the full result dict.
-    """
-    timeout_sec  = config["grading"]["timeout_minutes"] * 60
-    start        = time.time()
-
-    print(f"\n  📤 Submitting: {test_case['prompt']}")
+    timeout_sec = config["grading"]["timeout_minutes"] * 60
+    start       = time.time()
 
     result = {
-        "name":             test_case["name"],
-        "type":             test_case["type"],
-        "company":          test_case.get("company", ""),
-        "use_case":         test_case.get("use_case", ""),
-        "prompt":           test_case["prompt"],
-        "stages":           {},
-        "run_context":      {},
-        "stage_grading":    {},
-        "ai_grading":       {},
-        "total_score":      0.0,
-        "grade":            "F",
-        "error":            None,
-        "timed_out":        False,
-        "duration_seconds": 0,
+        "name": test_case["name"], "type": test_case["type"],
+        "company": test_case.get("company", ""),
+        "vertical": test_case.get("vertical", ""), "line": test_case.get("line", ""),
+        "function": test_case.get("function", ""), "company_url": test_case.get("company_url", ""),
+        "stages": {}, "run_context": {}, "stage_grading": {}, "ai_grading": {},
+        "total_score": 0.0, "grade": "F",
+        "error": None, "timed_out": False, "duration_seconds": 0,
     }
 
     try:
-        # Navigate to Chat sub-tab
-        page.goto(BASE_URL, timeout=90000)
-        page.wait_for_selector('button[role=tab]', timeout=60000)
-        page.wait_for_timeout(2000)
-
-        page.click('button[role=tab]:has-text("Chat")', timeout=10000)
-        page.wait_for_timeout(1000)
-
-        chat_input = page.locator('input[placeholder*="Amazon.com"]')
-        chat_input.wait_for(state='visible', timeout=30000)
-        chat_input.click()
-        chat_input.fill(test_case["prompt"])
-        page.click('button:has-text("Send")', timeout=10000)
-
+        submit_job(page, test_case)
         print(f"  ⏳ Monitoring pipeline (timeout: {config['grading']['timeout_minutes']}min)...")
 
         poll_interval = 15
@@ -655,28 +545,18 @@ def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
             if stages != last_stages:
                 done    = [k for k, v in stages.items() if v == "complete"]
                 running = [k for k, v in stages.items() if v == "running"]
-                failed  = [k for k, v in stages.items() if v == "failed"]
-                print(f"    ✅ {done}  ▶ {running}  ❌ {failed}")
+                print(f"    ✓ {done}  ▶ {running}")
                 last_stages = stages
             if pipeline_finished(stages):
-                print("  ✅ Pipeline finished")
+                print("  ✅ Pipeline complete")
                 break
         else:
             result["timed_out"] = True
-            print(f"  ⏰ Timed out after {config['grading']['timeout_minutes']} minutes")
+            print(f"  ⏰ Timed out after {config['grading']['timeout_minutes']} min")
 
         result["stages"] = last_stages or read_progress(page)
 
-        # Navigate back to Chat tab so the completion message (with URLs) is in the DOM
-        try:
-            page.click('button[role=tab]:has-text("📱 App")', timeout=5000)
-            page.wait_for_timeout(500)
-            page.get_by_role('tab', name='Chat', exact=True).click(timeout=5000)
-            page.wait_for_timeout(1500)
-        except Exception:
-            pass
-
-        # Extract GUIDs from the completion message
+        # Extract GUIDs — check body (progress panel may show URLs)
         run_ctx = extract_run_context(page)
         result["run_context"] = run_ctx
 
@@ -686,51 +566,48 @@ def run_single_test(page: Page, test_case: dict, config: dict) -> dict:
         try:
             result["stages"] = read_progress(page)
         except Exception:
-            result["stages"] = {}
+            pass
 
     result["duration_seconds"] = round(time.time() - start)
 
-    # --- Stage scoring (25 pts max) ---
+    # Stage scoring
     sg = grade_stages(result["stages"], config)
     result["stage_grading"] = sg
 
-    # --- AI quality grading (75 pts max) ---
+    # AI grading
     ag = {"data_points": 0.0, "liveboard_points": 0.0, "grading_errors": []}
     if result["run_context"].get("model_guid"):
         print("  🔬 Running AI quality grading...")
         ag = run_ai_grading(
             result["run_context"],
-            result["company"],
-            result["use_case"],
+            result["company"], result["vertical"], result["line"], result["function"],
         )
     else:
         ag["grading_errors"].append("Skipped — no model GUID (pipeline did not complete)")
     result["ai_grading"] = ag
 
-    # --- Total score (out of 100) ---
     total = sg["stage_total"] + ag.get("data_points", 0) + ag.get("liveboard_points", 0)
     result["total_score"] = round(total, 1)
     result["grade"]       = compute_grade(total, config)
-
     return result
 
 
 # ---------------------------------------------------------------------------
-# Results persistence
+# Results
 # ---------------------------------------------------------------------------
 def save_results(run: dict) -> Path:
     ts   = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
     path = RESULTS_DIR / f"{ts}_quality_run.json"
     with open(path, "w") as f:
         json.dump(run, f, indent=2, default=str)
-    print(f"\n💾 Results saved: {path}")
+    print(f"\n💾 Results: {path}")
     return path
 
 
 # ---------------------------------------------------------------------------
-# Main runner
+# Main
 # ---------------------------------------------------------------------------
-def run_quality_suite():
+def run_quality_suite(max_tests: int = None):
     if not TEST_USER or not TEST_PASSWORD:
         raise RuntimeError("TEST_USER and TEST_PASSWORD must be set in .env")
 
@@ -743,12 +620,11 @@ def run_quality_suite():
           f"{sum(1 for t in suite if t['type']=='fixed')} fixed  "
           f"{sum(1 for t in suite if t['type']=='random')} random  "
           f"{sum(1 for t in suite if t['type']=='ai_generated')} AI-generated")
-    print(f"  Scoring: stages(25) + data quality(50) + liveboard(25) = 100 pts")
+    print(f"  Scoring: stages(25) + data(50) + liveboard(25) = 100 pts")
     print(f"{'='*62}")
-
     for i, tc in enumerate(suite, 1):
-        label = {"fixed": "🔒 Fixed", "random": "🎲 Random", "ai_generated": "🤖 AI"}[tc["type"]]
-        print(f"  [{i}] {label}: {tc['name']}")
+        label = {"fixed": "🔒", "random": "🎲", "ai_generated": "🤖"}[tc["type"]]
+        print(f"  [{i}] {label} {tc['name']}")
 
     run_id  = str(uuid.uuid4())[:8]
     results = []
@@ -776,17 +652,15 @@ def run_quality_suite():
             result = run_single_test(page, test_case, config)
             results.append(result)
 
-            # Mini-report after each test
             ag = result["ai_grading"]
             sg = result["stage_grading"]
-            print(f"  Stage pts:     {sg.get('stage_total', 0)}/25")
+            print(f"  Stages: {sg.get('stage_total', 0)}/25")
             if ag.get("data_score") is not None:
-                print(f"  Data quality:  {ag['data_score']}/100 → {ag['data_points']} pts")
+                print(f"  Data:   {ag['data_score']}/100 → {ag['data_points']} pts")
             if ag.get("liveboard_score") is not None:
-                print(f"  Liveboard:     {ag['liveboard_score']}/100 → {ag['liveboard_points']} pts")
-            if ag.get("grading_errors"):
-                for err in ag["grading_errors"]:
-                    print(f"  ⚠️  {err}")
+                print(f"  Board:  {ag['liveboard_score']}/100 → {ag['liveboard_points']} pts")
+            for err in ag.get("grading_errors", []):
+                print(f"  ⚠️  {err}")
             print(f"  TOTAL: {result['total_score']}/100  Grade: {result['grade']}"
                   f"  ({result['duration_seconds']}s)"
                   f"{'  ⏰ TIMEOUT' if result['timed_out'] else ''}"
@@ -795,51 +669,38 @@ def run_quality_suite():
         ctx.close()
         browser.close()
 
-    # Compile run summary
-    total_pts = sum(r["total_score"] for r in results)
-    avg_score = round(total_pts / len(results), 1) if results else 0
-    overall_grade = compute_grade(avg_score, config)
+    avg   = round(sum(r["total_score"] for r in results) / len(results), 1) if results else 0
+    grade = compute_grade(avg, config)
 
     run = {
-        "run_id":        run_id,
-        "timestamp":     datetime.now().isoformat(),
-        "avg_score":     avg_score,
-        "overall_grade": overall_grade,
-        "test_count":    len(results),
-        "tests":         results,
+        "run_id": run_id, "timestamp": datetime.now().isoformat(),
+        "avg_score": avg, "overall_grade": grade,
+        "test_count": len(results), "tests": results,
     }
-
     save_results(run)
 
-    # Final report
     print(f"\n{'='*62}")
-    print(f"  QUALITY RUN COMPLETE")
-    print(f"{'='*62}")
-    print(f"  Average Score: {avg_score}/100  —  Grade: {overall_grade}")
-    print(f"\n  Per-test results:")
+    print(f"  COMPLETE  —  Avg: {avg}/100  Grade: {grade}")
     for r in results:
         label = {"fixed": "🔒", "random": "🎲", "ai_generated": "🤖"}[r["type"]]
         ag    = r["ai_grading"]
         ds    = f"{ag['data_score']}/100" if ag.get("data_score") is not None else "n/a"
         ls    = f"{ag['liveboard_score']}/100" if ag.get("liveboard_score") is not None else "n/a"
-        print(f"    {label} {r['name']}")
-        print(f"       Score: {r['total_score']}/100  Grade: {r['grade']}"
-              f"  data={ds}  lb={ls}  ({r['duration_seconds']}s)")
+        print(f"  {label} {r['name']}")
+        print(f"     {r['total_score']}/100  Grade:{r['grade']}  data={ds}  lb={ls}  ({r['duration_seconds']}s)")
     print(f"{'='*62}\n")
-
     return run
 
 
-# ---------------------------------------------------------------------------
-# Pytest entry point
-# ---------------------------------------------------------------------------
 def test_quality_run():
-    """Pytest wrapper. Fails if average grade is F."""
     run = run_quality_suite()
-    assert run["overall_grade"] != "F", (
-        f"Quality run averaged {run['avg_score']}% — too many pipeline failures."
-    )
+    assert run["overall_grade"] != "F", f"Quality run averaged {run['avg_score']}% — too many failures."
 
 
 if __name__ == "__main__":
-    run_quality_suite()
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--count", type=int, default=0,
+                        help="Run only N tests (default: all 6)")
+    args = parser.parse_args()
+    run_quality_suite(max_tests=args.count or None)

tests/quality_config.yaml

@@ -1,31 +1,26 @@
 # ============================================================
 # DemoPrep Quality Test Configuration
 # ============================================================
-# All grading weights and test cases live here — change them
-# without touching any Python code.
 
 # ------------------------------------------------------------
 # Grading
-# Each stage is worth points. Scores sum to max_score.
-# Adjust weights to reflect what matters most.
+# Stage completion = 25 pts total.
+# AI quality grades supply the other 75 pts (data 50, liveboard 25).
+# Each AI grade is 0-100, multiplied by weight/100 for points.
 # ------------------------------------------------------------
 grading:
-  # Stage completion points — these sum to 25.
-  # AI quality grades supply the other 75 pts (data 50, liveboard 25).
   stages:
-    research:    5    # Company researched successfully
-    ddl:         7    # Schema generated
-    data:        8    # Data populated in Snowflake
-    thoughtspot: 5    # Model deployed to ThoughtSpot
+    research:    5
+    ddl:         7
+    data:        8
+    thoughtspot: 5
 
-  # AI quality weights — each LLM score (0-100) is multiplied by weight/100
   ai_weights:
-    data_quality:      50   # LLM grades model TML + Snowflake sample
-    liveboard_quality: 25   # LLM grades liveboard TML
+    data_quality:      50
+    liveboard_quality: 25
 
-  max_score: 100      # stages(25) + data_quality(50) + liveboard_quality(25)
+  max_score: 100
 
-  # Letter grade thresholds (%)
   thresholds:
     A: 90
     B: 75
@@ -33,59 +28,75 @@ grading:
     D: 40
     F: 0
 
-  # How long to wait per test before declaring timeout
   timeout_minutes: 45
 
 # ------------------------------------------------------------
-# Fixed tests — same company/use case every run.
-# Use these as regression baselines.
+# Fixed tests — same every run, regression baselines.
+# vertical/line must match values in demo_personas.VERTICAL_LINES.
+# function must match DEMO_FUNCTIONS: Sales, Marketing, Finance, HR, IT, Legal
 # ------------------------------------------------------------
 fixed_tests:
   - name: "Nike — Retail Sales"
-    prompt: "nike.com, Retail Sales"
+    company_url: "nike.com"
     company: "Nike"
-    use_case: "Retail Sales"
+    vertical: "Retail & Consumer Goods"
+    line: "Fashion/Apparel"
+    function: "Sales"
 
   - name: "Wells Fargo — Banking Marketing"
-    prompt: "wellsfargo.com, Banking Marketing"
+    company_url: "wellsfargo.com"
     company: "Wells Fargo"
-    use_case: "Banking Marketing"
+    vertical: "Financial Services"
+    line: "Banking"
+    function: "Marketing"
 
 # ------------------------------------------------------------
-# Random pool — we pick a use case at random, then ask the LLM
-# to select a well-known company whose industry matches it.
+# Random pool — pick a use case combo, AI selects the company.
 # ------------------------------------------------------------
 random_pool:
   use_cases:
-    - "Retail Sales"
-    - "Retail Supply Chain"
-    - "Banking Sales"
-    - "Banking Supply Chain"
-    - "Software Sales"
-    - "Software Supply Chain"
+    - { vertical: "Retail & Consumer Goods",  line: "Department Stores",      function: "Sales",      label: "Retail Department Store Sales" }
+    - { vertical: "Retail & Consumer Goods",  line: "Consumer Electronics",   function: "Marketing",  label: "Consumer Electronics Marketing" }
+    - { vertical: "Retail & Consumer Goods",  line: "Grocery",                function: "Finance",    label: "Grocery Finance" }
+    - { vertical: "Financial Services",       line: "Insurance",              function: "Sales",      label: "Insurance Sales" }
+    - { vertical: "Financial Services",       line: "Asset & Wealth Management", function: "Marketing", label: "Wealth Management Marketing" }
+    - { vertical: "Technology",              line: "Software as a Service",   function: "Sales",      label: "SaaS Sales" }
+    - { vertical: "Technology",              line: "Software as a Service",   function: "Marketing",  label: "SaaS Marketing" }
+    - { vertical: "Transportation & Logistics", line: "Shipping",             function: "Finance",    label: "Shipping Finance" }
+    - { vertical: "Healthcare & Life Sciences", line: "Life Sciences",        function: "Sales",      label: "Life Sciences Sales" }
+    - { vertical: "Manufacturing",           line: "Automotive",              function: "Sales",      label: "Automotive Sales" }
 
-  # Prompt sent to the LLM to pick a matching company.
-  # {use_case} and {exclude_list} are filled in at runtime.
+  # Prompt for AI to pick a matching company for the chosen use case.
   company_prompt: |
-    Pick a well-known company whose primary industry matches the use case: {use_case}.
-    Do NOT pick any company in this list: {exclude_list}.
+    Pick a well-known company that is a strong fit for this use case: {label}.
+    Vertical: {vertical}, Line: {line}, Function: {function}.
+    Do NOT pick: {exclude_list}.
     Return ONLY valid JSON, no other text:
-    {{"company": "...", "company_url": "domain.com", "use_case": "{use_case}", "prompt": "domain.com, {use_case}"}}
+    {{"company": "...", "company_url": "domain.com"}}
 
 # ------------------------------------------------------------
-# AI-generated tests — Claude picks a novel company + use case.
-# Count must match what e2e_quality.py expects (default: 2).
+# AI-generated tests — LLM picks vertical, line, function, AND company.
 # ------------------------------------------------------------
 ai_generated:
   count: 2
-  # Prompt sent to Claude to generate a test case.
-  # Must return valid JSON with keys: company, company_url, use_case, prompt
   generation_prompt: |
-    You are generating test cases for a ThoughtSpot demo builder.
-    Pick a well-known company that is NOT in this list: Nike, Salesforce, Amazon, Target,
-    JPMorgan Chase, Microsoft, Pfizer, Walmart, Delta Airlines, UPS.
-    Also pick an analytics use case from: Retail Sales, Retail Supply Chain,
-    Banking Sales, Banking Supply Chain, Software Sales, Software Supply Chain.
-    Choose a company whose industry matches the use case.
+    You are generating a test case for a ThoughtSpot demo builder.
+
+    Pick a well-known company and a matching analytics use case.
+    The use case must be expressed as a vertical, line, and function from these options:
+
+    Verticals and their lines:
+    - "Retail & Consumer Goods": Fashion/Apparel, Consumer Electronics, Department Stores, Grocery, Specialty Retail
+    - "Financial Services": Banking, Insurance, Asset & Wealth Management
+    - "Technology": Software as a Service, Cybersecurity, Hardware, Artificial Intelligence, IT Services
+    - "Healthcare & Life Sciences": Life Sciences, Healthcare Providers, Healthcare Payers
+    - "Transportation & Logistics": Shipping, Supply Chain Management, Air Transport, Trucking, Freight
+    - "Manufacturing": Automotive, Aerospace, Chemical Production, Electronics Manufacturing
+
+    Functions: Sales, Marketing, Finance, HR, IT, Legal
+
+    Do NOT pick: Nike, Wells Fargo.
+    Choose a company whose industry clearly matches the vertical and line.
+
     Return ONLY valid JSON, no other text:
-    {"company": "...", "company_url": "domain.com", "use_case": "...", "prompt": "domain.com, Use Case"}
+    {{"company": "...", "company_url": "domain.com", "vertical": "...", "line": "...", "function": "..."}}