Update inference.py

inference.py

@@ -1,125 +1,223 @@
-"""
-inference.py — Email Sorting OpenEnv
-Mandatory plain-text [START] / [STEP] / [END] stdout format.
-"""
-
 import os
-import sys
 import json
-import urllib.request
 from openai import OpenAI
+from env import EmailSortingEnv
 
 # ============================================
-# SETUP
+# 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",     "")
-SPACE_URL    = os.environ.get("SPACE_URL",    "http://localhost:7860")
+MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
+# Initialize OpenAI client
 client = OpenAI(
     base_url=API_BASE_URL,
     api_key=HF_TOKEN if HF_TOKEN else "dummy-key"
 )
 
 # ============================================
-# ENV HELPERS
+# AI AGENT — Asks LLM to classify email
 # ============================================
 
-def env_reset():
-    req = urllib.request.Request(f"{SPACE_URL}/reset", method="POST")
-    req.add_header("Content-Type", "application/json")
-    with urllib.request.urlopen(req, data=b"{}") as r:
-        return json.loads(r.read())
+def ask_llm_to_classify(email: dict) -> str:
+    """
+    Send email to LLM and get classification.
+    Returns: 'spam', 'important', or 'promotion'
+    """
 
-def env_step(action: str):
-    payload = json.dumps({"action": action}).encode()
-    req = urllib.request.Request(f"{SPACE_URL}/step", method="POST")
-    req.add_header("Content-Type", "application/json")
-    with urllib.request.urlopen(req, data=payload) as r:
-        return json.loads(r.read())
+    prompt = f"""You are an email classification assistant.
 
-# ============================================
-# LLM CLASSIFIER
-# ============================================
+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
+- 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."""
 
-def classify_email(subject: str, body: str) -> str:
-    prompt = (
-        f"Classify this email into exactly one category: spam, important, or promotion.\n\n"
-        f"Subject: {subject}\nBody: {body}\n\n"
-        f"Reply with ONE word only: spam, important, or promotion."
-    )
     try:
-        resp = client.chat.completions.create(
+        response = client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
-                {"role": "system", "content": "You are an email classifier. Reply with only one word."},
-                {"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=5,
+            max_tokens=10,
             temperature=0.0
         )
-        answer = resp.choices[0].message.content.strip().lower()
-        for label in ("spam", "promotion", "important"):
-            if label in answer:
-                return label
+
+        # 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
+
+    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"]
+
+    promo_keywords = ["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)
+
+    suspicious_domain = any(d in sender for d in [".xyz", ".tk", "-secure", "-alert"])
+
+    if spam_score >= 2 or suspicious_domain:
         return "spam"
-    except Exception:
-        return fallback_classify(subject, body)
+    elif promo_score >= 2:
+        return "promotion"
+    else:
+        return "important"
 
-def fallback_classify(subject: str, body: str) -> str:
-    text = (subject + " " + body).lower()
-    spam_kw  = ["won", "free", "prize", "claim", "urgent", "congratulations", "selected"]
-    promo_kw = ["off", "sale", "deal", "discount", "offer", "save", "shop"]
-    if sum(1 for k in spam_kw  if k in text) >= 2: return "spam"
-    if sum(1 for k in promo_kw if k in text) >= 2: return "promotion"
-    return "important"
 
 # ============================================
-# MAIN
+# MAIN INFERENCE LOOP
 # ============================================
 
 def run_inference():
-    task_name = "email-sorting-openenv"
-
-    # [START]
-    print(f"[START] task={task_name} model={MODEL_NAME}", flush=True)
-
-    result   = env_reset()
-    state    = result.get("state", result)
-    total_reward = 0.0
-    step_num = 0
-
-    while not state.get("done", False):
-        step_num += 1
-        email   = state.get("email", {})
-        subject = email.get("subject", "")
-        body    = email.get("body",    "")
-
-        action = classify_email(subject, body)
-        result = env_step(action)
-
-        reward       = result.get("reward", 0.0)
-        total_reward += reward
-        correct      = result.get("info", {}).get("correct", "")
-        state        = result.get("state", result)
-
-        # [STEP]
-        print(
-            f"[STEP] step={step_num} action={action} correct={correct} "
-            f"reward={reward} total_reward={total_reward}",
-            flush=True
-        )
+    """
+    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()
+    state = env.reset()
+
+    print(f"\nStarting episode — max {state['max_steps']} steps\n")
+
+    step_results = []
+
+    # Run until episode is done
+    while not state["done"]:
+        current_step = state["step"] + 1
+        email = state["email"]
+
+        print(f"Step {current_step}/{state['max_steps']}")
+        print(f"Subject: {email['subject']}")
+        print(f"From:    {email['sender']}")
+
+        # Ask AI to classify
+        action = ask_llm_to_classify(email)
+        print(f"AI Decision: {action}")
+
+        # Take step in environment
+        next_state, reward, done, info = env.step(action)
+
+        print(f"Reward: {reward} | Result: {info.get('result', 'N/A')}")
+        print("-" * 40)
+
+        step_results.append({
+            "step": current_step,
+            "subject": email["subject"],
+            "action": action,
+            "reward": reward,
+            "result": info.get("result", "N/A")
+        })
+
+        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)
+
+    # Save results to file
+    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
+    }
+
+    with open("inference_results.json", "w") as f:
+        json.dump(results, f, indent=2)
+
+    print("\nResults saved to inference_results.json")
+    return results
+
+
+# ============================================
+# RUN GRADERS ALSO
+# ============================================
 
-    score = round(total_reward / step_num, 4) if step_num else 0.0
+def run_with_graders():
+    """Run inference + all graders and show combined score."""
+    from graders import run_all_graders
 
-    # [END]
-    print(
-        f"[END] task={task_name} score={score} "
-        f"total_reward={total_reward} steps={step_num}",
-        flush=True
-    )
+    print("\n--- Running Inference ---\n")
+    inference_results = run_inference()
+
+    print("\n--- Running Graders ---\n")
+    grader_results = run_all_graders()
+
+    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(f"All Graders Passed: {grader_results['all_passed']}")
+    print("=" * 50)
+
+
+# ============================================
+# ENTRY POINT
+# ============================================
 
 if __name__ == "__main__":
-    run_inference()
+    run_with_graders()