ashucode/metaxhuggingfacehackathon

1from __future__ import annotations2 3import argparse4import os5import re6import subprocess7import sys8import time9from typing import Any, Dict, Optional10 11import requests12from openai import OpenAI13 14 15DEFAULT_BASE_URL = "http://localhost:7860"16ACTION_NAMES = ["UP", "DOWN", "LEFT", "RIGHT", "COLLECT"]17ACTION_NAME_TO_ID = {name: idx for idx, name in enumerate(ACTION_NAMES)}18 19SYSTEM_PROMPT = """You are an expert shopping agent navigating a grid supermarket.20GRID LAYOUT (8 rows x 7 cols):21  - Spawn:    row 0, col 322  - Counter:  row 7, col 3  (EXIT -- go here after collecting all items)23  - Aisles:   A1(1,1) A2(3,1) A3(5,1) A4(1,5) A5(3,5) A6(5,5)24ACTIONS (reply with EXACTLY the action name -- nothing else):25  UP      move one row up    (row - 1)26  DOWN    move one row down  (row + 1)27  LEFT    move one col left  (col - 1)28  RIGHT   move one col right (col + 1)29  COLLECT pick up item at current aisle (only when standing on an aisle that has a target)30RULES:31  1. In Phase 0: navigate to each target aisle and COLLECT all required items.32  2. In Phase 1: return to the Counter at (7, 3) to complete the task.33  3. COLLECT is only valid when action_mask[4] is true (you are on a target aisle).34  4. On HARD level you MUST collect the CLOSEST item first.35  5. Every step costs reward; reach the counter as fast as possible.36Reply with ONE word -- the action name. No explanation, no punctuation."""37 38 39def build_user_message(state: Dict[str, Any]) -> str:40    mask_str = ", ".join(41        f"{name}={'OK' if ok else 'BLOCKED'}"42        for name, ok in zip(ACTION_NAMES, state["action_mask"])43    )44    lines = [45        f"Position : {state['agent_pos']}",46        f"Phase    : {state['phase']} ({state['phase_label']})",47        f"Targets  : {state['targets']}",48        f"Inventory: {state['inventory']}",49        f"Remaining: {[t for t in state['targets'] if t not in state['inventory']]}",50        f"Closest  : {state['closest_target']}",51        f"Steps    : {state['steps_taken']} / {state['steps_taken'] + state['steps_remaining']}",52        f"Total rwd: {state['total_reward']:.2f}",53        f"Masks    : {mask_str}",54        f"Last evt : {state.get('event', '')}",55    ]56    return "\n".join(lines)57 58 59def parse_action(text: str) -> Optional[int]:60    clean = text.strip().upper()61    if clean in ACTION_NAME_TO_ID:62        return ACTION_NAME_TO_ID[clean]63    for name, idx in ACTION_NAME_TO_ID.items():64        if re.search(rf"\b{name}\b", clean):65            return idx66    return None67 68 69def log_start(session_id: str, level: str, model: str):70    print(f"[START] session_id={session_id} level={level} model={model}", flush=True)71 72 73def log_step(step_num: int, action_name: str, state: Dict[str, Any]):74    print(75        f"[STEP] step={step_num} action={action_name} "76        f"step_reward={state.get('step_reward', 0.0):.4f} "77        f"score={state.get('normalised_score', 0.0):.4f}",78        flush=True,79    )80 81 82def log_end(session_id: str, state: Dict[str, Any]):83    print(84        f"[END] session_id={session_id} status={state.get('task_status', 'Unknown')} "85        f"steps={state.get('steps_taken', 0)} "86        f"final_score={state.get('normalised_score', 0.0):.4f}",87        flush=True,88    )89 90 91def ensure_server_running(base_url: str) -> None:92    """Start uvicorn server if not already reachable."""93    try:94        r = requests.get(f"{base_url}/healthz", timeout=3)95        if r.status_code == 200:96            print("[INFO] Server already running.", flush=True)97            return98    except requests.RequestException:99        pass100 101    print("[INFO] Starting uvicorn server...", flush=True)102    subprocess.Popen(103        [104            sys.executable, "-m", "uvicorn",105            "app:app",106            "--host", "0.0.0.0",107            "--port", "7860",108            "--workers", "1",109        ],110        stdout=subprocess.DEVNULL,111        stderr=subprocess.DEVNULL,112    )113    for _ in range(30):114        time.sleep(1)115        try:116            r = requests.get(f"{base_url}/healthz", timeout=2)117            if r.status_code == 200:118                print("[INFO] Server is up.", flush=True)119                return120        except requests.RequestException:121            continue122    raise RuntimeError("Server did not become reachable within 30 seconds.")123 124 125def run_agent(level: str, base_url: str, products=None, seed=None) -> float:126    ensure_server_running(base_url)127 128    api_base_url  = os.environ.get("API_BASE_URL", "").strip()129    model         = os.environ.get("MODEL_NAME", "gpt-4o-mini").strip()130    hf_token      = os.environ.get("HF_TOKEN", "").strip()131    effective_key = hf_token if hf_token else "placeholder-key"132 133    try:134        client = OpenAI(135            api_key=effective_key,136            base_url=api_base_url if api_base_url else None,137        )138    except Exception as exc:139        print(f"[ERROR] Failed to create OpenAI client: {exc}", flush=True)140        sys.exit(1)141 142    reset_payload: Dict[str, Any] = {"level": level}143    if products:144        reset_payload["products"] = products145    if seed is not None:146        reset_payload["seed"] = seed147 148    try:149        r = requests.post(f"{base_url}/reset", json=reset_payload, timeout=10)150        r.raise_for_status()151    except requests.RequestException as exc:152        print(f"[ERROR] Could not reach server at {base_url}: {exc}", flush=True)153        sys.exit(1)154 155    state      = r.json()156    session_id = state["session_id"]157    step_num   = 0158    conversation: list = []159 160    log_start(session_id, level, model)161 162    while (163        not state.get("terminated", False)164        and not state.get("truncated", False)165        and not state.get("done", False)166        and state.get("task_status") == "In-Progress"167    ):168        user_msg = build_user_message(state)169        conversation.append({"role": "user", "content": user_msg})170        trimmed_history = conversation[-20:]171 172        try:173            completion = client.chat.completions.create(174                model=model,175                messages=[{"role": "system", "content": SYSTEM_PROMPT}] + trimmed_history,176                max_tokens=10,177                temperature=0.0,178            )179            reply = completion.choices[0].message.content or ""180        except Exception as exc:181            print(f"[ERROR] LLM call failed: {exc}", flush=True)182            break183 184        action_id = parse_action(reply)185        if action_id is None:186            action_id = 1187 188        action_name = ACTION_NAMES[action_id]189        conversation.append({"role": "assistant", "content": action_name})190 191        try:192            r = requests.post(193                f"{base_url}/step/{session_id}",194                json={"action": action_id},195                timeout=10,196            )197            r.raise_for_status()198        except requests.RequestException as exc:199            print(f"[ERROR] Step request failed: {exc}", flush=True)200            break201 202        state    = r.json()203        step_num += 1204        log_step(step_num, action_name, state)205 206    log_end(session_id, state)207 208    # Clamp strictly between 0 and 1 as required by the evaluator209    raw_score = state.get("normalised_score", 0.0)210    return max(0.001, min(0.999, raw_score))211 212 213if __name__ == "__main__":214    parser = argparse.ArgumentParser(description="LLM agent for SupermarketNav")215    parser.add_argument("--level", default="easy", choices=["easy", "medium", "hard"])216    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, dest="base_url")217    parser.add_argument("--products", nargs="*", help="Optional product names")218    parser.add_argument("--seed", type=int, default=None)219    args = parser.parse_args()220 221    try:222        score = run_agent(223            level=args.level,224            base_url=args.base_url,225            products=args.products,226            seed=args.seed,227        )228        print(f"[DONE] score={score:.4f}", flush=True)229    except Exception as exc:230        print(f"[FATAL] {exc}", flush=True)231        sys.exit(1)232 233    sys.exit(0)