inference.py

"""
Inference Script for Email Triage Environment
- Uses OpenAI client for LLM calls
- Reads API credentials from environment variables
- Produces reproducible baseline scores on all 3 tasks
"""

import os
import json
import requests
from openai import OpenAI

# Required environment variables
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.3-70B-Instruct")
ENV_BASE_URL = os.getenv("ENV_BASE_URL", "https://niths13-email-triage-env.hf.space")

SYSTEM_PROMPT = """
You are an expert email triage assistant.
Given an email subject and body, you must:
1. Categorize it as one of: billing, technical, general, spam
2. Set priority as one of: high, medium, low
3. Write a short response subject line (for hard tasks)

Always respond in valid JSON format like this:
{
  "category": "billing",
  "priority": "high",
  "response_subject": "Re: Your billing inquiry"
}

Rules:
- billing: payment, invoice, charge, refund issues
- technical: app bugs, login issues, crashes
- general: questions, info requests
- spam: suspicious, prize, promotional
- high priority: urgent issues affecting service
- medium priority: important but not urgent
- low priority: general questions, spam
"""


def call_llm(client, email_subject: str, email_body: str, task_level: str) -> dict:
    """Call the LLM to triage an email."""
    user_prompt = f"""
Task Level: {task_level.upper()}
Email Subject: {email_subject}
Email Body: {email_body}

Respond with JSON only.
"""
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            temperature=0.1,
            max_tokens=200,
        )
        response_text = completion.choices[0].message.content or ""
        # Clean up response
        response_text = response_text.strip()
        if "```json" in response_text:
            response_text = response_text.split("```json")[1].split("```")[0]
        elif "```" in response_text:
            response_text = response_text.split("```")[1].split("```")[0]
        return json.loads(response_text)
    except Exception as e:
        print(f"LLM call failed: {e}")
        return {"category": "general", "priority": "low", "response_subject": "Re: Your email"}


def reset_env() -> dict:
    """Reset the environment."""
    response = requests.post(f"{ENV_BASE_URL}/reset", json={})
    return response.json()


def step_env(action: dict) -> dict:
    """Take a step in the environment."""
    response = requests.post(f"{ENV_BASE_URL}/step", json={"action": action})
    return response.json()


def run_episode(client, task_level: str = "easy") -> float:
    """Run one full episode and return total reward."""
    print(f"\n{'='*50}")
    print(f"Starting episode - Task: {task_level.upper()}")
    print(f"{'='*50}")

    # Reset environment
    obs = reset_env()
    print(f"Email: {obs.get('observation', {}).get('email_subject', '')}")
    print(f"Task Level: {obs.get('observation', {}).get('task_level', '')}")

    total_reward = 0.0
    max_steps = 5
    step = 0

    while not obs.get("done", False) and step < max_steps:
        observation = obs.get("observation", obs)
        email_subject = observation.get("email_subject", "")
        email_body = observation.get("email_body", "")
        current_level = observation.get("task_level", task_level)

        # Get LLM decision
        llm_response = call_llm(client, email_subject, email_body, current_level)
        print(f"\nStep {step + 1}:")
        print(f"  LLM Decision: {llm_response}")

        # Take action
        action = {
            "category": llm_response.get("category", "general"),
            "priority": llm_response.get("priority", "low"),
            "response_subject": llm_response.get("response_subject", ""),
            "metadata": {}
        }

        obs = step_env(action)
        reward = obs.get("reward", 0.0) or 0.0
        total_reward += reward
        print(f"  Reward: {reward:.2f}")
        print(f"  Message: {obs.get('observation', {}).get('message', '')}")

        if obs.get("done"):
            print("Episode complete!")
            break
        step += 1

    print(f"\nTotal Reward: {total_reward:.2f}")
    return total_reward


def main():
    print("Email Triage Environment - Baseline Inference Script")
    print("=" * 50)

    # Initialize OpenAI client
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    # Run 3 episodes and report scores
    scores = []
    for i in range(3):
        print(f"\nEpisode {i + 1}/3")
        score = run_episode(client)
        scores.append(score)

    print("\n" + "=" * 50)
    print("BASELINE SCORES:")
    for i, score in enumerate(scores):
        print(f"  Episode {i + 1}: {score:.2f}")
    print(f"  Average Score: {sum(scores) / len(scores):.2f}")
    print("=" * 50)


if __name__ == "__main__":
    main()