inference.py

"""Hackathon inference loop for the EmailTriage OpenEnv environment.

Runs all 3 tasks (easy, medium, hard) sequentially using the OpenAI client.
Emits structured [START]/[STEP]/[END] logs per the hackathon spec.
"""

import os
import json
import time
import asyncio
from urllib import error as urllib_error
from urllib import parse as urllib_parse
from urllib import request as urllib_request
from typing import List, Optional

from openai import OpenAI

try:
    from EmailTriage import EmailtriageAction, EmailtriageEnv
    _IMPORT_OK = True
    _IMPORT_ERROR = ""
except Exception as _import_err:
    _IMPORT_OK = False
    _IMPORT_ERROR = str(_import_err)

API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
ENV_BASE_URL = os.getenv("ENV_BASE_URL") or "http://localhost:8000"
LLM_API_BASE_URL = os.getenv("LLM_API_BASE_URL") or API_BASE_URL
MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
API_KEY = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")
BENCHMARK_NAME = "openenv-emailtriage"
REQUEST_TIMEOUT_SECONDS = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "8"))
REQUEST_MAX_RETRIES = int(os.getenv("REQUEST_MAX_RETRIES", "3"))
MAX_RUNTIME_SECONDS = int(os.getenv("INFERENCE_TIMEOUT_SECONDS", "1100"))

TASK_IDS = ["easy", "medium", "hard"]

# Per-task step budgets (must fit within 20min total runtime)
TASK_MAX_STEPS = {
    "easy": 6,
    "medium": 10,
    "hard": 12,
}

SPAM_HINTS = {
    "newsletter",
    "promo",
    "promotion",
    "sale",
    "discount",
    "offer",
    "webinar",
    "subscribe",
    "unsubscribe",
    "nigerian prince",
    "lottery",
    "free",
}

SCHEDULING_HINTS = {
    "schedule",
    "meeting",
    "calendar",
    "slot",
    "reschedule",
    "availability",
    "call",
}


# ---------------------------------------------------------------------------
# Structured stdout logging (hackathon spec)
# ---------------------------------------------------------------------------


def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(
    step: int,
    action: str,
    reward: float,
    done: bool,
    error: Optional[str],
) -> None:
    error_value = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} error={error_value}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{value:.2f}" for value in rewards)
    print(
        f"[END] success={str(success).lower()} "
        f"steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


def _normalize_base_url(base_url: str) -> str:
    return base_url.rstrip("/")


def _safe_request_json(
    method: str,
    base_url: str,
    path: str,
    payload: Optional[dict] = None,
    timeout: float = REQUEST_TIMEOUT_SECONDS,
    retries: int = REQUEST_MAX_RETRIES,
) -> tuple[bool, Optional[dict], str]:
    """Safely call endpoint with retries and JSON validation."""
    url = f"{_normalize_base_url(base_url)}{path}"
    body = None
    headers = {"Accept": "application/json"}
    if payload is not None:
        body = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"

    for attempt in range(1, retries + 1):
        req = urllib_request.Request(
            url,
            data=body,
            method=method.upper(),
            headers=headers,
        )
        try:
            with urllib_request.urlopen(req, timeout=timeout) as resp:
                status = int(getattr(resp, "status", 0) or 0)
                raw = resp.read().decode("utf-8", errors="replace")
                if status != 200:
                    return False, None, f"HTTP {status} from {path}"
                try:
                    parsed = json.loads(raw) if raw else {}
                except json.JSONDecodeError:
                    return False, None, f"Invalid JSON from {path}"
                if not isinstance(parsed, dict):
                    return False, None, f"Non-object JSON from {path}"
                return True, parsed, ""
        except (urllib_error.URLError, TimeoutError, ValueError) as exc:
            if attempt == retries:
                return False, None, f"{path} failed after {retries} attempts: {exc}"
            time.sleep(min(0.4 * attempt, 1.0))

    return False, None, f"{path} failed"


def preflight_env_endpoints(base_url: str) -> tuple[bool, str]:
    """Validate reset/step/state endpoints before creating the env client."""
    ok, reset_data, err = _safe_request_json(
        "POST",
        base_url,
        "/reset",
        payload={"task_id": TASK_IDS[0]},
    )
    if not ok:
        return False, err

    if "observation" not in (reset_data or {}):
        return False, "/reset missing observation"

    ok, _, err = _safe_request_json("GET", base_url, "/state")
    if not ok:
        return False, err

    # Minimal valid action envelope for compatibility check.
    ok, _, err = _safe_request_json(
        "POST",
        base_url,
        "/step",
        payload={
            "action": {
                "action_type": "query_calendar",
                "target_email_id": -1,
                "draft_content": "",
                "proposed_slot": "",
            }
        },
    )
    if not ok:
        return False, err

    return True, ""


# ---------------------------------------------------------------------------
# Prompt construction
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = (
    "You are an elite, proactive email triage assistant operating in a strictly structured environment. "
    "Your goal is to process the entire inbox efficiently, maximizing your rewards.\n"
    "CRITICAL RULES FOR STATE ADVANCEMENT:\n"
    "1. AVOID LOOPS: Check the 'Last action result' and 'Recently read emails'. If you just read an email, DO NOT read it again. You must take the next logical step (archive or draft_email).\n"
    "2. SPAM/NEWSLETTERS: If an unread email subject from the 'Inbox preview' clearly looks like spam, marketing, or a low-priority notification, immediately use action_type='archive'.\n"
    "3. IMPORTANT EMAILS: If an unread email is a client request, meeting, or escalation, use action_type='read' first to get the full text.\n"
    "4. RESPONDING: If 'Recently read emails' contains a client email that needs a reply, immediately use action_type='draft_email'. "
    "Your draft_content MUST be professional, mention 'thank', reference specific details from the subject, end firmly with a period, and be over 40 characters.\n"
    "5. SCHEDULING CALENDAR: If a read email asks for a meeting, first use action_type='query_calendar' (target_email_id=-1) to load availability. "
    "In your VERY NEXT turn, use action_type='draft_email' and provide one of the listed slots exactly as shown in the 'proposed_slot' field.\n"
    "6. JSON FORMAT: Respond ONLY with valid JSON. Keys required: action_type, target_email_id, draft_content, proposed_slot. No markdown, no conversational text."
)


def build_user_prompt(
    task_id: str,
    inbox_preview: List[dict],
    returned_emails: List[str],
    calendar_slots: List[str],
    last_action_result: str,
) -> str:
    slots = ", ".join(calendar_slots) if calendar_slots else "none"
    inbox_lines = [
        f"id={item.get('id')} sender={item.get('sender')} "
        f"priority={item.get('priority')} subject={item.get('subject')}"
        for item in inbox_preview
    ]
    inbox_block = (
        " | ".join(inbox_lines) if inbox_lines else "no unread emails"
    )
    reads_block = " | ".join(returned_emails) if returned_emails else "none"

    return (
        f"Task difficulty: {task_id}. "
        f"Inbox preview: {inbox_block}. "
        f"Recently read emails: {reads_block}. "
        f"Calendar slots: {slots}. "
        f"Last action result: {last_action_result}."
    )


# ---------------------------------------------------------------------------
# LLM action selection
# ---------------------------------------------------------------------------


def choose_action_with_llm(
    client: OpenAI,
    task_id: str,
    prompt: str,
) -> "EmailtriageAction":
    default_action = EmailtriageAction(
        action_type="query_calendar",
        target_email_id=-1,
        draft_content="",
        proposed_slot="",
    )

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.2,
            max_tokens=200,
            stream=False,
        )
        raw_content = (completion.choices[0].message.content or "").strip()
        if not raw_content:
            return default_action

        # Strip markdown fences if the model wraps JSON
        if raw_content.startswith("```"):
            lines = raw_content.split("\n")
            lines = [l for l in lines if not l.strip().startswith("```")]
            raw_content = "\n".join(lines)

        # Recover JSON object when model emits extra text around it.
        if "{" in raw_content and "}" in raw_content:
            start = raw_content.find("{")
            end = raw_content.rfind("}") + 1
            raw_content = raw_content[start:end]

        data = json.loads(raw_content)
        return EmailtriageAction(
            action_type=data.get("action_type", "query_calendar"),
            target_email_id=int(data.get("target_email_id", -1)),
            draft_content=data.get("draft_content", ""),
            proposed_slot=data.get("proposed_slot", ""),
        )
    except Exception:
        return default_action


def _is_spam_like(subject: str, sender: str) -> bool:
    text = f"{subject} {sender}".lower()
    return any(hint in text for hint in SPAM_HINTS)


def _needs_scheduling(read_emails: List[str]) -> bool:
    joined = " ".join(read_emails).lower()
    return any(hint in joined for hint in SCHEDULING_HINTS)


def _build_draft(read_emails: List[str], slot: str) -> str:
    if slot:
        return (
            "Thank you for your email. I reviewed your request and can confirm "
            f"{slot} works on our side. Please confirm and I will send the invite."
        )
    return (
        "Thank you for the detailed context. I reviewed your request and will "
        "follow up with the next concrete steps shortly."
    )


def choose_action_with_fallback(
    llm_action: "EmailtriageAction",
    inbox_preview: List[dict],
    returned_emails: List[str],
    calendar_slots: List[str],
    recent_actions: List[str],
    last_read_email_id: int,
    has_queried_calendar: bool,
    closed_email_ids: set,
) -> "EmailtriageAction":
    valid_ids = {
        int(item.get("id"))
        for item in inbox_preview
        if item.get("id") is not None
    }

    repeated_calendar_loop = (
        len(recent_actions) >= 1
        and recent_actions[-1] == "query_calendar"
        and llm_action.action_type == "query_calendar"
    )

    llm_invalid = (
        llm_action.action_type not in {"read", "archive", "query_calendar", "draft_email"}
        or (
            llm_action.action_type in {"read", "archive"}
            and llm_action.target_email_id not in valid_ids
        )
    )

    # Prefer drafting immediately after reading if we have content.
    if last_read_email_id != -1 and returned_emails:
        needs_slot = _needs_scheduling(returned_emails)
        if needs_slot and (not has_queried_calendar):
            return EmailtriageAction(
                action_type="query_calendar",
                target_email_id=-1,
                draft_content="",
                proposed_slot="",
            )

        slot = calendar_slots[0] if needs_slot and calendar_slots else ""
        return EmailtriageAction(
            action_type="draft_email",
            target_email_id=last_read_email_id,
            draft_content=_build_draft(returned_emails, slot),
            proposed_slot=slot,
        )

    # Only allow query_calendar when it contributes to scheduling.
    if llm_action.action_type == "query_calendar":
        if has_queried_calendar and not returned_emails:
            repeated_calendar_loop = True

    if not repeated_calendar_loop and not llm_invalid:
        if llm_action.action_type in {"archive", "draft_email"}:
            if llm_action.target_email_id in closed_email_ids:
                llm_invalid = True
            else:
                return llm_action
        elif llm_action.action_type == "read":
            if llm_action.target_email_id in closed_email_ids:
                llm_invalid = True
            else:
                return llm_action
        else:
            return llm_action

    # Archive obvious spam/newsletters.
    for item in inbox_preview:
        email_id = item.get("id")
        subject = str(item.get("subject", ""))
        sender = str(item.get("sender", ""))
        if (
            email_id is not None
            and int(email_id) not in closed_email_ids
            and _is_spam_like(subject, sender)
        ):
            return EmailtriageAction(
                action_type="archive",
                target_email_id=int(email_id),
                draft_content="",
                proposed_slot="",
            )

    # Otherwise read the highest-priority available message.
    priority_rank = {"high": 0, "medium": 1, "low": 2}
    if inbox_preview:
        sorted_preview = sorted(
            [
                item
                for item in inbox_preview
                if int(item.get("id", -1)) not in closed_email_ids
            ]
            or inbox_preview,
            key=lambda item: priority_rank.get(
                str(item.get("priority", "low")).lower(), 3
            ),
        )
        pick_id = sorted_preview[0].get("id", -1)
        return EmailtriageAction(
            action_type="read",
            target_email_id=int(pick_id) if pick_id is not None else -1,
            draft_content="",
            proposed_slot="",
        )

    return EmailtriageAction(
        action_type="query_calendar",
        target_email_id=-1,
        draft_content="",
        proposed_slot="",
    )


# ---------------------------------------------------------------------------
# Single-task runner
# ---------------------------------------------------------------------------


async def run_task(
    llm_client: OpenAI,
    env: "EmailtriageEnv",
    task_id: str,
    start_time: float,
) -> None:
    """Run a single task (easy/medium/hard) and emit structured logs."""
    max_steps = TASK_MAX_STEPS[task_id]
    task_name = f"email-triage-{task_id}"
    rewards: List[float] = []
    steps_taken = 0
    score = 0.01
    success = False
    recent_actions: List[str] = []
    last_read_email_id = -1
    has_queried_calendar = False
    closed_email_ids: set = set()

    log_start(task=task_name, env=BENCHMARK_NAME, model=MODEL_NAME)

    try:
        try:
            result = await env.reset(options={"task_id": task_id})
        except Exception as exc:
            log_step(
                step=1,
                action="reset()",
                reward=0.0,
                done=True,
                error=str(exc),
            )
            return

        for step in range(1, max_steps + 1):
            elapsed = time.time() - start_time
            if elapsed >= MAX_RUNTIME_SECONDS:
                log_step(
                    step=step,
                    action="timeout_guard",
                    reward=0.0,
                    done=True,
                    error="runtime limit reached",
                )
                break

            obs = result.observation
            if result.done or obs.inbox_remaining <= 0:
                break

            prompt = build_user_prompt(
                task_id=task_id,
                inbox_preview=obs.inbox_preview,
                returned_emails=obs.returned_emails,
                calendar_slots=obs.calendar_slots,
                last_action_result=obs.last_action_result,
            )
            action = choose_action_with_llm(llm_client, task_id, prompt)
            action = choose_action_with_fallback(
                llm_action=action,
                inbox_preview=obs.inbox_preview,
                returned_emails=obs.returned_emails,
                calendar_slots=obs.calendar_slots,
                recent_actions=recent_actions,
                last_read_email_id=last_read_email_id,
                has_queried_calendar=has_queried_calendar,
                closed_email_ids=closed_email_ids,
            )

            if action.action_type == "read":
                last_read_email_id = action.target_email_id
            elif action.action_type == "draft_email":
                closed_email_ids.add(action.target_email_id)
                last_read_email_id = -1
            elif action.action_type == "archive":
                closed_email_ids.add(action.target_email_id)

            if action.action_type == "query_calendar":
                has_queried_calendar = True

            recent_actions.append(action.action_type)
            if len(recent_actions) > 6:
                recent_actions.pop(0)

            try:
                result = await env.step(action)
            except Exception as exc:
                log_step(
                    step=step,
                    action="env.step()",
                    reward=0.0,
                    done=True,
                    error=str(exc),
                )
                break

            reward = float(result.reward or 0.0)
            rewards.append(reward)
            steps_taken = step

            action_str = (
                f"{action.action_type}("
                f"target_email_id={action.target_email_id},"
                f"proposed_slot={action.proposed_slot})"
            )
            log_step(
                step=step,
                action=action_str,
                reward=reward,
                done=bool(result.done),
                error=None,
            )

            if result.done:
                break

        score = sum(rewards) / len(rewards) if rewards else 0.01
        score = min(max(score, 0.01), 0.99)
        # Success threshold is 0.5 avg reward
        success = score >= 0.5

    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


async def main() -> None:
    start_time = time.time()

    if not _IMPORT_OK:
        for task_id in TASK_IDS:
            log_start(task=f"email-triage-{task_id}", env=BENCHMARK_NAME, model=MODEL_NAME)
            log_step(1, "import", 0.0, True, error=_IMPORT_ERROR)
            log_end(False, 1, 0.01, [0.01])
        return

    try:
        # Environment variable safety checks (do not crash validator).
        if not API_BASE_URL:
            for task_id in TASK_IDS:
                log_start(task=f"email-triage-{task_id}", env=BENCHMARK_NAME, model=MODEL_NAME)
                log_step(
                    step=1,
                    action="preflight",
                    reward=0.0,
                    done=True,
                    error="API_BASE_URL is missing",
                )
                log_end(success=False, steps=1, score=0.01, rewards=[0.01])
            return

        if not API_KEY:
            for task_id in TASK_IDS:
                log_start(task=f"email-triage-{task_id}", env=BENCHMARK_NAME, model=MODEL_NAME)
                log_step(
                    step=1,
                    action="preflight",
                    reward=0.0,
                    done=True,
                    error="HF_TOKEN/OPENAI_API_KEY is missing",
                )
                log_end(success=False, steps=1, score=0.01, rewards=[0.01])
            return

        ok, error_message = preflight_env_endpoints(ENV_BASE_URL)
        if not ok:
            for task_id in TASK_IDS:
                log_start(task=f"email-triage-{task_id}", env=BENCHMARK_NAME, model=MODEL_NAME)
                log_step(
                    step=1,
                    action="endpoint_preflight",
                    reward=0.0,
                    done=True,
                    error=error_message,
                )
                log_end(success=False, steps=1, score=0.01, rewards=[0.01])
            return

        llm_client = OpenAI(base_url=LLM_API_BASE_URL, api_key=API_KEY)
        env = EmailtriageEnv(base_url=ENV_BASE_URL)

        try:
            for task_id in TASK_IDS:
                await run_task(llm_client, env, task_id, start_time)
                if time.time() - start_time >= MAX_RUNTIME_SECONDS:
                    break
        except Exception:
            # Keep validator-safe behavior: no crash propagation.
            pass
        finally:
            try:
                await env.close()
            except Exception:
                pass

    except Exception:
        # Catch anything from preflight or env setup
        pass


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except BaseException:
        # Ensure sandbox validator always receives exit code 0.
        pass