import json import httpx import traceback import gradio as gr from typing import Any, Dict # Charcoal Black Theme Palette BG_COLOR = "#121212" TEXT_PRIMARY = "#E0E0E0" TEXT_SECONDARY = "#B0B0B0" BORDER_COLOR = "#444444" ACCENT_COLOR = "#888888" custom_theme = gr.themes.Base().set( body_background_fill=BG_COLOR, body_background_fill_dark=BG_COLOR, body_text_color=TEXT_PRIMARY, body_text_color_dark=TEXT_PRIMARY, background_fill_primary="#181818", background_fill_primary_dark="#181818", background_fill_secondary=BG_COLOR, background_fill_secondary_dark=BG_COLOR, border_color_primary=BORDER_COLOR, border_color_primary_dark=BORDER_COLOR, block_background_fill="#181818", block_background_fill_dark="#181818", button_primary_background_fill=ACCENT_COLOR, button_primary_background_fill_dark=ACCENT_COLOR, button_primary_background_fill_hover="#aaaaaa", button_primary_background_fill_hover_dark="#aaaaaa", button_primary_text_color=BG_COLOR, button_primary_text_color_dark=BG_COLOR, button_primary_text_color_hover=BG_COLOR, button_primary_text_color_hover_dark=BG_COLOR, slider_color=ACCENT_COLOR, slider_color_dark=ACCENT_COLOR, ) custom_css = f""" .gradio-container {{ font-family: 'Inter', sans-serif; }} .stat-box {{ background-color: #181818; border: 1px solid {BORDER_COLOR}; border-radius: 8px; padding: 15px; margin-bottom: 10px; text-align: center; }} .stat-value {{ font-size: 24px; font-weight: bold; color: {TEXT_PRIMARY}; }} .stat-label {{ font-size: 12px; text-transform: uppercase; color: {TEXT_SECONDARY}; letter-spacing: 1px; }} .header-title {{ color: {TEXT_PRIMARY} !important; text-align: center; border-bottom: 2px solid {BORDER_COLOR}; padding-bottom: 10px; margin-bottom: 20px; }} textarea, input, select {{ background-color: {BG_COLOR} !important; color: {TEXT_PRIMARY} !important; border: 1px solid {BORDER_COLOR} !important; }} /* Tier Cards */ .tier-card {{ justify-content: flex-start !important; text-align: left !important; padding: 15px !important; border-radius: 8px !important; border: 2px solid #333 !important; background-color: #181818 !important; white-space: pre-wrap !important; line-height: 1.4 !important; transition: all 0.2s !important; }} .tier-card:hover {{ border-color: #666 !important; background-color: #222 !important; }} .tier-selected-image {{ border-color: #22c55e !important; background-color: rgba(34, 197, 94, 0.05) !important; }} .tier-selected-clip {{ border-color: #f59e0b !important; background-color: rgba(245, 158, 11, 0.05) !important; }} .tier-selected-weight {{ border-color: #ef4444 !important; background-color: rgba(239, 68, 68, 0.05) !important; }} /* Subenv Coverage Badges */ .coverage-row {{ display: flex; gap: 10px; margin-bottom: 15px; font-family: 'Inter', sans-serif; }} .cov-badge {{ padding: 4px 12px; border-radius: 16px; font-size: 0.85em; font-weight: bold; cursor: help; }} .cov-active-image {{ background-color: #22c55e; color: #fff; border: 1px solid #22c55e; }} .cov-active-clip {{ background-color: #f59e0b; color: #fff; border: 1px solid #f59e0b; }} .cov-active-weight {{ background-color: #ef4444; color: #fff; border: 1px solid #ef4444; }} .cov-inactive {{ background-color: transparent; color: #666; border: 1px solid #444; cursor: default; }} """ def format_observation(obs: Dict[str, Any]) -> str: """Formats the current observation dict to a readable markdown format.""" if not obs: return "Not initialized. Please click **Initialize Scenario**." node = obs.get("node", "Unknown") step = obs.get("step", 0) instr = obs.get("instruction", "No instructions") signals = obs.get("signals", {}) scores = obs.get("scores", None) md = f"### Step {step}: {node}\n\n**Instructions:** {instr}\n\n" if signals: md += "#### 📡 Available Signals\n```json\n" + json.dumps(signals, indent=2) + "\n```\n" if scores: md += "#### 🏆 Final Score Report\n" for k, v in scores.items(): if isinstance(v, float): md += f"- **{k}**: {v:.3f}\n" else: md += f"- **{k}**: {v}\n" return md def generate_badges(tier: str) -> str: color_class = { "image_audit": "cov-active-image", "clip_audit": "cov-active-clip", "weight_audit": "cov-active-weight" }.get(tier, "cov-active-image") b1_active = tier == "image_audit" b2_active = tier == "clip_audit" b3_active = tier == "weight_audit" def render_badge(name, active, weight): if active: return f'
{name} ✓
' else: return f'
{name} —
' html = '
' html += render_badge("Image Audit", b1_active, "n/a") html += render_badge("Clip Audit", b2_active, "n/a") html += render_badge("Weight Audit", b3_active, "n/a") html += '
' return html async def handle_ingestion(ref_img, clip_vids, lora, tok, prompt, param_json, tier): if tier == "image_audit": if not ref_img: return "Error: reference_image is required for image_audit.", "unknown" if not str(prompt or "").strip(): return "Error: prompt is required for image_audit.", "unknown" if tier == "clip_audit" and not clip_vids: return "Error: clips (1-12 .mp4) are required for clip_audit.", "unknown" if tier == "weight_audit" and not lora: return "Error: lora_weights (.safetensors) are required for weight_audit.", "unknown" try: url = "http://127.0.0.1:8000/ingest-artifacts" files_data = [] if ref_img is not None and tier == "image_audit": files_data.append(("reference_image", open(ref_img.name, "rb"))) if lora is not None and tier == "weight_audit": files_data.append(("lora_weights", open(lora.name, "rb"))) if tok is not None and tier == "weight_audit": files_data.append(("tokenizer_config", open(tok.name, "rb"))) if clip_vids is not None and tier == "clip_audit": for clip in clip_vids: files_data.append(("clips", open(clip.name, "rb"))) data = { "prompt": str(prompt or ""), "param_config_json": str(param_json or "") } async with httpx.AsyncClient() as client: resp = await client.post(url, data=data, files=files_data, timeout=60.0) if resp.status_code == 200: return resp.json().get("ingestion_id", "Error: No ID returned"), tier else: return f"Error {resp.status_code}: {resp.text}", "unknown" except Exception as e: return f"Exception: {str(e)}\n{traceback.format_exc()}", "unknown" def format_tier_label(tier: str) -> str: if tier == "image_audit": return '
[Task: Image Audit ●]
' elif tier == "clip_audit": return '
[Task: Clip Audit ●]
' elif tier == "weight_audit": return '
[Task: Weight Audit ●]
' else: return '
[Task: Unknown ○]
' async def handle_analysis(ingestion_id, model_id, provider, api_key, max_tokens, temp, tier): if not ingestion_id: return "Error: Provide an Ingestion ID first." try: url = "http://127.0.0.1:8000/analyze-ingestion" data = { "ingestion_id": str(ingestion_id), "provider": str(provider), "max_tokens": int(max_tokens), "temperature": float(temp) } if model_id: data["model_id"] = str(model_id) if api_key: data["api_key"] = str(api_key) if tier in ["image_audit", "clip_audit", "weight_audit"]: data["task_tier"] = str(tier) async with httpx.AsyncClient() as client: resp = await client.post(url, json=data, timeout=120.0) if tier not in ["image_audit", "clip_audit", "weight_audit"]: tier = "weight_audit" # Backend default if resp.status_code == 200: report = resp.json().get("report", "Error: No report returned") badges = generate_badges(tier) return badges + "\n\n" + report else: return f"Error {resp.status_code}: {resp.text}" except Exception as e: return f"Exception: {str(e)}" TIER_INFO = { "image_audit": { "title": "Image Audit — Node 1 + Node 2", "body": "The agent receives image quality signals and prompt/config context. It must:\n(1) classify image regime,\n(2) identify image/prompt risks,\n(3) detect parameter anomalies with directional fixes.\nEpisode completes after Node 2 with a standalone Sub-env 1 score." }, "clip_audit": { "title": "Clip Audit — Node 5", "body": "The agent receives clip evidence + dataset context and must recommend\naccept/reject/fix/defer with clear impact reasoning. Episode completes after\nclip disposition with a standalone Sub-env 2 score." }, "weight_audit": { "title": "Weight Audit — Node 8", "body": "The agent receives weight-derived phoneme risk signals and must rank\nrisky phonemes, predict behavior triggers, identify risky clusters, and\nrecommend mitigations. Episode completes after Node 8 with a standalone\nSub-env 3 score." } } def build_custom_ui(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md): with gr.Blocks(theme=custom_theme, css=custom_css, title="TalkingHeadBench Console") as demo: gr.Markdown(f"# {title} 📡", elem_classes=["header-title"]) with gr.Tabs(): # TAB 1: INGESTION & AUTOMATED ANALYSIS with gr.Tab("Workspace Initialization & LLM Analysis"): with gr.Row(): with gr.Column(scale=1): gr.Markdown("### 1. Select Audit Tier & Ingest") task_tier = gr.Textbox(value="image_audit", visible=False) with gr.Row(): btn_image = gr.Button("🟢 IMAGE AUDIT\nReference + Prompt\nNode 1+2", elem_classes=["tier-card", "tier-selected-image"]) btn_clip = gr.Button("🟡 CLIP AUDIT\nDataset Clips Only\nNode 5", elem_classes=["tier-card"]) btn_weight = gr.Button("🔴 WEIGHT AUDIT\nLoRA Weights Only\nNode 8", elem_classes=["tier-card"]) with gr.Accordion(label=TIER_INFO["image_audit"]["title"], open=False) as eval_accordion: eval_markdown = gr.Markdown(value=TIER_INFO["image_audit"]["body"]) ref_file = gr.File(label="Reference Image (.jpg / .png)") clip_files = gr.File(label="Video Clips (.mp4)", file_count="multiple", visible=False) lora_file = gr.File(label="LoRA Weights (.safetensors)", visible=False) tok_file = gr.File(label="Tokenizer Config (.json)", visible=False) prompt_text = gr.Textbox(label="Text Prompt", value="A man speaking directly to the camera.") param_text = gr.Textbox(label="Config JSON (Optional)", value='{"cfg": 7.0}') ingest_btn = gr.Button("Ingest Reference Image", variant="primary") ingestion_id_out = gr.Textbox(label="Generated Ingestion ID", interactive=False) # Tier Selection Logic def set_image_audit(): return [ "image_audit", gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), "Ingest Image Audit Artifacts", gr.update(elem_classes=["tier-card", "tier-selected-image"]), gr.update(elem_classes=["tier-card"]), gr.update(elem_classes=["tier-card"]), gr.update(label=TIER_INFO["image_audit"]["title"], open=True), TIER_INFO["image_audit"]["body"], ] def set_clip_audit(): return [ "clip_audit", gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), "Ingest Clip Audit Artifacts", gr.update(elem_classes=["tier-card"]), gr.update(elem_classes=["tier-card", "tier-selected-clip"]), gr.update(elem_classes=["tier-card"]), gr.update(label=TIER_INFO["clip_audit"]["title"], open=True), TIER_INFO["clip_audit"]["body"], ] def set_weight_audit(): return [ "weight_audit", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), "Ingest Weight Audit Artifacts", gr.update(elem_classes=["tier-card"]), gr.update(elem_classes=["tier-card"]), gr.update(elem_classes=["tier-card", "tier-selected-weight"]), gr.update(label=TIER_INFO["weight_audit"]["title"], open=True), TIER_INFO["weight_audit"]["body"], ] tier_outputs = [ task_tier, ref_file, prompt_text, clip_files, lora_file, tok_file, ingest_btn, btn_image, btn_clip, btn_weight, eval_accordion, eval_markdown, ] btn_image.click(set_image_audit, inputs=[], outputs=tier_outputs) btn_clip.click(set_clip_audit, inputs=[], outputs=tier_outputs) btn_weight.click(set_weight_audit, inputs=[], outputs=tier_outputs) current_task_tier = gr.Textbox(value="unknown", visible=False) with gr.Column(scale=1): gr.Markdown("### 2. Automated LLM Analysis") gr.Markdown( "Use an LLM Diagnostician to automatically analyze the ingested bundle. " "If credentials are configured on the server (HF_TOKEN / API_BASE_URL / MODEL_NAME), " "leave model and API key blank to use those defaults." ) with gr.Row(): analysis_ingestion_id = gr.Textbox(label="Ingestion ID (auto-filled from step 1)", scale=4) active_tier_label = gr.HTML(value='
[Task: Unknown ○]
', scale=1) analysis_provider = gr.Dropdown(choices=["auto", "openai", "anthropic", "huggingface", "local"], value="auto", label="LLM Provider") analysis_model = gr.Textbox(label="Model ID (Optional formatting)", placeholder="e.g. meta-llama/Llama-3.1-70B-Instruct") analysis_api_key = gr.Textbox( label="API Key (Optional override)", type="password", placeholder="Leave blank to use server-side HF_TOKEN", ) analysis_max_tokens = gr.Slider(minimum=64, maximum=4096, value=700, step=1, label="Max Tokens") analysis_temp = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.01, label="Temperature") analyze_btn = gr.Button("Run Diagnostic Analysis", variant="primary") analysis_report_out = gr.Markdown(label="LLM Final Report") # Link step 1 output to step 2 input ingestion_id_out.change(fn=lambda x: x, inputs=ingestion_id_out, outputs=analysis_ingestion_id) def handle_ingestion_id_change(user_typed_id, last_ingested_id, current_state_tier): if user_typed_id != last_ingested_id and user_typed_id: return "unknown" return current_state_tier analysis_ingestion_id.change( fn=handle_ingestion_id_change, inputs=[analysis_ingestion_id, ingestion_id_out, current_task_tier], outputs=[current_task_tier] ) current_task_tier.change( fn=format_tier_label, inputs=[current_task_tier], outputs=[active_tier_label] ) ingest_btn.click( handle_ingestion, inputs=[ref_file, clip_files, lora_file, tok_file, prompt_text, param_text, task_tier], outputs=[ingestion_id_out, current_task_tier] ) analyze_btn.click( handle_analysis, inputs=[analysis_ingestion_id, analysis_model, analysis_provider, analysis_api_key, analysis_max_tokens, analysis_temp, current_task_tier], outputs=[analysis_report_out] ) # TAB 2: MANUAL EVALUATION CONSOLE with gr.Tab("Manual Evaluation Console"): current_obs = gr.State({}) with gr.Row(): # LEFT SIDEBAR - Environment Status & Actions with gr.Column(scale=1): with gr.Group(): gr.Markdown("### Dashboard") with gr.Row(): with gr.Column(elem_classes=["stat-box"]): step_counter = gr.Markdown("
0/2
Step
") with gr.Column(elem_classes=["stat-box"]): node_display = gr.Markdown("
Idle
Node
") manual_task_tier = gr.Dropdown( choices=["image_audit", "clip_audit", "weight_audit"], value="image_audit", label="Manual Audit Mode", ) manual_ingestion_id = gr.Textbox(label="Enter Ingestion ID to load session") reset_btn = gr.Button("Initialize Scenario", variant="primary") # Auto-fill manual ingestion id if one was generated in Tab 1 ingestion_id_out.change(fn=lambda x: x, inputs=ingestion_id_out, outputs=manual_ingestion_id) current_task_tier.change( fn=lambda x: x if x in ["image_audit", "clip_audit", "weight_audit"] else "image_audit", inputs=[current_task_tier], outputs=[manual_task_tier], ) with gr.Accordion("Raw Observation Payload", open=False): raw_obs_view = gr.JSON(label="Current Context") # RIGHT MAIN PANEL - Interactive Data with gr.Column(scale=2): obs_markdown = gr.Markdown(value=format_observation({})) # STEP 1: ImageDiagnosticsAction with gr.Group(visible=False) as grp_step1: gr.Markdown("## Action 1: Image Diagnostics") s1_regime = gr.Dropdown( choices=["frontal_simple", "non_frontal", "complex_background", "occluded", "low_quality"], label="Regime Classification", value="frontal_simple" ) s1_risk_factors = gr.Textbox(label="Identified Risk Factors (JSON list)", value='["lateral_pose_risk"]') s1_usability = gr.Slider(0.0, 1.0, value=0.5, label="Image Usability Score") s1_reasoning = gr.Textbox(label="Reasoning", value="Image analysis completed.") s1_submit = gr.Button("Execute Diagnostics", variant="primary") # STEP 2: ParamAnomalyAction with gr.Group(visible=False) as grp_step2: gr.Markdown("## Action 2: Parameter Anomaly Detection") s2_risk = gr.Dropdown( choices=["safe", "marginal", "risky", "dangerous"], label="Config Risk Level", value="marginal" ) s2_anomalies = gr.Code( label="Anomalies (JSON array of ParameterAnomaly dicts)", language="json", value='[]' ) s2_summary = gr.Textbox(label="Risk Summary", value="Parameter analysis completed.") s2_submit = gr.Button("Execute Parameter Check", variant="primary") # CLIP STEP: ClipDispositionAction with gr.Group(visible=False) as grp_clip: gr.Markdown("## Action: Clip Disposition") c_disposition = gr.Dropdown( choices=["accept", "reject", "fix", "defer"], label="Disposition", value="accept", ) c_confidence = gr.Slider(0.0, 1.0, value=0.6, label="Confidence") c_rejection = gr.Code( label="Rejection Reasons (JSON array)", language="json", value="[]", ) c_fixes = gr.Code( label="Fix Instructions (JSON array)", language="json", value="[]", ) c_fix_effort = gr.Dropdown( choices=["", "trivial", "moderate", "high"], value="", label="Estimated Fix Effort (optional)", ) c_defer = gr.Textbox(label="Defer Reason (optional)", value="") c_reasoning = gr.Textbox( label="Dataset Impact Reasoning", value="Disposition reasoning based on dossier quality and coverage impact.", ) c_override = gr.Dropdown( choices=["not_applicable", "declined", "applied"], value="not_applicable", label="Override Decision", ) c_override_just = gr.Textbox( label="Override Justification (optional)", value="", ) c_submit = gr.Button("Execute Clip Disposition", variant="primary") # STEP 3: PhonemeRiskAction with gr.Group(visible=False) as grp_step3: gr.Markdown("## Action 3: Phoneme Behavioral Risk Assessment") s3_safety = gr.Dropdown( choices=["safe", "minor_concerns", "moderate_risk", "high_risk", "unsafe"], label="Model Behavioral Safety", value="minor_concerns" ) s3_risks = gr.Code( label="Phoneme Risk Ranking (JSON array)", language="json", value='[]' ) s3_summary = gr.Textbox(label="Overall Behavioral Summary", value="Behavioral evaluate completed.") s3_submit = gr.Button("Execute Final Assessment", variant="primary") # DONE VIEW with gr.Group(visible=False) as grp_done: gr.Markdown("## 🏁 Episode Complete") score_markdown = gr.Markdown() # State update flow def update_ui_state(obs: Dict[str, Any]): if not obs: return [{}, format_observation({}), "
0/2
Step
", "
Error
Node
", False, False, False, False, False, ""] step = obs.get("step", 0) node = obs.get("node", "Unknown") is_done = getattr(obs, "done", obs.get("done", False)) # Handle both object getattr and dict get schema = obs.get("expected_action_schema", "") scores = obs.get("scores", None) step_str = f"
{step}/2
Step
" node_short = node.split(" ")[0] if " " in node else node node_str = f"
{node_short}
Node
" score_md = "" if is_done and scores: html_scores = "" for k, v in scores.items(): val = f"{v:.3f}" if isinstance(v, float) else str(v) html_scores += f"
  • {k}: {val}
    • " score_md = f"

    Evaluation finished. Change ingestion ID and Reset to run new scenario.

    " return [ obs, format_observation(obs), step_str, node_str, gr.update(visible=(not is_done and schema == "ImageDiagnosticsAction")), gr.update(visible=(not is_done and schema == "ParamAnomalyAction")), gr.update(visible=(not is_done and schema == "ClipDispositionAction")), gr.update(visible=(not is_done and schema == "PhonemeRiskAction")), gr.update(visible=is_done), gr.update(value=score_md) ] async def do_reset(ingestion_id, selected_tier): mode_map = { "image_audit": "image", "clip_audit": "clips", "weight_audit": "weights", } payload = {"mode": mode_map.get(selected_tier, "image")} if str(ingestion_id or "").strip(): payload["ingestion_id"] = str(ingestion_id) res = await web_manager.reset_environment(payload) obs = res.get("observation", {}) return update_ui_state(obs) async def do_step1(regime, risks_str, usability, reasoning): try: risks = json.loads(risks_str) except: risks = [] action_payload = { "regime_classification": regime, "identified_risk_factors": risks, "image_usability_score": float(usability), "reasoning": reasoning, "prompt_issues": [], "recommended_prompt_modifications": [], } res = await web_manager.step_environment(action_payload) obs = res.get("observation", {}) obs["done"] = res.get("done", False) return update_ui_state(obs) async def do_step2(risk_level, anomalies_str, summary): try: anoms = json.loads(anomalies_str) except: anoms = [] action_payload = { "config_risk_level": risk_level, "anomalies": anoms, "summary": summary, "predicted_failure_modes": [], "directional_fixes": [], } res = await web_manager.step_environment(action_payload) obs = res.get("observation", {}) obs["done"] = res.get("done", False) return update_ui_state(obs) async def do_step3(safety_level, risks_str, summary): try: risks = json.loads(risks_str) except: risks = [] action_payload = { "model_behavioral_safety": safety_level, "phoneme_risk_ranking": risks, "summary": summary, "predicted_behavior_triggers": [], "risky_phoneme_clusters": [], "mitigation_recommendations": [], } res = await web_manager.step_environment(action_payload) obs = res.get("observation", {}) obs["done"] = res.get("done", False) return update_ui_state(obs) async def do_clip_step( disposition, confidence, rejection_str, fixes_str, fix_effort, defer_reason, reasoning, override_decision, override_justification, ): try: rejection_reasons = json.loads(rejection_str) if not isinstance(rejection_reasons, list): rejection_reasons = [] except: rejection_reasons = [] try: fix_instructions = json.loads(fixes_str) if not isinstance(fix_instructions, list): fix_instructions = [] except: fix_instructions = [] action_payload = { "disposition": disposition, "confidence": float(confidence), "rejection_reasons": rejection_reasons or None, "fix_instructions": fix_instructions or None, "estimated_fix_effort": fix_effort or None, "defer_reason": (defer_reason or "").strip() or None, "dataset_impact_reasoning": reasoning, "override_decision": override_decision, "override_justification": (override_justification or "").strip() or None, } res = await web_manager.step_environment(action_payload) obs = res.get("observation", {}) obs["done"] = res.get("done", False) return update_ui_state(obs) outputs = [ raw_obs_view, obs_markdown, step_counter, node_display, grp_step1, grp_step2, grp_clip, grp_step3, grp_done, score_markdown ] reset_btn.click(do_reset, inputs=[manual_ingestion_id, manual_task_tier], outputs=outputs) s1_submit.click(do_step1, inputs=[s1_regime, s1_risk_factors, s1_usability, s1_reasoning], outputs=outputs) s2_submit.click(do_step2, inputs=[s2_risk, s2_anomalies, s2_summary], outputs=outputs) c_submit.click( do_clip_step, inputs=[ c_disposition, c_confidence, c_rejection, c_fixes, c_fix_effort, c_defer, c_reasoning, c_override, c_override_just, ], outputs=outputs, ) s3_submit.click(do_step3, inputs=[s3_safety, s3_risks, s3_summary], outputs=outputs) return demo