SHUBHAMOS/meta-pytorch-hackathon
0
1"""2SHUBHAMOS: AI Email Operations & Triage Environment3inference.py — LLM Agent Loop (Phase 4)4 5Runs an AI agent against the email triage environment using the OpenAI client6pointed at Hugging Face Router (Qwen/Qwen2.5-72B-Instruct).7 8Usage:9 HF_TOKEN=<your_token> python inference.py --task easy10 HF_TOKEN=<your_token> python inference.py --task medium --max-steps 5011 HF_TOKEN=<your_token> python inference.py --task all # run all 3 tasks12 13Environment variables:14 HF_TOKEN - Hugging Face API token (required)15 API_BASE_URL - Override API base (default: https://router.huggingface.co/v1)16 MODEL_NAME - Override model name (default: Qwen/Qwen2.5-72B-Instruct)17"""18 19import argparse20import json21import os22import sys23import time24import re25from typing import Any, Dict, List, Optional26from dotenv import load_dotenv27 28# Load secrets from .env file29load_dotenv()30 31from openai import OpenAI32 33from server.environment import EmailTriageEnv34from server.models import Action, Observation35from server.reward import RewardEngine36from server.tasks import TASKS37from server.graders import EasyGrader, MediumGrader, HardGrader, PeacefulGrader, ExtremeGrader38 39# ── AI Client Setup ──────────────────────────────────────────────────────────40 41# Load default config from environment42API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")43MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")44HF_TOKEN = os.getenv("HF_TOKEN")45LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")46 47# Global storage for clients to avoid re-initializing if token hasn't changed.48_primary_client = None49_internal_client = None50_current_token = None51 52def setup_clients(hf_token: Optional[str] = None):53 """Initializes or updates the OpenAI clients with a specific token."""54 global _primary_client, _internal_client, _current_token55 token = hf_token or HF_TOKEN56 57 if not token and not os.getenv("INTERNAL_AI_KEY"):58 print(" [Setup Error] No HF_TOKEN provided.")59 return None, None60 61 if _primary_client and _current_token == token:62 return _primary_client, _internal_client63 64 _current_token = token65 base_url = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1")66 model_name = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")67 68 _primary_client = OpenAI(base_url=base_url, api_key=token, max_retries=0)69 70 # Internal AI Fallback71 internal_url = os.environ.get("INTERNAL_AI_URL", "https://integrate.api.nvidia.com/v1")72 internal_model = os.environ.get("INTERNAL_AI_MODEL", "qwen/qwen3.5-122b-a10b")73 internal_key = os.environ.get("INTERNAL_AI_KEY", "")74 75 _internal_client = OpenAI(base_url=internal_url, api_key=internal_key, max_retries=0) if internal_key else None76 77 return _primary_client, _internal_client78 79GRADERS = {80 "peaceful": PeacefulGrader,81 "easy": EasyGrader,82 "medium": MediumGrader,83 "hard": HardGrader,84 "extreme": ExtremeGrader,85}86 87# ── System prompt ─────────────────────────────────────────────────────────────88 89SYSTEM_PROMPT = """You are an AI email triage agent.90 91Your GOAL is to process emails efficiently. 92PRIORitize classify_email for incoming emails before other actions.93 94You MUST respond ONLY with valid JSON.95No explanations, no markdown, no extra text.96 97STRICT FORMAT:98{99"action_type": "<one of: classify_email, set_priority, draft_reply, mark_resolved, escalate_email, ignore_email>",100"email_id": "<email_id>",101"category": "<optional>",102"priority": "<optional>"103}104 105EXAMPLES:106Example 1 (Billing):107Input: billing issue from customer108Output: { "action_type": "classify_email", "email_id": "email_1", "category": "billing" }109 110Example 2 (Urgency):111Input: urgent complaint112Output: { "action_type": "set_priority", "email_id": "email_2", "priority": "high" }113 114Return ONLY JSON. 115DO NOT include any explanation. 116DO NOT include text before or after JSON.117"""118 119# ── Observation → prompt ──────────────────────────────────────────────────────120 121def obs_to_prompt(obs: Observation, email_id: str, prev_action_result: str = "") -> str:122 """Focuses the LLM on exactly ONE email to maximize precision and avoid confusion."""123 email = next((e for e in obs.emails if e.id == email_id), None)124 if not email:125 return "No emails pending. Respond with {'action_type': 'ignore_email', 'email_id': 'none'}"126 127 lines = [128 f"=== FOCUS: EMAIL {email.id} ===",129 f"Subject: {email.subject}",130 f"From: {email.sender}",131 f"Current Status: {email.category or 'UNCLASSIFIED'} / {email.priority or 'NONE'}",132 f"Body: {email.body_preview[:250]}",133 "",134 "HISTORY:",135 f"Last action result: {prev_action_result or 'Start of flow'}",136 "",137 "GOAL for this email:",138 "1. If UNCLASSIFIED -> action_type: classify_email (e.g. billing_issue, tech_support)",139 "2. If Priority NONE -> action_type: set_priority (e.g. high, medium, low)",140 "3. If Classified & Prioritized -> action_type: mark_resolved",141 "",142 "Return ONLY JSON. No explanation. No extra text."143 ]144 return "\n".join(lines)145 146 147# ── API Health Check ─────────────────────────────────────────────────────────148 149def handle_rate_limit(provider_name: str) -> None:150 """Requested 429 handling with loasing animation."""151 print(f"\n [Rate Limit] {provider_name}: Ai got rate limitws")152 print(" Waiting 10 seconds...")153 pass # Removed mock loading for production speed154 155def check_client_health() -> bool:156 """Run a single test prompt to see if the Primary AI is active."""157 if not _primary_client:158 return False159 160 model_name = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")161 print(f" [Health Check] Testing Primary AI ({model_name}) with JSON Mode...")162 try:163 response = _primary_client.chat.completions.create(164 model=model_name,165 messages=[{"role": "user", "content": "Respond with {'status': 'ok'} in JSON format."}],166 max_tokens=20,167 timeout=10,168 response_format={"type": "json_object"}169 )170 if response.choices[0].message.content:171 print(" [Health Check] Primary AI is HEALTHY ✅")172 return True173 except Exception as e:174 print(f" [Health Check] Primary AI failed: {str(e)[:80]} ❌")175 return False176 177# ── LLM call ─────────────────────────────────────────────────────────────────178 179def _safe_llm_call(client, model, messages, timeout=20):180 """Internal helper to try JSON mode with a fallback to standard completions."""181 try:182 # Final optimization: Strict Token & Temperature Control183 response = client.chat.completions.create(184 model=model,185 messages=messages,186 temperature=0.05, # Ultra-low for consistency187 max_tokens=80, # Lean JSON only188 timeout=timeout,189 response_format={"type": "json_object"}190 )191 return response.choices[0].message.content192 except Exception as e:193 # If response_format is NOT supported, fallback to standard call194 if "response_format" in str(e) or "json_object" in str(e):195 response = client.chat.completions.create(196 model=model,197 messages=messages,198 temperature=0.05,199 max_tokens=80,200 timeout=timeout201 )202 return response.choices[0].message.content203 raise e204 205def call_llm(conversation: List[Dict[str, str]], primary_healthy: bool = True, clients: tuple = (None, None)) -> str:206 """207 Call the LLM with three-tier logic: Primary AI -> Experimental AI (Fallback).208 JSON Mode enforcement included.209 """210 p_client, i_client = clients211 primary_model = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")212 internal_model = os.environ.get("INTERNAL_AI_MODEL", "qwen/qwen3.5-122b-a10b")213 214 # ── TIER 1: Primary AI if healthy ──────────────────────────────────215 if p_client and primary_healthy:216 try:217 content = _safe_llm_call(p_client, primary_model, conversation, timeout=15)218 if content:219 time.sleep(1) # Fast throttle for success220 return content.strip()221 except Exception as e:222 if "429" in str(e):223 handle_rate_limit("Primary AI")224 elif "402" in str(e):225 print(" [Primary Error] Credits depleted (402). Switching to Fallback.")226 else:227 print(f" [Primary Error] {str(e)[:100]}")228 229 # ── TIER 2: Secondary AI Fallback (Secret) ───────────────────────────────────230 if i_client:231 try:232 content = _safe_llm_call(i_client, internal_model, conversation, timeout=20)233 if content:234 return content.strip()235 except Exception as e:236 print(f" [Fallback Error] {str(e)[:80]}")237 238 # ── TIER 3: Desperation Primary (even if failed health) ──────────────────239 # Only try this if we haven't already tried it in Tier 1240 if p_client and not primary_healthy:241 try:242 content = _safe_llm_call(p_client, primary_model, conversation, timeout=25)243 if content:244 return content.strip()245 except Exception:246 pass247 248 return ""249 250def get_fallback_action(obs: Observation, email_id: Optional[str] = None) -> Action:251 """Deterministic SMART fallback logic when LLM fails."""252 if not obs.emails:253 return Action(action_type="ignore_email", email_id="none")254 255 # Try to find target email256 target_id = email_id257 if not target_id:258 # Default to first unresolved email if possible259 target_id = obs.emails[0].id260 261 email = next((e for e in obs.emails if e.id == target_id), obs.emails[0])262 text = (email.subject + " " + email.body_preview).lower()263 264 # Priority Keywords265 is_urgent = any(k in text for k in ["urgent", "asap", "critical", "blocking", "emergency", "immediately"])266 267 # Classification Keywords268 if "billing" in text or "invoice" in text or "payment" in text or "charged" in text:269 return Action(action_type="classify_email", email_id=email.id, category="billing_issue")270 if "issue" in text or "bug" in text or "broken" in text or "not working" in text:271 return Action(action_type="classify_email", email_id=email.id, category="tech_support")272 273 if is_urgent and email.priority != "high":274 return Action(action_type="set_priority", email_id=email.id, level="high")275 276 # Final smart safety: resolve if it looks handled, otherwise classify general277 if email.category and email.priority:278 return Action(action_type="mark_resolved", email_id=email.id)279 280 return Action(action_type="classify_email", email_id=email.id, category="general_inquiry")281 282 # Robust default: classify as general283 return Action(action_type="classify_email", email_id=email.id, category="general_inquiry")284 285def parse_action(raw: str, obs: Observation) -> Action:286 """Parse LLM output focusing on structured JSON Mode response."""287 print(f" [DEBUG] RAW MODEL OUTPUT:\n{raw}\n{'-'*40}")288 289 if not raw or "{" not in raw:290 return get_fallback_action(obs)291 292 try:293 # Try direct load first (ideal for JSON Mode)294 try:295 data = json.loads(raw.strip())296 except json.JSONDecodeError:297 # Simple single-pattern extraction if chatter persists298 match = re.search(r'\{.*\}', raw, re.DOTALL)299 if not match: return get_fallback_action(obs)300 data = json.loads(match.group(0))301 302 if not isinstance(data, dict) or "action_type" not in data:303 # Handle "action" hallucination304 if "action" in data:305 data["action_type"] = data.pop("action")306 else:307 return get_fallback_action(obs)308 309 action_type = data.get("action_type")310 email_id = data.get("email_id")311 312 valid_actions = {313 "classify_email", "set_priority", "draft_reply", 314 "mark_resolved", "escalate_email", "ignore_email"315 }316 317 if action_type not in valid_actions or not email_id:318 return get_fallback_action(obs, email_id)319 320 known_ids = {e.id for e in obs.emails}321 if email_id not in known_ids:322 return get_fallback_action(obs)323 324 # Normalization for Category325 cat = data.get("category")326 if cat:327 cat = str(cat).lower()328 if "bill" in cat: cat = "billing_issue"329 elif "urgent" in cat or "complaint" in cat: cat = "urgent_complaint"330 elif "tech" in cat or "support" in cat: cat = "tech_support"331 elif "general" in cat: cat = "general_inquiry"332 elif "spam" in cat: cat = "spam"333 else: cat = "general_inquiry" # fallback safe334 335 # Normalization for Priority Level336 level = data.get("priority") or data.get("level") or data.get("priority_level")337 if level:338 level = str(level).lower()339 if "high" in level or "urgent" in level or "critical" in level: level = "high"340 elif "med" in level: level = "medium"341 elif "low" in level: level = "low"342 else: level = "unknown"343 344 action = Action(345 action_type=action_type,346 email_id=email_id,347 category=cat,348 level=level,349 text=data.get("reply_text") or data.get("text") or data.get("reply")350 )351 return action352 except Exception as e:353 print(f" [Parse Error] {e}")354 return get_fallback_action(obs)355 356 357# ── Agent loop ────────────────────────────────────────────────────────────────358 359def run_agent(task_id: str, hf_token: Optional[str] = None, verbose: bool = True) -> Dict[str, Any]:360 """361 Run one full episode of the email triage agent.362 363 Returns: grading report dict364 """365 clients = setup_clients(hf_token)366 p_client, i_client = clients367 368 if not p_client:369 return {"error": "AI client not initialized. Check HF_TOKEN."}370 371 # Step: Check AI Health (User's "Check then Use" request)372 primary_healthy = check_client_health_with_client(p_client)373 374 task_cls = TASKS[task_id]375 task_config = task_cls.config()376 377 # Setup environment + reward engine378 env = EmailTriageEnv()379 engine = RewardEngine()380 env.attach_reward_engine(engine)381 382 obs = env.reset(task_config)383 384 if verbose:385 model_name_env = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")386 print(f"[START] task={task_id} env=shubhamos model={model_name_env}", flush=True)387 print(f"\n{'='*60}")388 print(f" SHUBHAMOS — Task: {task_id.upper()} | {task_cls.email_count} emails | max {task_cls.max_steps} steps")389 print(f"{'='*60}")390 391 total_reward = 0.0392 parse_success = 0393 parse_failure = 0394 failure_streak = 0395 last_action_summary = ""396 step_times: List[float] = []397 rewards_list: List[float] = []398 399 for step in range(task_cls.max_steps):400 # STOP EARLY: All handled?401 pending = [e for e in obs.emails if not (e.resolved or e.escalated or e.ignored)]402 if not pending:403 if verbose: print(" [Early Stop] All emails processed. Ending episode.")404 break405 406 # TARGET SELECTION: Focus on the first pending email407 target_email = pending[0]408 409 # Build user message (Stateful & Focused)410 user_msg = obs_to_prompt(obs, target_email.id, prev_action_result=last_action_summary)411 412 # FAST FAIL: If LLM is failing repeatedly, switch to pure fallback for this step413 if failure_streak >= 2:414 if verbose: print(" [Fast Fail] Consecutive failures. Reverting to Smart Fallback.")415 action = get_fallback_action(obs, target_email.id)416 setattr(action, "_is_fallback", True)417 raw = "{}" # Dummy418 failure_streak = 0 # reset streak after one fallback419 else:420 conversation = [421 {"role": "system", "content": SYSTEM_PROMPT},422 {"role": "user", "content": user_msg}423 ]424 t0 = time.time()425 raw = call_llm(conversation, primary_healthy=primary_healthy, clients=clients)426 step_times.append(time.time() - t0)427 action = parse_action(raw, obs)428 429 # Trace failure for metrics & streak430 is_llm_failure = not raw or raw == "{}" or getattr(action, "_is_fallback", False) or action.action_type not in raw431 432 if is_llm_failure:433 parse_failure += 1434 failure_streak += 1435 if not getattr(action, "_is_fallback", False): # if parse failed but streak not yet triggered436 action = get_fallback_action(obs, target_email.id)437 else:438 parse_success += 1439 failure_streak = 0440 441 # Apply action442 done = False443 error = None444 try:445 obs, reward, done, info = env.step(action)446 total_reward += reward447 last_action_summary = f"Success: {action.action_type} on {action.email_id}"448 except Exception as e:449 error = str(e)[:50].replace('\n', ' ')450 last_action_summary = f"Error: {error}"451 # One last try with fallback452 try:453 action = get_fallback_action(obs, target_email.id)454 obs, reward, done, info = env.step(action)455 total_reward += reward456 error = None457 except Exception as e2:458 reward = 0.0459 error = str(e2)[:50].replace('\n', ' ')460 done = True461 462 rewards_list.append(reward)463 464 if verbose:465 done_val = "true" if done else "false"466 error_val = "null" if not error else f"'{error}'"467 print(f"[STEP] step={step+1} action={action.action_type} reward={reward:.2f} done={done_val} error={error_val}", flush=True)468 print(f" Step {step+1:02d} | Action: {action.action_type} | Reward: {reward:+.2f} | Fallback: {'Yes' if is_llm_failure else 'No'}")469 470 if done: break471 time.sleep(2) # Key protection472 473 # Final Output...474 final_state = env.state()475 report = GRADERS[task_id]().grade(final_state)476 477 result = report.to_dict()478 score = result.get("scores", {}).get("final_score", 0.0)479 480 if verbose:481 rewards_str = ",".join(f"{r:.2f}" for r in rewards_list)482 success_val = "true" if score > 0 else "false"483 steps_taken = len(rewards_list)484 print(f"[END] success={success_val} steps={steps_taken} score={score:.3f} rewards={rewards_str}", flush=True)485 486 result["total_reward"] = round(total_reward, 4)487 return result488 489def check_client_health_with_client(client) -> bool:490 """Run a single test prompt to see if the client is active."""491 if not client: return False492 model_name = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")493 print(f" [Health Check] Testing AI ({model_name}) with JSON Mode...")494 try:495 response = client.chat.completions.create(496 model=model_name,497 messages=[{"role": "user", "content": "Respond with {'status': 'ok'} in JSON format."}],498 max_tokens=20,499 timeout=10,500 response_format={"type": "json_object"}501 )502 if response.choices[0].message.content:503 print(" [Health Check] Provider is HEALTHY ✅")504 return True505 except Exception as e:506 print(f" [Health Check] AI failed: {str(e)[:80]} ❌")507 return False508 return False509 510 511# ── Entry point ───────────────────────────────────────────────────────────────512 513def main() -> None:514 parser = argparse.ArgumentParser(515 description="SHUBHAMOS Email Triage Agent — runs LLM agent against environment"516 )517 parser.add_argument(518 "--task",519 choices=["peaceful", "easy", "medium", "hard", "extreme", "all"],520 default="all",521 help="Task difficulty to run (default: all — runs all 5 tasks)",522 )523 parser.add_argument("--verbose", action="store_true", default=True)524 parser.add_argument("--quiet", action="store_true", help="Suppress step-by-step output")525 parser.add_argument("--output", type=str, help="Write results JSON to this file")526 args = parser.parse_args()527 528 p_client, i_client = setup_clients()529 if not p_client and not i_client:530 print("ERROR: No AI clients configured. Check your .env file.")531 sys.exit(1)532 533 # Use the global clients initialized above534 verbose = not args.quiet535 536 tasks_to_run = ["peaceful", "easy", "medium", "hard", "extreme"] if args.task == "all" else [args.task]537 all_results: Dict[str, Any] = {}538 539 for task_id in tasks_to_run:540 result = run_agent(task_id, hf_token=None, verbose=verbose)541 all_results[task_id] = result542 543 # Print exact required output formats for final scores544 print()545 for task_id in tasks_to_run:546 score = all_results.get(task_id, {}).get("scores", {}).get("final_score", 0.0)547 print(f"FINAL SCORE ({task_id}): {score:.2f}")548 549 if args.output:550 with open(args.output, "w") as f:551 json.dump(all_results, f, indent=2)552 print(f"\nResults written to: {args.output}")553 554 555if __name__ == "__main__":556 main()557