Spaces:

Code-Father
/

store

Sleeping

store / inference.py

updated infrance to remove thinking token

060c03b 7 months ago

15.2 kB

"""
Inference script for OpenEnv Store environment.
Uses an LLM (via OpenAI-compatible client) to play the kirana-store simulation.
Emits structured [START], [STEP], [END] logs for evaluation scoring.
Required env vars:
    API_BASE_URL  – LLM API endpoint
    MODEL_NAME    – model identifier
    HF_TOKEN      – API key
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import time
from typing import Any

from openai import AsyncOpenAI

from client import StoreEnv
from models import (
    OrderRequest,
    PriceUpdate,
    StoreAction,
    StoreObservation,
    dump_model,
)
from dotenv import load_dotenv
load_dotenv()

API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.5-397B-A17B")
HF_TOKEN = os.getenv("HF_TOKEN")

BENCHMARK = "store"

TASKS: list[dict[str, Any]] = [
    {
        "name": "store_short_run",
        "description": "Run the store profitably over 5 decision windows (15 days).",
        "max_steps": 5,
        "success_threshold": 0.5,
    },
    {
        "name": "store_medium_run",
        "description": "Run the store profitably over 10 decision windows (30 days).",
        "max_steps": 10,
        "success_threshold": 0.55,
    },
    {
        "name": "store_long_run",
        "description": (
            "Run the store profitably over 15 decision windows (45 days). "
            "Requires sustained capital growth and careful long-term planning."
        ),
        "max_steps": 15,
        "success_threshold": 0.6,
    },
]

ENV_URL = os.environ.get("ENV_URL", "http://localhost:8000")


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: dict[str, Any] | str,
    reward: float,
    done: bool,
    error: str | None,
) -> None:
    done_str = "true" if done else "false"
    error_str = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={done_str} error={error_str}",
        flush=True,
    )


def log_end(
    *,
    success: bool,
    steps: int,
    score: float,
    rewards: list[float],
) -> None:
    success_str = "true" if success else "false"
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={success_str} steps={steps} score={score:.2f} "
        f"rewards={rewards_str}",
        flush=True,
    )


def _format_observation(obs: StoreObservation) -> str:
    """Build a concise text summary of the observation for the LLM prompt."""
    store = obs.store
    lines = [
        "=== Store Snapshot ===",
        f"Day: {store.day}/{store.episode_horizon_days}",
        f"Bank Balance: ₹{store.bank_balance_inr:,.2f}",
        f"Customer Satisfaction: {store.customer_satisfaction}",
        f"Order Processing Capacity: {store.order_processing_capacity}",
        f"Daily Expense: ₹{store.current_daily_expense_inr:,.2f}",
        f"Cumulative Profit: ₹{store.cumulative_profit_inr:,.2f}",
        f"Active Recurring Cost: ₹{store.active_recurring_cost_inr:,.2f}",
        f"Active Action IDs: {store.active_action_ids}",
        "",
        "=== Current Stock (top items by value) ===",
    ]
    sorted_stock = sorted(obs.stock, key=lambda s: s.stock_value_inr, reverse=True)
    for item in sorted_stock[:15]:
        margin = (
            ((item.selling_price_inr - item.buying_price_inr) / item.selling_price_inr) * 100
            if item.selling_price_inr > 0
            else 0
        )
        lines.append(
            f"  {item.product_name}: qty={item.quantity_on_hand:.1f}, "
            f"buy=₹{item.buying_price_inr}, sell=₹{item.selling_price_inr}, "
            f"demand={item.demand_score}, margin={margin:.1f}%, "
            f"value=₹{item.stock_value_inr:.0f}"
        )

    low_stock = [s for s in obs.stock if 0 < s.quantity_on_hand < 15]
    if low_stock:
        lines.append("")
        lines.append("=== Low Stock Alert ===")
        for item in low_stock:
            lines.append(f"  {item.product_name}: qty={item.quantity_on_hand:.1f}")

    if obs.pending_orders:
        lines.append("")
        lines.append("=== Pending Orders ===")
        for po in obs.pending_orders:
            lines.append(
                f"  {po.order_id}: {po.product_name} x{po.quantity} from {po.seller_name}, "
                f"arrives in {po.days_remaining}d"
            )

    lines.append("")
    lines.append("=== Marketing Actions (inactive only) ===")
    for ma in obs.marketing_actions:
        if not ma.is_active:
            lines.append(
                f"  ID {ma.action_id}: {ma.description} | "
                f"cost=₹{ma.cost_inr}, recurring=₹{ma.recurring_cost_inr}/day, "
                f"satisfaction_delta={ma.customer_satisfaction_delta:+d}, "
                f"capacity_delta={ma.order_processing_capacity_delta:+d}"
            )

    if obs.metadata:
        rb = obs.metadata.get("reward_breakdown")
        if rb:
            lines.append("")
            lines.append("=== Last Reward Breakdown ===")
            lines.append(f"  Starting Capital: ₹{rb.get('starting_capital_inr', 'N/A')}")
            lines.append(f"  Ending Capital: ₹{rb.get('ending_capital_inr', 'N/A')}")
            lines.append(f"  Difference: ₹{rb.get('difference_inr', 'N/A')}")
            lines.append(f"  Recurring Costs: ₹{rb.get('recurring_costs_paid_inr', 'N/A')}")

    return "\n".join(lines)


SYSTEM_PROMPT = """\
You are a store management AI agent controlling a neighborhood kirana (grocery) store in India.
## Simulation Mechanics
Each step you submit ONE bundled action for the next 3-day decision window:
1. orders: list of supplier orders [{"product": "...", "seller": "Seller X", "quantity": N}]
2. prices: list of price updates [{"product": "...", "price": N}] (in-stock products only)
3. action_ids: list of marketing/operational action IDs to activate
## Ground-Truth Formulas
DAILY ORDERS PROCESSED = order_processing_capacity × (customer_satisfaction / 100)
For each in-stock product with quantity > 0:
  margin_pct = ((selling_price - buying_price) / selling_price) × 100
  product_selling_score = max(0, demand_score - margin_pct × 3)
TOTAL_SELLING_SCORE = sum of all product_selling_scores
Each product's daily sales = (product_selling_score / TOTAL_SELLING_SCORE) × DAILY_ORDERS_PROCESSED
  (capped at available quantity)
REWARD = 0.5 + (ending_capital - starting_capital - recurring_costs_6days) / starting_capital
  where capital = bank_balance + total_inventory_value
CUSTOMER SATISFACTION UPDATE after price/stock changes:
  new_satisfaction = (1 + (score_after - score_before) / score_after) × base_satisfaction
  Capped at [1, 100].
## Seller Profiles
- Seller A: Cheapest on staples but OVERCOMMITS — deliveries often 1-2 days late
- Seller B: Reliable baseline, matches quoted ETA consistently
- Seller C: Quotes conservatively, often arrives EARLY — best for urgent fresh restocks
- Seller D: Competitive on snacks/beverages but execution is uneven
## Key Constraints
- Episode: 30 days, each step commits 3 days, reward horizon = 6 days
- Bank balance < 0 → episode ends immediately
- Marketing actions are DURABLE: one-time cost + permanent recurring daily cost
- Starting cash is tight (₹100k) relative to inventory (~₹240k)
## Response Format
You MUST respond with ONLY valid JSON:
{
  "reasoning": "brief explanation of your decision",
  "orders": [{"product": "...", "seller": "Seller X", "quantity": N}, ...],
  "prices": [{"product": "...", "price": N}, ...],
  "action_ids": [N, ...]
}
No extra text outside JSON. Product names must exactly match the catalog.
"""


def _build_task_system_prompt(task: dict[str, Any]) -> str:
    max_steps = task["max_steps"]
    threshold = task["success_threshold"]
    goal = (
        f"## Task: {task['name']}\n\n"
        f"You have {max_steps} decision windows. "
        f"Target average reward per step ≥ {threshold:.1f}.\n\n"
        "STRATEGY PRIORITIES:\n"
        "1. Activate zero-recurring marketing early: IDs 1, 10, 6 (+75 capacity, no daily cost)\n"
        "2. Restock high-demand cheap items: Biscuits (demand 90), Bottled Water (demand 89), "
        "Eggs (demand 84), Bread (demand 86) — use Seller B/C for reliability\n"
        "3. Use Seller A for bulk staples (Rice, Oil, Spices) only when stock cover is safe\n"
        "4. Protect anchor item prices: do NOT raise Milk/Bread/Eggs/Sugar/Bottled Water\n"
        "5. NEVER reorder Kitchen Appliances (dead stock, demand 28, margin ~35%)\n"
        "6. NEVER activate ID 15 (social ads: ₹1,200/day for +0 capacity)\n"
        "7. Keep bank_balance > ₹15k safety buffer at all times\n"
        "8. After foundations, consider WhatsApp ordering (ID 4) and delivery partner (ID 7)"
    )
    return SYSTEM_PROMPT + f"\n\n{goal}"


async def get_model_action(
    client: AsyncOpenAI,
    step: int,
    observation_text: str,
    last_reward: float,
    history: list[str],
    task_system_prompt: str,
) -> dict[str, Any]:
    """Query the LLM for the next action."""
    history_text = "\n".join(history[-5:]) if history else "No history yet."

    user_msg = (
        f"Step {step} | Last reward: {last_reward:+.4f}\n\n"
        f"Recent history:\n{history_text}\n\n"
        f"Current observation:\n{observation_text}\n\n"
        "Decide your next action. Respond with ONLY valid JSON."
    )

    response = await client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": task_system_prompt},
            {"role": "user", "content": user_msg},
        ],
        temperature=0.3,
        max_tokens=4096,
    )

    content = response.choices[0].message.content or "{}"
    content = content.strip()

    print(f"[DEBUG] Raw LLM response (first 500 chars): {content[:500]}", flush=True)

    think_end = content.find("</think>")
    if think_end != -1:
        content = content[think_end + len("</think>"):].strip()

    if content.startswith("```"):
        content = content.split("\n", 1)[-1]
        if content.endswith("```"):
            content = content[: -len("```")]
        content = content.strip()

    return json.loads(content)


def _parse_action(raw: dict[str, Any]) -> StoreAction:
    """Convert the LLM JSON response into a StoreAction."""
    orders = [
        OrderRequest(
            product=o["product"],
            seller=o["seller"],
            quantity=int(o["quantity"]),
        )
        for o in raw.get("orders", [])
    ]
    prices = [
        PriceUpdate(product=p["product"], price=int(p["price"]))
        for p in raw.get("prices", [])
    ] or None
    action_ids = [int(a) for a in raw.get("action_ids", [])]
    return StoreAction(orders=orders, prices=prices, action_ids=action_ids)


async def run_task(
    llm_client: AsyncOpenAI,
    task: dict[str, Any],
) -> float:
    """Run a single task episode and return the clamped score."""
    task_name = task["name"]
    max_steps = task["max_steps"]
    threshold = task["success_threshold"]
    task_system_prompt = _build_task_system_prompt(task)

    history: list[str] = []
    rewards: list[float] = []
    steps_taken = 0
    score = 0.0
    success = False

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

    env = StoreEnv(base_url=ENV_URL)

    try:
        result = await env.reset()
        obs = result.observation
        obs_text = _format_observation(obs)
        last_reward = 0.0

        for step in range(1, max_steps + 1):
            if obs.done:
                break

            raw_action: dict[str, Any] = {}
            action_str = ""
            error = None

            try:
                raw_action = await get_model_action(
                    llm_client,
                    step,
                    obs_text,
                    last_reward,
                    history,
                    task_system_prompt,
                )
                action = _parse_action(raw_action)
                action_str = json.dumps(dump_model(action, exclude_none=True), separators=(",", ":"))
            except Exception as e:
                print(f"[DEBUG] Action parse error at step {step}: {e}", flush=True)
                action = StoreAction()
                action_str = "{}"
                error = str(e)

            try:
                result = await env.step(action)
            except Exception as e:
                print(f"[DEBUG] Step error at step {step}: {e}", flush=True)
                action = StoreAction()
                action_str = "{}"
                error = str(e)
                try:
                    result = await env.step(action)
                except Exception as e2:
                    print(f"[DEBUG] Pass-step also failed: {e2}", flush=True)
                    error = str(e2)
                    break

            obs = result.observation
            reward = result.reward or 0.0
            done = result.done

            rewards.append(reward)
            steps_taken = step
            obs_text = _format_observation(obs)
            last_reward = reward

            log_step(
                step=step,
                action=action_str,
                reward=reward,
                done=done,
                error=error,
            )

            reasoning = raw_action.get("reasoning", "")
            history.append(
                f"Step {step}: reward={reward:+.4f}, done={done}, "
                f"reasoning={reasoning!r}"
            )

            if done:
                break

        score = sum(rewards) / (max_steps * 1.0) if max_steps > 0 else 0.0
        score = min(max(score, 0.0), 1.0)
        success = score >= threshold

    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return score


async def main() -> None:
    llm_client = AsyncOpenAI(
        base_url=API_BASE_URL,
        api_key=HF_TOKEN,
    )

    print(f"[DEBUG] Starting inference with model={MODEL_NAME}", flush=True)
    print(f"[DEBUG] API_BASE_URL={API_BASE_URL}", flush=True)
    print(f"[DEBUG] ENV_URL={ENV_URL}", flush=True)
    print(f"[DEBUG] Tasks: {[t['name'] for t in TASKS]}", flush=True)

    all_scores: list[float] = []

    for task in TASKS:
        print(f"\n[DEBUG] === Running task: {task['name']} ===", flush=True)
        t0 = time.time()
        task_score = await run_task(llm_client, task)
        elapsed = time.time() - t0
        all_scores.append(task_score)
        print(
            f"[DEBUG] Task {task['name']} completed: score={task_score:.4f}, "
            f"elapsed={elapsed:.1f}s",
            flush=True,
        )

    avg_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
    print(f"\n[DEBUG] All task scores: {all_scores}", flush=True)
    print(f"[DEBUG] Average score: {avg_score:.4f}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())

	"""
	Inference script for OpenEnv Store environment.

	Uses an LLM (via OpenAI-compatible client) to play the kirana-store simulation.
	Emits structured [START], [STEP], [END] logs for evaluation scoring.

	Required env vars:
	API_BASE_URL – LLM API endpoint
	MODEL_NAME – model identifier
	HF_TOKEN – API key
	"""

	from __future__ import annotations

	import asyncio
	import json
	import os
	import sys
	import time
	from typing import Any

	from openai import AsyncOpenAI

	from client import StoreEnv
	from models import (
	OrderRequest,
	PriceUpdate,
	StoreAction,
	StoreObservation,
	dump_model,
	)
	from dotenv import load_dotenv
	load_dotenv()

	API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
	MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.5-397B-A17B")
	HF_TOKEN = os.getenv("HF_TOKEN")

	BENCHMARK = "store"

	TASKS: list[dict[str, Any]] = [
	{
	"name": "store_short_run",
	"description": "Run the store profitably over 5 decision windows (15 days).",
	"max_steps": 5,
	"success_threshold": 0.5,
	},
	{
	"name": "store_medium_run",
	"description": "Run the store profitably over 10 decision windows (30 days).",
	"max_steps": 10,
	"success_threshold": 0.55,
	},
	{
	"name": "store_long_run",
	"description": (
	"Run the store profitably over 15 decision windows (45 days). "
	"Requires sustained capital growth and careful long-term planning."
	),
	"max_steps": 15,
	"success_threshold": 0.6,
	},
	]

	ENV_URL = os.environ.get("ENV_URL", "http://localhost:8000")


	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: dict[str, Any] \| str,
	reward: float,
	done: bool,
	error: str \| None,
	) -> None:
	done_str = "true" if done else "false"
	error_str = error if error else "null"
	print(
	f"[STEP] step={step} action={action} reward={reward:.2f} "
	f"done={done_str} error={error_str}",
	flush=True,
	)


	def log_end(
	*,
	success: bool,
	steps: int,
	score: float,
	rewards: list[float],
	) -> None:
	success_str = "true" if success else "false"
	rewards_str = ",".join(f"{r:.2f}" for r in rewards)
	print(
	f"[END] success={success_str} steps={steps} score={score:.2f} "
	f"rewards={rewards_str}",
	flush=True,
	)


	def _format_observation(obs: StoreObservation) -> str:
	"""Build a concise text summary of the observation for the LLM prompt."""
	store = obs.store
	lines = [
	"=== Store Snapshot ===",
	f"Day: {store.day}/{store.episode_horizon_days}",
	f"Bank Balance: ₹{store.bank_balance_inr:,.2f}",
	f"Customer Satisfaction: {store.customer_satisfaction}",
	f"Order Processing Capacity: {store.order_processing_capacity}",
	f"Daily Expense: ₹{store.current_daily_expense_inr:,.2f}",
	f"Cumulative Profit: ₹{store.cumulative_profit_inr:,.2f}",
	f"Active Recurring Cost: ₹{store.active_recurring_cost_inr:,.2f}",
	f"Active Action IDs: {store.active_action_ids}",
	"",
	"=== Current Stock (top items by value) ===",
	]
	sorted_stock = sorted(obs.stock, key=lambda s: s.stock_value_inr, reverse=True)
	for item in sorted_stock[:15]:
	margin = (
	((item.selling_price_inr - item.buying_price_inr) / item.selling_price_inr) * 100
	if item.selling_price_inr > 0
	else 0
	)
	lines.append(
	f" {item.product_name}: qty={item.quantity_on_hand:.1f}, "
	f"buy=₹{item.buying_price_inr}, sell=₹{item.selling_price_inr}, "
	f"demand={item.demand_score}, margin={margin:.1f}%, "
	f"value=₹{item.stock_value_inr:.0f}"
	)

	low_stock = [s for s in obs.stock if 0 < s.quantity_on_hand < 15]
	if low_stock:
	lines.append("")
	lines.append("=== Low Stock Alert ===")
	for item in low_stock:
	lines.append(f" {item.product_name}: qty={item.quantity_on_hand:.1f}")

	if obs.pending_orders:
	lines.append("")
	lines.append("=== Pending Orders ===")
	for po in obs.pending_orders:
	lines.append(
	f" {po.order_id}: {po.product_name} x{po.quantity} from {po.seller_name}, "
	f"arrives in {po.days_remaining}d"
	)

	lines.append("")
	lines.append("=== Marketing Actions (inactive only) ===")
	for ma in obs.marketing_actions:
	if not ma.is_active:
	lines.append(
	f" ID {ma.action_id}: {ma.description} \| "
	f"cost=₹{ma.cost_inr}, recurring=₹{ma.recurring_cost_inr}/day, "
	f"satisfaction_delta={ma.customer_satisfaction_delta:+d}, "
	f"capacity_delta={ma.order_processing_capacity_delta:+d}"
	)

	if obs.metadata:
	rb = obs.metadata.get("reward_breakdown")
	if rb:
	lines.append("")
	lines.append("=== Last Reward Breakdown ===")
	lines.append(f" Starting Capital: ₹{rb.get('starting_capital_inr', 'N/A')}")
	lines.append(f" Ending Capital: ₹{rb.get('ending_capital_inr', 'N/A')}")
	lines.append(f" Difference: ₹{rb.get('difference_inr', 'N/A')}")
	lines.append(f" Recurring Costs: ₹{rb.get('recurring_costs_paid_inr', 'N/A')}")

	return "\n".join(lines)


	SYSTEM_PROMPT = """\
	You are a store management AI agent controlling a neighborhood kirana (grocery) store in India.

	## Simulation Mechanics

	Each step you submit ONE bundled action for the next 3-day decision window:
	1. orders: list of supplier orders [{"product": "...", "seller": "Seller X", "quantity": N}]
	2. prices: list of price updates [{"product": "...", "price": N}] (in-stock products only)
	3. action_ids: list of marketing/operational action IDs to activate

	## Ground-Truth Formulas

	DAILY ORDERS PROCESSED = order_processing_capacity × (customer_satisfaction / 100)

	For each in-stock product with quantity > 0:
	margin_pct = ((selling_price - buying_price) / selling_price) × 100
	product_selling_score = max(0, demand_score - margin_pct × 3)

	TOTAL_SELLING_SCORE = sum of all product_selling_scores
	Each product's daily sales = (product_selling_score / TOTAL_SELLING_SCORE) × DAILY_ORDERS_PROCESSED
	(capped at available quantity)

	REWARD = 0.5 + (ending_capital - starting_capital - recurring_costs_6days) / starting_capital
	where capital = bank_balance + total_inventory_value

	CUSTOMER SATISFACTION UPDATE after price/stock changes:
	new_satisfaction = (1 + (score_after - score_before) / score_after) × base_satisfaction
	Capped at [1, 100].

	## Seller Profiles
	- Seller A: Cheapest on staples but OVERCOMMITS — deliveries often 1-2 days late
	- Seller B: Reliable baseline, matches quoted ETA consistently
	- Seller C: Quotes conservatively, often arrives EARLY — best for urgent fresh restocks
	- Seller D: Competitive on snacks/beverages but execution is uneven

	## Key Constraints
	- Episode: 30 days, each step commits 3 days, reward horizon = 6 days
	- Bank balance < 0 → episode ends immediately
	- Marketing actions are DURABLE: one-time cost + permanent recurring daily cost
	- Starting cash is tight (₹100k) relative to inventory (~₹240k)

	## Response Format
	You MUST respond with ONLY valid JSON:
	{
	"reasoning": "brief explanation of your decision",
	"orders": [{"product": "...", "seller": "Seller X", "quantity": N}, ...],
	"prices": [{"product": "...", "price": N}, ...],
	"action_ids": [N, ...]
	}
	No extra text outside JSON. Product names must exactly match the catalog.
	"""


	def _build_task_system_prompt(task: dict[str, Any]) -> str:
	max_steps = task["max_steps"]
	threshold = task["success_threshold"]
	goal = (
	f"## Task: {task['name']}\n\n"
	f"You have {max_steps} decision windows. "
	f"Target average reward per step ≥ {threshold:.1f}.\n\n"
	"STRATEGY PRIORITIES:\n"
	"1. Activate zero-recurring marketing early: IDs 1, 10, 6 (+75 capacity, no daily cost)\n"
	"2. Restock high-demand cheap items: Biscuits (demand 90), Bottled Water (demand 89), "
	"Eggs (demand 84), Bread (demand 86) — use Seller B/C for reliability\n"
	"3. Use Seller A for bulk staples (Rice, Oil, Spices) only when stock cover is safe\n"
	"4. Protect anchor item prices: do NOT raise Milk/Bread/Eggs/Sugar/Bottled Water\n"
	"5. NEVER reorder Kitchen Appliances (dead stock, demand 28, margin ~35%)\n"
	"6. NEVER activate ID 15 (social ads: ₹1,200/day for +0 capacity)\n"
	"7. Keep bank_balance > ₹15k safety buffer at all times\n"
	"8. After foundations, consider WhatsApp ordering (ID 4) and delivery partner (ID 7)"
	)
	return SYSTEM_PROMPT + f"\n\n{goal}"


	async def get_model_action(
	client: AsyncOpenAI,
	step: int,
	observation_text: str,
	last_reward: float,
	history: list[str],
	task_system_prompt: str,
	) -> dict[str, Any]:
	"""Query the LLM for the next action."""
	history_text = "\n".join(history[-5:]) if history else "No history yet."

	user_msg = (
	f"Step {step} \| Last reward: {last_reward:+.4f}\n\n"
	f"Recent history:\n{history_text}\n\n"
	f"Current observation:\n{observation_text}\n\n"
	"Decide your next action. Respond with ONLY valid JSON."
	)

	response = await client.chat.completions.create(
	model=MODEL_NAME,
	messages=[
	{"role": "system", "content": task_system_prompt},
	{"role": "user", "content": user_msg},
	],
	temperature=0.3,
	max_tokens=4096,
	)

	content = response.choices[0].message.content or "{}"
	content = content.strip()

	print(f"[DEBUG] Raw LLM response (first 500 chars): {content[:500]}", flush=True)

	think_end = content.find("</think>")
	if think_end != -1:
	content = content[think_end + len("</think>"):].strip()

	if content.startswith("```"):
	content = content.split("\n", 1)[-1]
	if content.endswith("```"):
	content = content[: -len("```")]
	content = content.strip()

	return json.loads(content)


	def _parse_action(raw: dict[str, Any]) -> StoreAction:
	"""Convert the LLM JSON response into a StoreAction."""
	orders = [
	OrderRequest(
	product=o["product"],
	seller=o["seller"],
	quantity=int(o["quantity"]),
	)
	for o in raw.get("orders", [])
	]
	prices = [
	PriceUpdate(product=p["product"], price=int(p["price"]))
	for p in raw.get("prices", [])
	] or None
	action_ids = [int(a) for a in raw.get("action_ids", [])]
	return StoreAction(orders=orders, prices=prices, action_ids=action_ids)


	async def run_task(
	llm_client: AsyncOpenAI,
	task: dict[str, Any],
	) -> float:
	"""Run a single task episode and return the clamped score."""
	task_name = task["name"]
	max_steps = task["max_steps"]
	threshold = task["success_threshold"]
	task_system_prompt = _build_task_system_prompt(task)

	history: list[str] = []
	rewards: list[float] = []
	steps_taken = 0
	score = 0.0
	success = False

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

	env = StoreEnv(base_url=ENV_URL)

	try:
	result = await env.reset()
	obs = result.observation
	obs_text = _format_observation(obs)
	last_reward = 0.0

	for step in range(1, max_steps + 1):
	if obs.done:
	break

	raw_action: dict[str, Any] = {}
	action_str = ""
	error = None

	try:
	raw_action = await get_model_action(
	llm_client,
	step,
	obs_text,
	last_reward,
	history,
	task_system_prompt,
	)
	action = _parse_action(raw_action)
	action_str = json.dumps(dump_model(action, exclude_none=True), separators=(",", ":"))
	except Exception as e:
	print(f"[DEBUG] Action parse error at step {step}: {e}", flush=True)
	action = StoreAction()
	action_str = "{}"
	error = str(e)

	try:
	result = await env.step(action)
	except Exception as e:
	print(f"[DEBUG] Step error at step {step}: {e}", flush=True)
	action = StoreAction()
	action_str = "{}"
	error = str(e)
	try:
	result = await env.step(action)
	except Exception as e2:
	print(f"[DEBUG] Pass-step also failed: {e2}", flush=True)
	error = str(e2)
	break

	obs = result.observation
	reward = result.reward or 0.0
	done = result.done

	rewards.append(reward)
	steps_taken = step
	obs_text = _format_observation(obs)
	last_reward = reward

	log_step(
	step=step,
	action=action_str,
	reward=reward,
	done=done,
	error=error,
	)

	reasoning = raw_action.get("reasoning", "")
	history.append(
	f"Step {step}: reward={reward:+.4f}, done={done}, "
	f"reasoning={reasoning!r}"
	)

	if done:
	break

	score = sum(rewards) / (max_steps * 1.0) if max_steps > 0 else 0.0
	score = min(max(score, 0.0), 1.0)
	success = score >= threshold

	finally:
	try:
	await env.close()
	except Exception as e:
	print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
	log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

	return score


	async def main() -> None:
	llm_client = AsyncOpenAI(
	base_url=API_BASE_URL,
	api_key=HF_TOKEN,
	)

	print(f"[DEBUG] Starting inference with model={MODEL_NAME}", flush=True)
	print(f"[DEBUG] API_BASE_URL={API_BASE_URL}", flush=True)
	print(f"[DEBUG] ENV_URL={ENV_URL}", flush=True)
	print(f"[DEBUG] Tasks: {[t['name'] for t in TASKS]}", flush=True)

	all_scores: list[float] = []

	for task in TASKS:
	print(f"\n[DEBUG] === Running task: {task['name']} ===", flush=True)
	t0 = time.time()
	task_score = await run_task(llm_client, task)
	elapsed = time.time() - t0
	all_scores.append(task_score)
	print(
	f"[DEBUG] Task {task['name']} completed: score={task_score:.4f}, "
	f"elapsed={elapsed:.1f}s",
	flush=True,
	)

	avg_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
	print(f"\n[DEBUG] All task scores: {all_scores}", flush=True)
	print(f"[DEBUG] Average score: {avg_score:.4f}", flush=True)


	if __name__ == "__main__":
	asyncio.run(main())