Update inference.py

inference.py

@@ -1,5 +1,4 @@
 import os
-import json
 from openai import OpenAI
 from env import EmailSortingEnv
 
@@ -7,14 +6,17 @@ from env import EmailSortingEnv
 # SETUP — Read environment variables
 # ============================================
 
-API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
-MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME   = os.getenv("MODEL_NAME", "gpt-4o-mini")
+HF_TOKEN     = os.getenv("HF_TOKEN")
+
+if HF_TOKEN is None:
+    raise ValueError("HF_TOKEN environment variable is required")
 
 # Initialize OpenAI client
 client = OpenAI(
     base_url=API_BASE_URL,
-    api_key=HF_TOKEN if HF_TOKEN else "dummy-key"
+    api_key=HF_TOKEN
 )
 
 # ============================================
@@ -22,88 +24,53 @@ client = OpenAI(
 # ============================================
 
 def ask_llm_to_classify(email: dict) -> str:
-    """
-    Send email to LLM and get classification.
-    Returns: 'spam', 'important', or 'promotion'
-    """
-
     prompt = f"""You are an email classification assistant.
-
 Classify the following email into exactly ONE of these categories:
 - spam: unwanted, scam, phishing, prize winning, fake offers
-- important: work emails, order updates, bank alerts from real banks, newsletters
+- important: work emails, order updates, bank alerts, urgent notices
 - promotion: genuine sale offers, discount emails from real shops
 
-Email Details:
 Subject: {email['subject']}
 From: {email['sender']}
 Body: {email['body']}
 
-Reply with ONLY one word — either: spam, important, or promotion
-Do not explain. Just one word."""
+Reply with ONLY one word: spam, important, or promotion"""
 
     try:
         response = client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
-                {
-                    "role": "system",
-                    "content": "You are an email classifier. Reply with only one word: spam, important, or promotion."
-                },
-                {
-                    "role": "user",
-                    "content": prompt
-                }
+                {"role": "system", "content": "You are an email classifier. Reply with only one word: spam, important, or promotion."},
+                {"role": "user", "content": prompt}
             ],
             max_tokens=10,
             temperature=0.0
         )
-
-        # Extract the answer
         answer = response.choices[0].message.content.strip().lower()
-
-        # Clean up answer — only keep valid categories
-        if "spam" in answer:
-            return "spam"
-        elif "promotion" in answer:
-            return "promotion"
-        elif "important" in answer:
-            return "important"
-        else:
-            return "spam"  # Default fallback
-
+        if "spam" in answer:        return "spam"
+        elif "promotion" in answer: return "promotion"
+        elif "important" in answer: return "important"
+        else:                       return "spam"
     except Exception as e:
-        print(f"LLM Error: {e}")
-        # Fallback to simple rule-based classification
         return fallback_classify(email)
 
 
 def fallback_classify(email: dict) -> str:
-    """
-    Simple rule-based fallback if LLM fails.
-    """
     subject = email["subject"].lower()
-    body = email["body"].lower()
-    sender = email["sender"].lower()
-
-    spam_keywords = ["won", "free", "prize", "urgent", "money",
-                    "congratulations", "claim", "earn", "selected",
-                    "suspended", "verify", "action required"]
+    body    = email["body"].lower()
+    sender  = email["sender"].lower()
 
-    promo_keywords = ["off", "sale", "deal", "discount", "shop",
-                     "offer", "save", "limited time"]
+    spam_kw  = ["won", "free", "prize", "urgent", "money", "congratulations",
+                "claim", "earn", "selected", "suspended", "verify"]
+    promo_kw = ["off", "sale", "deal", "discount", "shop", "offer", "save", "limited time"]
 
-    spam_score = sum(1 for kw in spam_keywords if kw in subject or kw in body)
-    promo_score = sum(1 for kw in promo_keywords if kw in subject or kw in body)
+    spam_score  = sum(1 for kw in spam_kw  if kw in subject or kw in body)
+    promo_score = sum(1 for kw in promo_kw if kw in subject or kw in body)
+    suspicious  = any(d in sender for d in [".xyz", ".tk", "-secure", "-alert"])
 
-    suspicious_domain = any(d in sender for d in [".xyz", ".tk", "-secure", "-alert"])
-
-    if spam_score >= 2 or suspicious_domain:
-        return "spam"
-    elif promo_score >= 2:
-        return "promotion"
-    else:
-        return "important"
+    if spam_score >= 2 or suspicious: return "spam"
+    elif promo_score >= 2:             return "promotion"
+    else:                              return "important"
 
 
 # ============================================
@@ -111,105 +78,55 @@ def fallback_classify(email: dict) -> str:
 # ============================================
 
 def run_inference():
-    """
-    Main function — runs the AI agent for one full episode.
-    """
-    print("=" * 50)
-    print("Email Sorting Environment — Inference Script")
-    print("=" * 50)
-    print(f"Model: {MODEL_NAME}")
-    print(f"API Base: {API_BASE_URL}")
-    print("=" * 50)
-
-    # Initialize environment
-    env = EmailSortingEnv()
+    env   = EmailSortingEnv()
     state = env.reset()
 
-    print(f"\nStarting episode — max {state['max_steps']} steps\n")
-
+    rewards      = []
     step_results = []
+    error        = None
 
-    print(f"[START] task=email_sorting", flush=True)
-
-    # Run until episode is done
-    while not state["done"]:
-        current_step = state["step"] + 1
-        email = state["email"]
-
-        print(f"Subject: {email['subject']}")
-        print(f"From:    {email['sender']}")
-
-        # Ask AI to classify
-        action = ask_llm_to_classify(email)
-
-        # Take step in environment
-        next_state, reward, done, info = env.step(action)
-
-        result = info.get("result", "N/A")
-        print(f"[STEP] step={current_step} action={action} reward={reward} result={result}", flush=True)
-
-        step_results.append({
-            "step": current_step,
-            "subject": email["subject"],
-            "action": action,
-            "reward": reward,
-            "result": result
-        })
-
-        state = next_state
-
-    # ============================================
-    # FINAL RESULTS
-    # ============================================
-
-    total_reward = state["total_reward"]
-    total_steps = state["step"]
-    correct_count = sum(1 for r in step_results if r["result"] == "correct")
-
-    print("\n" + "=" * 50)
-    print("EPISODE COMPLETE")
-    print("=" * 50)
-    print(f"Total Steps:    {total_steps}")
-    print(f"Correct:        {correct_count}/{total_steps}")
-    print(f"Accuracy:       {round(correct_count/total_steps*100, 1)}%")
-    print(f"Total Reward:   {total_reward}")
-    print("=" * 50)
-
-    results = {
-        "model": MODEL_NAME,
-        "total_steps": total_steps,
-        "correct": correct_count,
-        "accuracy": round(correct_count / total_steps * 100, 1),
-        "total_reward": total_reward,
-        "step_details": step_results
-    }
+    # [START] — required format
+    print(f"[START] task=email_sorting env=email-sorting-openenv model={MODEL_NAME}", flush=True)
 
-    score = round(correct_count / total_steps, 4) if total_steps > 0 else 0.0
-    print(f"[END] task=email_sorting score={score} steps={total_steps}", flush=True)
+    try:
+        while not state["done"]:
+            current_step = state["step"] + 1
+            email        = state["email"]
 
-    return results
+            action = ask_llm_to_classify(email)
 
+            next_state, reward, done, info = env.step(action)
 
-# ============================================
-# RUN GRADERS ALSO
-# ============================================
+            error_msg = info.get("error", "null") or "null"
+            done_str  = "true" if done else "false"
 
-def run_with_graders():
-    """Run inference + all graders and show combined score."""
-    from graders import run_all_graders
+            # [STEP] — exact required format
+            print(f"[STEP] step={current_step} action={action} reward={reward:.2f} done={done_str} error={error_msg}", flush=True)
 
-    print("\n--- Running Inference ---\n")
-    inference_results = run_inference()
+            rewards.append(reward)
+            step_results.append({"step": current_step, "action": action, "reward": reward, "result": info.get("result", "N/A")})
 
-    print("\n--- Running Graders ---\n")
-    grader_results = run_all_graders()
+            state = next_state
 
-    print("\n" + "=" * 50)
-    print("FINAL COMBINED RESULTS")
-    print("=" * 50)
-    print(f"Inference Accuracy: {inference_results['accuracy']}%")
-    print(f"Grader Average Score: {grader_results['average_score']}")
-    print("=" * 50)
+        success = True
+
+    except Exception as e:
+        error   = str(e)
+        success = False
+
+    # [END] — exact required format
+    total_steps   = len(rewards)
+    rewards_str   = ",".join(f"{r:.2f}" for r in rewards)
+    success_str   = "true" if success else "false"
+    print(f"[END] success={success_str} steps={total_steps} rewards={rewards_str}", flush=True)
+
+    return {
+        "model":        MODEL_NAME,
+        "total_steps":  total_steps,
+        "total_reward": round(sum(rewards), 2),
+        "success":      success,
+        "step_details": step_results
+    }
 
 
 # ============================================
@@ -217,4 +134,4 @@ def run_with_graders():
 # ============================================
 
 if __name__ == "__main__":
-    run_with_graders()
+    run_inference()