raunakratan/priority-mind-lite
0
1"""2PriorityMind-Lite: Gradio Web Interface for Hugging Face Spaces3================================================================4This provides an interactive demo of the PriorityMind-Lite environment5where users can interact with the AI customer support triage agent.6 7Deployed at: https://huggingface.co/spaces/raunakratan/priority-mind-lite8"""9 10from __future__ import annotations11 12import asyncio13import json14import os15from pathlib import Path16from typing import Any17 18if os.name == "nt":19 try:20 asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())21 except AttributeError:22 pass23 24import gradio as gr25from dotenv import load_dotenv26 27from environment import PriorityMindEnv28from inference import mock_action29from models import Action, Observation30from utils import format_partial_signal31 32load_dotenv(Path(__file__).parent / ".env")33 34API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")35MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")36HF_TOKEN = os.getenv("HF_TOKEN", "").strip()37 38# Try to import openai for live mode39try:40 from openai import OpenAI41 42 OPENAI_AVAILABLE = True43except ImportError:44 OPENAI_AVAILABLE = False45 46 47APP_THEME = gr.themes.Soft()48 49 50def get_client() -> OpenAI | None:51 """Get OpenAI client for HF Router if available."""52 if not HF_TOKEN or not OPENAI_AVAILABLE:53 return None54 return OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN, timeout=12)55 56 57def configure_runtime() -> None:58 """Use a Windows-friendly event loop policy for Gradio/Uvicorn."""59 if os.name != "nt":60 return61 try:62 asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())63 except (AttributeError, RuntimeError):64 pass65 66 67def _llm_prompt(obs: Observation) -> str:68 """Build the LLM prompt for action selection."""69 return f"""You are a customer support triage agent.70 71Observation:72- ticket_text: {obs.ticket_text}73- sentiment: {obs.sentiment:.2f}74- category: {obs.category}75- priority: {obs.priority}76- attempts: {obs.attempts}77- resolved: {str(obs.resolved).lower()}78 79Choose exactly one next action:80- categorize with content billing/technical/general/complaint81- prioritize with priority low/medium/high/urgent82- respond with a concise, empathetic support message83- escalate84- resolve85 86Return JSON only with keys action_type, content, priority.87"""88 89 90def choose_action(obs: Observation, task: str, client: OpenAI | None) -> tuple[Action, str]:91 """Choose the next action using LLM if available, else use heuristic.92 93 Returns:94 A tuple of (action, action_source) where action_source is one of:95 "LLM", "heuristic", or "heuristic fallback".96 """97 if client is not None:98 try:99 response = client.chat.completions.create(100 model=MODEL_NAME,101 messages=[{"role": "user", "content": _llm_prompt(obs)}],102 temperature=0.0,103 response_format={"type": "json_object"},104 )105 payload = response.choices[0].message.content106 if payload:107 data = json.loads(payload)108 action = Action(**data)109 return action, "LLM"110 except Exception:111 return mock_action(obs, task), "heuristic fallback"112 113 return mock_action(obs, task), "heuristic"114 115 116def format_action(action: Action) -> str:117 """Format action for display."""118 parts = [action.action_type.upper()]119 if action.content:120 parts.append(f'"{action.content[:60]}{"..." if len(action.content) > 60 else ""}"')121 if action.priority:122 parts.append(f"priority={action.priority}")123 return " -> ".join(parts)124 125 126def _action_mode_label(mode: str) -> str:127 labels = {128 "llm_only": "LLM policy",129 "mixed": "LLM policy with heuristic fallback",130 "heuristic_only": "Heuristic policy",131 "heuristic_fallback": "Heuristic policy after LLM fallback",132 }133 return labels.get(mode, mode)134 135 136def _reward_mode_label(mode: str) -> str:137 labels = {138 "hybrid_llm": "Hybrid grading (LLM + rules)",139 "hybrid_mixed": "Hybrid grading with programmatic fallback",140 "programmatic_fallback": "Programmatic fallback grading",141 "programmatic_only": "Programmatic grading only",142 }143 return labels.get(mode, mode)144 145 146def run_episode(task: str, use_llm: bool) -> tuple[list[dict[str, Any]], dict[str, Any]]:147 """Run a complete episode and return step-by-step results plus mode summary."""148 client = get_client() if use_llm else None149 # Use a fixed seed for reproducible demo results150 env = PriorityMindEnv(task=task, seed=42, enable_llm=client is not None)151 obs = env.reset()152 153 steps = []154 step_num = 0155 llm_action_steps = 0156 reward_llm_steps = 0157 reward_fallback_steps = 0158 159 while True:160 step_num += 1161 action, action_source = choose_action(obs, task, client)162 if action_source == "LLM":163 llm_action_steps += 1164 165 obs, reward, done, info = env.step(action)166 reward_used_fallback = bool(info.get("used_fallback", False))167 if not env.grader.enable_llm:168 reward_source = "Programmatic grading"169 else:170 reward_source = "Programmatic fallback" if reward_used_fallback else "LLM judge"171 if reward_used_fallback:172 reward_fallback_steps += 1173 else:174 reward_llm_steps += 1175 176 step_info = {177 "step": step_num,178 "observation": {179 "ticket_text": obs.ticket_text,180 "sentiment": f"{obs.sentiment:+.2f}",181 "category": obs.category or "Not set",182 "priority": obs.priority or "Not set",183 "attempts": obs.attempts,184 "resolved": obs.resolved,185 },186 "action": format_action(action),187 "action_source": action_source,188 "action_error": info.get("last_action_error"),189 "reward": {190 "score": f"{reward.score:.2f}",191 "reasoning": reward.reasoning[:200] + "..." if len(reward.reasoning) > 200 else reward.reasoning,192 },193 "reward_source": reward_source,194 "partial_signals": {195 "empathy": format_partial_signal(reward.partial_signals.get("empathy")),196 "efficiency": format_partial_signal(reward.partial_signals.get("efficiency")),197 "strategy": format_partial_signal(reward.partial_signals.get("strategy")),198 },199 "done": done,200 }201 steps.append(step_info)202 203 if done:204 break205 206 if not use_llm or client is None:207 action_mode = "heuristic_only"208 elif llm_action_steps == len(steps):209 action_mode = "llm_only"210 elif llm_action_steps > 0:211 action_mode = "mixed"212 else:213 action_mode = "heuristic_fallback"214 215 if not env.grader.enable_llm:216 reward_mode = "programmatic_only"217 elif reward_llm_steps == len(steps):218 reward_mode = "hybrid_llm"219 elif reward_llm_steps > 0:220 reward_mode = "hybrid_mixed"221 else:222 reward_mode = "programmatic_fallback"223 224 summary = {225 "action_mode": action_mode,226 "reward_mode": reward_mode,227 "client_available": client is not None,228 "llm_requested": use_llm,229 "total_steps": len(steps),230 "llm_action_steps": llm_action_steps,231 "heuristic_action_steps": len(steps) - llm_action_steps,232 "reward_llm_steps": reward_llm_steps,233 "reward_fallback_steps": reward_fallback_steps,234 }235 return steps, summary236 237 238def run_demo(task: str, use_llm: bool) -> str:239 """Run demo and return formatted HTML results."""240 try:241 steps, summary = run_episode(task, use_llm)242 243 output = []244 output.append(f"<h1>PriorityMind-Lite Demo</h1>")245 output.append(f"<h2>Task: {task.upper()}</h2>")246 output.append(f"<p><strong>Action Policy:</strong> {_action_mode_label(summary['action_mode'])}</p>")247 output.append(f"<p><strong>Reward Grading:</strong> {_reward_mode_label(summary['reward_mode'])}</p>")248 249 if use_llm and not summary["client_available"]:250 output.append("<div class='alert alert-warning'>[!] Note: HF_TOKEN or OpenAI client unavailable, using heuristic actions and programmatic grading only.</div>")251 elif summary["action_mode"] == "mixed":252 output.append(f"<div class='alert alert-info'>[i] Action generation used LLM in {summary['llm_action_steps']}/{summary['total_steps']} steps, heuristic fallback for the rest.</div>")253 elif summary["action_mode"] == "heuristic_fallback":254 output.append("<div class='alert alert-warning'>[!] Every action fell back to heuristic policy after LLM failure.</div>")255 256 if summary["reward_mode"] == "hybrid_mixed":257 output.append(f"<div class='alert alert-info'>[i] Reward grading used LLM judge in {summary['reward_llm_steps']}/{summary['total_steps']} steps, programmatic fallback in the rest.</div>")258 elif summary["reward_mode"] == "programmatic_fallback":259 output.append("<div class='alert alert-warning'>[!] Reward grading attempted live evaluation, but every step fell back to programmatic grader.</div>")260 261 output.append("<div class='steps-container'>")262 263 for step_info in steps:264 output.append(f"<div class='step-card'>")265 output.append(f"<h3>Step {step_info['step']}</h3>")266 output.append(f"<p><strong>Action:</strong> {step_info['action']} <span class='badge badge-secondary'>({step_info['action_source']})</span></p>")267 output.append(f"<p><strong>Reward Source:</strong> {step_info['reward_source']}</p>")268 output.append(f"<p><strong>Reward Score:</strong> <span class='score'>{step_info['reward']['score']}</span></p>")269 output.append(f"<p><strong>Reasoning:</strong> {step_info['reward']['reasoning']}</p>")270 271 # Progress bars for partial signals272 empathy_val = float(step_info['partial_signals']['empathy'].split('/')[0])273 efficiency_val = float(step_info['partial_signals']['efficiency'].split('/')[0])274 strategy_val = float(step_info['partial_signals']['strategy'].split('/')[0])275 276 def get_color(val):277 if val >= 0.8:278 return "#28a745" # green279 elif val >= 0.6:280 return "#ffc107" # yellow281 else:282 return "#dc3545" # red283 284 output.append("<div class='signals'>")285 output.append(f"<div class='signal'><span class='icon'>EMPATHY</span> <div class='progress-bar'><div class='progress-fill' style='width: {empathy_val*100}%; background-color: {get_color(empathy_val)}'></div></div> <span class='signal-value'>{step_info['partial_signals']['empathy']}</span></div>")286 output.append(f"<div class='signal'><span class='icon'>EFFICIENCY</span> <div class='progress-bar'><div class='progress-fill' style='width: {efficiency_val*100}%; background-color: {get_color(efficiency_val)}'></div></div> <span class='signal-value'>{step_info['partial_signals']['efficiency']}</span></div>")287 output.append(f"<div class='signal'><span class='icon'>STRATEGY</span> <div class='progress-bar'><div class='progress-fill' style='width: {strategy_val*100}%; background-color: {get_color(strategy_val)}'></div></div> <span class='signal-value'>{step_info['partial_signals']['strategy']}</span></div>")288 output.append("</div>")289 290 if step_info["action_error"]:291 output.append(f"<p class='error'>[ERROR] Action Error: {step_info['action_error']}</p>")292 if step_info['done']:293 output.append("<p class='success'>[OK] Episode Complete</p>")294 output.append("</div>")295 296 output.append("</div>")297 298 # Summary with score chart299 scores = [float(s['reward']['score']) for s in steps]300 avg_score = sum(scores) / len(scores) if scores else 0301 min_score = min(scores) if scores else 0302 max_score = max(scores) if scores else 0303 304 # Simple bar chart for scores305 chart_html = "<div class='score-chart'>"306 for i, score in enumerate(scores, 1):307 color = get_color(score)308 chart_html += f"<div class='chart-bar' style='height: {score*100}px; background-color: {color}' title='Step {i}: {score:.2f}'></div>"309 chart_html += "</div>"310 311 output.append("<div class='summary'>")312 output.append("---")313 output.append("<h3>Summary</h3>")314 output.append(f"<p><strong>Average Score:</strong> <span class='avg-score'>{avg_score:.2f}</span></p>")315 output.append(f"<p><strong>Min Score:</strong> {min_score:.2f} | <strong>Max Score:</strong> {max_score:.2f}</p>")316 output.append(f"<p><strong>Total Steps:</strong> {len(steps)}</p>")317 output.append("<h4>Score Progression</h4>")318 output.append(chart_html)319 output.append("</div>")320 321 # Add CSS322 css = """323 <style>324 body { color: #000; }325 h1, h2, h3, h4, h5, h6 { color: #000 !important; }326 .alert { padding: 12px 15px; margin: 10px 0; border-radius: 5px; font-weight: 500; }327 .alert-warning { background-color: #fff3cd; border: 2px solid #ffc107; color: #000; }328 .alert-info { background-color: #d1ecf1; border: 2px solid #17a2b8; color: #000; }329 .alert-danger { background-color: #f8d7da; border: 2px solid #dc3545; color: #000; }330 .step-card { border: 2px solid #ddd; border-radius: 8px; padding: 20px; margin: 15px 0; background-color: #fff; }331 .step-card h3 { color: #000; margin: 0 0 10px 0; }332 .step-card p { color: #000; margin: 8px 0; }333 .badge { display: inline-block; padding: 4px 8px; font-size: 0.85em; border-radius: 3px; background-color: #6c757d; color: white; }334 .badge-secondary { background-color: #6c757d; }335 .score { font-size: 1.2em; font-weight: bold; color: #0056b3; }336 .signals { margin-top: 15px; }337 .signal { display: flex; align-items: center; margin: 10px 0; }338 .signal .icon { width: 100px; font-weight: bold; font-size: 0.9em; color: #000; }339 .progress-bar { flex: 1; height: 15px; background-color: #e9ecef; border-radius: 5px; margin: 0 10px; overflow: hidden; border: 1px solid #ccc; }340 .progress-fill { height: 100%; transition: width 0.3s; }341 .signal-value { width: 80px; text-align: right; font-weight: bold; color: #000; }342 .error { color: #000; font-weight: bold; }343 .success { color: #000; font-weight: bold; }344 .summary { margin-top: 30px; padding: 25px; background-color: #fff; border: 3px solid #007bff; border-radius: 8px; }345 .summary h3 { color: #000; margin: 0 0 15px 0; font-size: 1.5em; }346 .summary h4 { color: #000; margin: 15px 0 10px 0; }347 .summary p { color: #000; font-size: 1.1em; margin: 10px 0; font-weight: 500; }348 .avg-score { font-size: 1.3em; font-weight: bold; color: #000; }349 .score-chart { display: flex; align-items: end; justify-content: space-around; height: 150px; margin: 15px 0; padding: 15px; background-color: #f8f9fa; border-radius: 5px; border: 2px solid #ddd; }350 .chart-bar { width: 25px; min-height: 10px; border-radius: 3px 3px 0 0; margin: 0 3px; }351 .steps-container { margin: 20px 0; }352 </style>353 """354 355 return css + "\n".join(output)356 357 except Exception as e:358 return f"<div class='alert alert-danger'>Error: {str(e)}</div>"359 360 361def compare_modes() -> str:362 """Run comparison between LLM-rewarded and heuristic modes."""363 output = []364 output.append("<h2>Mode Comparison</h2>")365 366 for task in ["easy", "medium", "hard"]:367 output.append(f"<h3>Task: {task.upper()}</h3>")368 369 output.append("<div class='mode-section'>")370 output.append("<h4>LLM-Enabled Mode</h4>")371 try:372 steps_llm, summary_llm = run_episode(task, use_llm=True)373 scores_llm = [float(s['reward']['score']) for s in steps_llm]374 avg_llm = sum(scores_llm) / len(scores_llm) if scores_llm else 0375 output.append(f"<p><strong>Action Policy:</strong> {_action_mode_label(summary_llm['action_mode'])}</p>")376 output.append(f"<p><strong>Reward Grading:</strong> {_reward_mode_label(summary_llm['reward_mode'])}</p>")377 output.append(f"<p><strong>Average Score:</strong> <span class='score'>{avg_llm:.2f}</span></p>")378 output.append(f"<p><strong>Steps:</strong> {len(steps_llm)}</p>")379 except Exception as e:380 output.append(f"<p class='error'>Error: {str(e)}</p>")381 output.append("</div>")382 383 output.append("<div class='mode-section'>")384 output.append("<h4>Heuristic Mode</h4>")385 try:386 steps_heur, summary_heur = run_episode(task, use_llm=False)387 scores_heur = [float(s['reward']['score']) for s in steps_heur]388 avg_heur = sum(scores_heur) / len(scores_heur) if scores_heur else 0389 output.append(f"<p><strong>Action Policy:</strong> {_action_mode_label(summary_heur['action_mode'])}</p>")390 output.append(f"<p><strong>Reward Grading:</strong> {_reward_mode_label(summary_heur['reward_mode'])}</p>")391 output.append(f"<p><strong>Average Score:</strong> <span class='score'>{avg_heur:.2f}</span></p>")392 output.append(f"<p><strong>Steps:</strong> {len(steps_heur)}</p>")393 except Exception as e:394 output.append(f"<p class='error'>Error: {str(e)}</p>")395 output.append("</div>")396 397 output.append("<hr>")398 output.append("<p><em>Note: LLM-enabled mode requires HF_TOKEN to be configured. The app reports action generation and reward grading separately so fallback is explicit.</em></p>")399 400 css = """401 <style>402 h1, h2, h3, h4, h5, h6 { color: #000 !important; }403 .mode-section { border: 1px solid #ddd; border-radius: 8px; padding: 15px; margin: 10px 0; background-color: #f9f9f9; }404 .score { font-size: 1.2em; font-weight: bold; color: #000; }405 .error { color: #000; }406 </style>407 """408 409 return css + "\n".join(output)410 411 412# Create Gradio Interface413with gr.Blocks(title="PriorityMind-Lite") as demo:414 gr.Markdown("""415 # PriorityMind-Lite416 ### LLM-Rewarded Customer Support Ticket Triage Environment417 418 **Meta PyTorch OpenEnv Hackathon 2026 | Team Axiom (IIT Madras)**419 420 This demo showcases an AI agent that learns to triage customer support tickets421 using rewards evaluated by Llama — not hardcoded rules.422 """)423 424 with gr.Tab("Interactive Demo"):425 with gr.Row():426 with gr.Column():427 task_dropdown = gr.Dropdown(428 choices=["easy", "medium", "hard"],429 value="hard",430 label="Select Task Difficulty"431 )432 llm_toggle = gr.Checkbox(433 value=True,434 label="Use LLM for actions and grading (requires HF_TOKEN)"435 )436 run_button = gr.Button("Run Demo", variant="primary")437 438 with gr.Column():439 output_html = gr.HTML(label="Demo Results")440 441 run_button.click(442 fn=run_demo,443 inputs=[task_dropdown, llm_toggle],444 outputs=output_html445 )446 447 with gr.Tab("Mode Comparison"):448 compare_button = gr.Button("Compare enabled vs heuristic", variant="secondary")449 comparison_output = gr.HTML()450 451 with gr.Tab("Mode Comparison"):452 compare_button = gr.Button("Compare enabled vs heuristic", variant="secondary")453 comparison_output = gr.HTML()454 455 compare_button.click(456 fn=compare_modes,457 inputs=[],458 outputs=comparison_output459 )460 461 with gr.Tab("About"):462 gr.Markdown("""463 ### How It Works464 465 1. **Environment Reset**: Generate a customer support ticket with sentiment and true category466 2. **Agent Takes Action**: Choose from categorize, prioritize, respond, escalate, or resolve467 3. **Hybrid Grader Evaluates**:468 - LLM Evaluation (60%): Llama judges on empathy, efficiency, strategy469 - Fallback (40%): Deterministic rules ensure reliability470 4. **Reward Signal**: Normalized score [0.0, 1.0] with partial signals471 472 ### Key Innovation473 474 Instead of defining "good customer service" mathematically, we let a language model475 (Llama) judge each agent action on dimensions like empathy, efficiency, and strategy.476 477 ### Technology Stack478 479 - **Framework**: OpenEnv (Meta)480 - **Model**: Llama 3.1 8B Instruct (via HF Router)481 - **Language**: Python 3.10482 """)483 484# Mount OpenEnv server endpoints for hackathon validation485try:486 from environment import PriorityMindEnv487 from models import Action, Observation488 from openenv.core.env_server.http_server import create_app489 490 # Create OpenEnv HTTP server491 openenv_app = create_app(492 PriorityMindEnv,493 Action,494 Observation,495 env_name="priority-mind-lite",496 max_concurrent_envs=1,497 )498 499 # Mount OpenEnv endpoints at root level for hackathon validation500 demo.app.mount("/", openenv_app)501 502except Exception as e:503 print(f"Warning: Could not mount OpenEnv server endpoints: {e}")504 print("This is expected if openenv-core is not installed or if running locally.")505 506if __name__ == "__main__":507 configure_runtime()508 demo.launch(509 server_name=os.getenv("SERVER_NAME", "0.0.0.0"),510 server_port=int(os.getenv("PORT", "7860")),511 theme=APP_THEME,512 )513