Spaces:
Sleeping
Sleeping
| """ | |
| Inference script for the Store Manager environment. | |
| Runs an LLM agent through all three difficulty tasks (easy → medium → hard) | |
| and reports a deterministic grade (0.0–1.0) per task. | |
| Usage: | |
| python inference.py | |
| python inference.py --url http://localhost:8000 | |
| Required environment variables: | |
| HF_TOKEN API key for the LLM provider (Hugging Face token or OpenAI key) | |
| API_BASE_URL Base URL of the LLM API (e.g. https://api-inference.huggingface.co/v1/) | |
| MODEL_NAME Model identifier (e.g. meta-llama/Llama-3.1-8B-Instruct) | |
| Optional environment variables: | |
| STORE_ENV_URL Base URL of the running Store Manager server (default: http://localhost:8000) | |
| Structured log format emitted to stdout: | |
| [START] {"task": "...", "seed": N, "model": "...", "max_steps": N, "profit_target": N} | |
| [STEP] {"task": "...", "step": N, "action": {...}, "reward": ..., "cumulative_profit": ...} | |
| [END] {"task": "...", "cumulative_profit": ..., "score": ..., "profit_target": ...} | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import time | |
| try: | |
| import websocket # websocket-client library | |
| except ImportError: | |
| print( | |
| "ERROR: websocket-client not installed. Run: pip install websocket-client", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(1) | |
| try: | |
| from openai import OpenAI | |
| except ImportError: | |
| print( | |
| "ERROR: openai not installed. Run: pip install openai", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(1) | |
| # Add project root to path so tasks.py is importable | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from tasks import TASKS, grade | |
| # ── Prompt helpers ───────────────────────────────────────────────────────────── | |
| _SYSTEM_PROMPT = """You are an expert retail store manager optimizing profit. | |
| Each step you choose EXACTLY ONE action. There are three action types: | |
| ── 1. DISCOUNT ────────────────────────────────────────────────────────────────── | |
| Apply a price discount to a product for this step only (ephemeral). | |
| Discounts boost customer pick probability and clear near-expiry stock. | |
| {"action_type": "discount", "product_id": <int>, "discount_pct": <0|10|20|50>} | |
| Discount tiers and effect: | |
| 0% -> no change | 10% -> +10% relative boost | 20% -> +20% boost | 50% -> +50% boost | |
| PRICE FLOOR RULES (selling below cost is penalized unless near expiry): | |
| In(steps) >= 4 -> floor = cost_price (no below-cost selling allowed) | |
| In(steps) == 3 -> floor = 70% of cost (max 30% loss/unit) | |
| In(steps) == 2 -> floor = 40% of cost (max 60% loss/unit) | |
| In(steps) == 1 -> floor = 10% of cost (max 90% loss/unit — clear it!) | |
| Selling below cost when In(steps) >= 4 incurs an EXTRA penalty equal to the | |
| per-unit loss × units sold. Avoid below-cost discounts unless near expiry. | |
| When to discount: near-expiry items (avoid expiry penalty = remaining_qty x cost), | |
| high-margin products (Margin > $0) to drive volume, or oversupplied items. | |
| Near expiry (In <= 3): aggressive discounts are worthwhile to avoid the full expiry penalty. | |
| ── 2. RESTOCK ─────────────────────────────────────────────────────────────────── | |
| Place a reorder for a product. Stock arrives after restock_lead_time steps. | |
| Costs restock_cost_per_unit per unit (upfront, this step). One order per product | |
| at a time -- wait for delivery before reordering the same product. | |
| {"action_type": "restock", "product_id": <int>, "restock_quantity": <int>} | |
| When to restock: low stock on a high-demand product, enough steps left before expiry | |
| to sell through the restocked units, and available budget. | |
| ── 3. PLACE ───────────────────────────────────────────────────────────────────── | |
| Move a product to a different shelf zone. Costs a flat placement fee ($3.00). | |
| {"action_type": "place", "product_id": <int>, "target_zone": <1|2|3>} | |
| Zone multipliers on pick probability: | |
| Zone 1 (premium/entrance) -> x1.5 [capacity: 2 products max] | |
| Zone 2 (standard) -> x1.0 [capacity: 4 products max] | |
| Zone 3 (back/low-traffic) -> x0.6 [unlimited] | |
| When to place: move high-value or near-expiry products to zone 1 to increase | |
| their pick probability. Move slow-movers to zone 3. | |
| ── COSTS CHARGED EVERY STEP ───────────────────────────────────────────────────── | |
| - Holding cost: each unit in stock costs holding_cost_rate x cost_price per step. | |
| Minimize idle inventory by balancing restocks with sales velocity. | |
| - Expiry penalty: remaining_qty x cost_price for each expired product. | |
| Goal: maximize cumulative profit = revenue - COGS - expiry_penalty - holding_cost | |
| - placement_cost - restock_cost. | |
| Respond ONLY with a single JSON object (one of the three formats above). | |
| No explanation, no extra text -- just the JSON. | |
| """ | |
| def _format_inventory(obs: dict) -> str: | |
| """Format inventory state as a readable table for the LLM prompt.""" | |
| # JSON serializes int dict keys as strings; coerce back to int for lookup | |
| zm_raw = obs.get("zone_multipliers", {}) | |
| zc_raw = obs.get("zone_capacity", {}) | |
| zo_raw = obs.get("zone_occupancy", {}) | |
| zm = {int(k): v for k, v in zm_raw.items()} if zm_raw else {1: 1.5, 2: 1.0, 3: 0.6} | |
| zc = {int(k): v for k, v in zc_raw.items()} if zc_raw else {1: 2, 2: 4, 3: 99} | |
| zo = {int(k): v for k, v in zo_raw.items()} if zo_raw else {1: 0, 2: 0, 3: 0} | |
| lines = [ | |
| f"Step {obs['current_step'] + 1}/{obs['max_steps']} " | |
| f"| Steps remaining: {obs['steps_remaining']} " | |
| f"| Cumulative profit: ${obs['cumulative_profit']:.2f}", | |
| f"Task difficulty: {obs.get('task_name', 'unknown')}", | |
| f"Holding cost this step: ${obs.get('holding_cost_this_step', 0.0):.4f}", | |
| "", | |
| "Zone info: " | |
| + " ".join( | |
| f"Zone {z}: x{zm.get(z, 1.0)} ({zo.get(z, 0)}/{zc.get(z, '?')} slots)" | |
| for z in [1, 2, 3] | |
| ), | |
| "", | |
| f"{'ID':<4} {'Name':<10} {'Cost':>6} {'Price':>7} {'Margin':>7} {'Qty':>5} " | |
| f"{'In(steps)':>9} {'Zone':>5} {'PickProb':>9} {'PopMult':>8} {'HoldRate':>9} " | |
| f"{'RestockCost':>11} {'Lead':>5} {'PendingQty':>10} {'Active':>7}", | |
| "-" * 122, | |
| ] | |
| for p in obs["inventory"]: | |
| pending_str = ( | |
| f"{p['pending_restock_quantity']}@s{p['pending_restock_arrives_at_step']}" | |
| if p.get("pending_restock_quantity", 0) > 0 | |
| else "-" | |
| ) | |
| margin = p.get("margin_per_unit", round(p["selling_price"] - p["cost_price"], 4)) | |
| steps_left = p.get("steps_until_expiry", p.get("expiry_step", 0) - obs.get("current_step", 0)) | |
| lines.append( | |
| f"{p['product_id']:<4} {p['name']:<10} " | |
| f"${p['cost_price']:>5.2f} ${p['selling_price']:>6.2f} " | |
| f"${margin:>+6.2f} " | |
| f"{p['quantity']:>5} " | |
| f"{steps_left:>9} " | |
| f"{p.get('zone', 2):>5} " | |
| f"{p['effective_pick_prob']:>9.4f} " | |
| f"{p.get('popularity_multiplier', 1.0):>8.3f} " | |
| f"{p.get('holding_cost_rate', 0.0):>9.4f} " | |
| f"${p.get('restock_cost_per_unit', 0.0):>10.2f} " | |
| f"{p.get('restock_lead_time', 2):>5} " | |
| f"{pending_str:>10} " | |
| f"{'YES' if p['is_active'] else 'no':>7}" | |
| ) | |
| if obs.get("last_expired_products"): | |
| lines.append(f"\nJust expired: {', '.join(obs['last_expired_products'])}") | |
| return "\n".join(lines) | |
| def _parse_action(text: str, active_ids: list[int]) -> dict: | |
| """Extract a valid action from LLM output; fall back to safe default.""" | |
| def _validate(data: dict) -> dict | None: | |
| action_type = data.get("action_type", "discount") | |
| pid = int(data.get("product_id", active_ids[0])) | |
| if pid not in active_ids: | |
| return None | |
| if action_type == "discount": | |
| disc = int(data.get("discount_pct", 0)) | |
| if disc in (0, 10, 20, 50): | |
| return {"action_type": "discount", "product_id": pid, "discount_pct": disc} | |
| elif action_type == "restock": | |
| qty = data.get("restock_quantity") | |
| if qty is not None and int(qty) > 0: | |
| return {"action_type": "restock", "product_id": pid, "restock_quantity": int(qty)} | |
| elif action_type == "place": | |
| zone = data.get("target_zone") | |
| if zone is not None and int(zone) in (1, 2, 3): | |
| return {"action_type": "place", "product_id": pid, "target_zone": int(zone)} | |
| return None | |
| # Try direct JSON parse | |
| try: | |
| result = _validate(json.loads(text.strip())) | |
| if result: | |
| return result | |
| except (json.JSONDecodeError, ValueError, KeyError): | |
| pass | |
| # Try extracting JSON from surrounding text | |
| match = re.search(r'\{[^{}]+\}', text) | |
| if match: | |
| try: | |
| result = _validate(json.loads(match.group(0))) | |
| if result: | |
| return result | |
| except (json.JSONDecodeError, ValueError, KeyError): | |
| pass | |
| # Safe fallback | |
| return {"action_type": "discount", "product_id": active_ids[0], "discount_pct": 0} | |
| # ── WebSocket helpers ────────────────────────────────────────────────────────── | |
| def _ws_url(http_url: str) -> str: | |
| """Convert http(s) URL to ws(s).""" | |
| return http_url.replace("https://", "wss://").replace("http://", "ws://") + "/ws" | |
| def _recv_json(ws: websocket.WebSocket) -> dict: | |
| """Receive one WebSocket message and parse as JSON.""" | |
| return json.loads(ws.recv()) | |
| def _unpack(msg: dict) -> dict: | |
| """ | |
| Unpack a WebSocket message into a flat observation dict. | |
| The server sends: {"type": "observation", "data": {"observation": {...}, "done": bool, "reward": float}} | |
| This extracts the inner observation and injects done/reward at the top level. | |
| """ | |
| data = msg.get("data", msg) | |
| obs = data.get("observation", data) | |
| obs["done"] = data.get("done", obs.get("done", False)) | |
| obs["reward"] = data.get("reward", obs.get("reward", 0.0)) | |
| return obs | |
| # ── Episode runner ───────────────────────────────────────────────────────────── | |
| def run_episode( | |
| ws_url: str, | |
| task_name: str, | |
| client: OpenAI, | |
| model: str, | |
| ) -> tuple[float, float]: | |
| """ | |
| Run a full episode for `task_name` using the LLM agent. | |
| Returns: | |
| (cumulative_profit, score) | |
| """ | |
| task = TASKS[task_name] | |
| ws = None | |
| cumulative = 0.0 | |
| try: | |
| ws = websocket.create_connection(ws_url, timeout=60) | |
| # ── Reset ────────────────────────────────────────────────────────────── | |
| ws.send(json.dumps({"type": "reset", "data": {"task": task_name, "seed": task.seed}})) | |
| obs = _unpack(_recv_json(ws)) | |
| print( | |
| f"[START] {json.dumps({'task': task_name, 'seed': task.seed, 'model': model, 'max_steps': task.max_steps, 'profit_target': task.profit_target})}", | |
| flush=True, | |
| ) | |
| # ── Episode loop ─────────────────────────────────────────────────────── | |
| step_num = 0 | |
| while True: | |
| active_ids = [ | |
| p["product_id"] | |
| for p in obs.get("inventory", []) | |
| if p.get("is_active") and p.get("quantity", 0) > 0 | |
| ] | |
| if not active_ids or obs.get("done", False): | |
| break | |
| # Build LLM prompt | |
| user_msg = _format_inventory(obs) | |
| # Call LLM | |
| response = client.chat.completions.create( | |
| model=model, | |
| messages=[ | |
| {"role": "system", "content": _SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_msg}, | |
| ], | |
| temperature=0.2, | |
| max_tokens=64, | |
| ) | |
| llm_text = response.choices[0].message.content or "" | |
| action = _parse_action(llm_text, active_ids) | |
| # Send action to environment | |
| ws.send(json.dumps({"type": "step", "data": action})) | |
| obs = _unpack(_recv_json(ws)) | |
| step_num += 1 | |
| print( | |
| f"[STEP] {json.dumps({'task': task_name, 'step': step_num, 'action': action, 'reward': round(obs.get('reward') if obs.get('reward') is not None else 0.5, 4), 'cumulative_profit': round(float(obs.get('cumulative_profit') or 0.0), 4), 'holding_cost': round(float(obs.get('holding_cost_this_step') or 0.0), 4), 'placement_cost': float(obs.get('placement_cost_this_step') or 0.0), 'restock_cost': round(float(obs.get('restock_cost_this_step') or 0.0), 4), 'unjustified_penalty': round(float(obs.get('unjustified_discount_penalty') or 0.0), 4)})}", | |
| flush=True, | |
| ) | |
| if obs.get("done", False): | |
| break | |
| cumulative_raw = obs.get("cumulative_profit") | |
| cumulative = float(cumulative_raw) if cumulative_raw is not None else 0.0 | |
| return cumulative, grade(task_name, cumulative) | |
| finally: | |
| score = grade(task_name, cumulative) | |
| # Extra safety: ensure score is strictly within (0, 1) exclusive | |
| score = max(1e-9, min(1.0 - 1e-9, score)) | |
| print( | |
| f"[END] {json.dumps({'task': task_name, 'cumulative_profit': round(cumulative, 4), 'score': round(score, 6), 'profit_target': task.profit_target})}", | |
| flush=True, | |
| ) | |
| if ws is not None: | |
| try: | |
| ws.close() | |
| except Exception: | |
| pass | |
| # ── Main ─────────────────────────────────────────────────────────────────────── | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Run LLM agent on Store Manager tasks") | |
| parser.add_argument( | |
| "--url", | |
| default=os.getenv("STORE_ENV_URL", "http://localhost:8000"), | |
| help="Base URL of the running Store Manager server (env: STORE_ENV_URL)", | |
| ) | |
| args = parser.parse_args() | |
| # ── Mandatory env vars (per hackathon spec) ──────────────────────────────── | |
| hf_token = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") | |
| if not hf_token: | |
| print( | |
| "ERROR: HF_TOKEN environment variable not set. " | |
| "Set it to your Hugging Face token or OpenAI API key.", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(1) | |
| api_base_url = os.getenv("API_BASE_URL") | |
| model = os.getenv("MODEL_NAME") or os.getenv("OPENAI_MODEL", "gpt-4o-mini") | |
| # Build OpenAI-compatible client | |
| client_kwargs: dict = {"api_key": hf_token} | |
| if api_base_url: | |
| client_kwargs["base_url"] = api_base_url | |
| client = OpenAI(**client_kwargs) | |
| ws_url = _ws_url(args.url) | |
| results = {} | |
| for task_name in ["easy", "medium", "hard", "expert"]: | |
| try: | |
| profit, score = run_episode(ws_url, task_name, client, model) | |
| results[task_name] = {"profit": round(profit, 4), "score": round(score, 4)} | |
| except Exception as exc: | |
| print(f"[ERROR] task={task_name} error={exc}", file=sys.stderr, flush=True) | |
| results[task_name] = {"profit": 0.0, "score": grade(task_name, 0.0)} | |
| # Brief pause between tasks to avoid overwhelming the server | |
| time.sleep(1) | |
| # Final summary | |
| print("\n=== Summary ===") | |
| for task_name, r in results.items(): | |
| print(f" {task_name:<8} profit={r['profit']:>10.2f} score={r['score']:.4f}") | |
| overall = sum(r["score"] for r in results.values()) / len(results) | |
| print(f" {'AVERAGE':<8} score={overall:.4f}") | |
| if __name__ == "__main__": | |
| main() | |