nik-55/medchain-openenv-hackathon

1"""2MedChain Env — Inference Script3================================4Runs all tasks sequentially and reports scores.5 6MANDATORY environment variables:7    API_BASE_URL        The API endpoint for the LLM8    MODEL_NAME / MODEL  The model identifier for inference9    HF_TOKEN / API_KEY  Your Hugging Face / API key10 11OPTIONAL environment variables:12    LOCAL_IMAGE_NAME    Docker image tag; if set, Docker is used (highest priority).13    BASE_URL            URL of a running MedChain server (e.g. an HF Space).14    TASK_NAMES          Comma-separated list of tasks to run.15                        Default: orientation_ward,single_ward_stable,multi_ward_seasonal16    LOG_LEVEL           INFO (default) or DEBUG (writes a timestamped log to logs/)17 18Environment connection priority (per task):19    1. LOCAL_IMAGE_NAME → spin up a Docker container20    2. BASE_URL          → connect directly to that server URL21    3. Default HF Space  → https://nik-55-medchain-openenv-hackathon.hf.space22    4. Default image     → nik-55_medchain-openenv  (last-resort Docker fallback)23 24STDOUT FORMAT25- The script emits exactly three line types to stdout, in this order:26 27    [START] task=<task_name> env=medchain model=<model_name>28    [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>29    [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...,rn>30 31  Rules:32    - One [START] line at episode begin.33    - One [STEP] line per step, immediately after env.step() returns.34    - One [END] line after env.close(), always emitted (even on exception).35    - reward and rewards are formatted to 2 decimal places; score to 3.36    - done and success are lowercase booleans: true or false.37    - error is the raw error string, or null if none.38    - All fields on a single line with no newlines within a line.39"""40 41import asyncio42import json43import logging44import os45import sys46import time47import urllib.request48from datetime import datetime49from pathlib import Path50from typing import Any, Dict, List, Optional51 52from openai import BadRequestError, OpenAI, RateLimitError53 54sys.path.insert(0, str(Path(__file__).parent.parent))55 56from medchain_env import CallToolAction, MedchainEnv57 58LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()59 60_log_fmt = logging.Formatter(61    "[%(levelname)s] %(asctime)s %(message)s", datefmt="%H:%M:%S"62)63_stream_handler = logging.StreamHandler(sys.stdout)64_stream_handler.setFormatter(_log_fmt)65_handlers: list = [_stream_handler]66 67if LOG_LEVEL == "DEBUG":68    os.makedirs("logs", exist_ok=True)69    _log_filename = datetime.now().strftime("logs/inference_%Y%m%d_%H%M%S.log")70    _file_handler = logging.FileHandler(_log_filename)71    _file_handler.setFormatter(_log_fmt)72    _handlers.append(_file_handler)73    print(f"[DEBUG] Logging to file: {_log_filename}", flush=True)74 75logging.basicConfig(level=logging.WARNING, handlers=_handlers)76log = logging.getLogger(__name__)77log.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))78 79API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")80API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")81MODEL_NAME = os.getenv("MODEL_NAME") or os.getenv("MODEL", "openai/gpt-oss-120b:groq")82SMALL_MODEL = "openai/gpt-oss-20b:groq"83 84LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")85BASE_URL = os.getenv("BASE_URL")86DEFAULT_BASE_URL = "https://nik-55-medchain-openenv-hackathon.hf.space"87DEFAULT_IMAGE_NAME = "nik-55_medchain-openenv"88 89# All available tasks:90#   "orientation_ward", "single_ward_stable", "multi_ward_seasonal", "hospital_network_crisis"91_task_names_env = os.getenv(92    "TASK_NAMES",93    "orientation_ward,single_ward_stable,multi_ward_seasonal",94)95TASKS = [t.strip() for t in _task_names_env.split(",") if t.strip()]96 97# Per-task step limits (actions_per_shift × max_days + generous error headroom)98MAX_STEPS_PER_TASK = {99    "orientation_ward": 30,  # 8 actions × 2 days + headroom100    "single_ward_stable": 45,  # 10 actions × 3 days + headroom101    "multi_ward_seasonal": 75,  # 14 actions × 6 days + headroom102    "hospital_network_crisis": 180,  # 18 actions × 12 days + headroom103}104MAX_TOKENS = 6000105TEMPERATURE = 0.1106MAX_CONSECUTIVE_ERRORS = 5107SLEEP_BETWEEN_STEPS = 2108SHIFT_HISTORY_KEEP = 6109 110# 429 rate-limit handling111_429_WINDOW = 60  # seconds to track 429 count112_429_DOWNGRADE_THRESHOLD = 3  # downgrade model after this many 429s in window113_429_BASE_BACKOFF = 5  # initial backoff seconds114_429_MAX_BACKOFF = 30  # cap backoff to stay within 20-min budget115 116BENCHMARK = "medchain"117 118SYSTEM_PROMPT = """You are an experienced hospital supply chain manager operating a legacy ERP system.119Your goal is to maintain adequate medical supplies across all locations while controlling costs.120 121CRITICAL — ACTION BUDGET: You have a strictly limited number of actions per shift.122Budget does NOT roll over. Unspent actions are lost at end_shift().123 124Recommended budget allocation (highest priority first):125  1. read_inbox()                         — ALWAYS do this first to catch urgent alerts126  2. query_erp(table='inventory')         — check current stock levels across all locations127  3. submit_po(...)                        — place orders for items below safety stock (PRIORITY)128  4. end_shift()                           — call this when budget is exhausted OR tasks are done129 130Query tools (query_erp expiry/pipeline, query_forecast, query_supplier) are LOW PRIORITY.131Only use them if you have budget remaining AFTER placing critical orders.132 133MANDATORY RULES:134- If you receive "Action budget exhausted" → call end_shift() as your VERY NEXT action.135  Do NOT call any other tool. The budget cannot be restored until end_shift() is called.136- Order early: factor in lead times. If lead time is 2 days, order today to avoid stockout in 2 days.137- Expedited orders require file_justification(ticket_id=...) with a real clinical reason.138- FEFO: oldest stock consumed first — check expiry and rotate perishables proactively.139- Recalls: quarantine the recalled lot immediately, then order a replacement.140- MCI events: pre-emptive ordering beats reactive ordering. Order extra blood/critical supplies NOW.141 142Safety stock target: aim for at least (lead_time + 1) × daily_demand units on hand.143 144When calling tools, use the EXACT parameter names shown in the tool descriptions.145"""146 147 148def log_start(task: str, model: str) -> None:149    print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True)150 151 152def log_step(153    step: int, action: str, reward: float, done: bool, error: Optional[str]154) -> None:155    error_val = error if error else "null"156    print(157        f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}",158        flush=True,159    )160 161 162def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:163    rewards_str = ",".join(f"{r:.2f}" for r in rewards)164    print(165        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",166        flush=True,167    )168 169 170def _tools_to_openai_format(tools) -> List[dict]:171    """Convert MCP tools to OpenAI function-calling format."""172    openai_tools = []173    for tool in tools:174        properties = {}175        required = []176        if tool.input_schema and "properties" in tool.input_schema:177            for name, schema in tool.input_schema["properties"].items():178                properties[name] = {179                    "type": schema.get("type", "string"),180                    "description": schema.get("description", ""),181                }182            required = tool.input_schema.get("required", [])183        openai_tools.append(184            {185                "type": "function",186                "function": {187                    "name": tool.name,188                    "description": tool.description or "",189                    "parameters": {190                        "type": "object",191                        "properties": properties,192                        "required": required,193                    },194                },195            }196        )197        log.debug("Tool registered: %s (required=%s)", tool.name, required)198    return openai_tools199 200 201def _make_shift_summary(shift_day: int, end_shift_result: str) -> str:202    """Build a compact summary of a completed shift for the context window."""203    lines = []204    for line in (end_shift_result or "").splitlines():205        stripped = line.strip()206        if stripped and any(207            kw in stripped208            for kw in [209                "DEMAND:",210                "FULFILLED:",211                "DELIVERIES:",212                "EXPIRED:",213                "Spend:",214                "Waste",215                "Service Level",216                "END OF SHIFT",217                "Day ",218                "Score:",219            ]220        ):221            lines.append(stripped)222        if len(lines) >= 35:223            break224    summary_body = "\n".join(lines) if lines else (end_shift_result or "")[:900]225    return f"[SHIFT DAY {shift_day} SUMMARY]\n{summary_body}"226 227 228async def _is_url_reachable(url: str, timeout: float = 30.0) -> bool:229    loop = asyncio.get_event_loop()230 231    def _check() -> bool:232        try:233            data = b"{}"234            req = urllib.request.Request(235                url.rstrip("/") + "/reset",236                data=data,237                headers={"Content-Type": "application/json"},238                method="POST",239            )240            with urllib.request.urlopen(req, timeout=timeout) as resp:241                return resp.status == 200242        except Exception:243            return False244 245    try:246        return await loop.run_in_executor(None, _check)247    except Exception:248        return False249 250 251async def create_env(task_name: str) -> MedchainEnv:252    if LOCAL_IMAGE_NAME:253        log.info("Using Docker image '%s' for task '%s'", LOCAL_IMAGE_NAME, task_name)254        return await MedchainEnv.from_docker_image(LOCAL_IMAGE_NAME)255 256    if BASE_URL:257        log.info("Using BASE_URL '%s' for task '%s'", BASE_URL, task_name)258        env = MedchainEnv(base_url=BASE_URL)259        await env.connect()260        return env261 262    log.info("Probing default URL: %s", DEFAULT_BASE_URL)263    if await _is_url_reachable(DEFAULT_BASE_URL):264        log.info("Default URL reachable; connecting: %s", DEFAULT_BASE_URL)265        env = MedchainEnv(base_url=DEFAULT_BASE_URL)266        await env.connect()267        return env268 269    log.warning(270        "Default URL '%s' not reachable. Falling back to Docker image: %s",271        DEFAULT_BASE_URL,272        DEFAULT_IMAGE_NAME,273    )274    try:275        return await MedchainEnv.from_docker_image(DEFAULT_IMAGE_NAME)276    except Exception as docker_err:277        raise RuntimeError(278            f"All environment connection methods failed for task '{task_name}'.\n"279            f"  1. LOCAL_IMAGE_NAME: not set\n"280            f"  2. BASE_URL: not set\n"281            f"  3. Default URL ({DEFAULT_BASE_URL}): not reachable\n"282            f"  4. Default Docker image ({DEFAULT_IMAGE_NAME}): {docker_err}\n"283            "\nFix: set LOCAL_IMAGE_NAME (Docker image name) or BASE_URL (running server URL), "284            "or ensure Docker is running with the image available."285        ) from docker_err286 287 288async def run_task_episode(289    env: MedchainEnv,290    client: OpenAI,291    tools: List[dict],292    task_name: str,293) -> Dict[str, Any]:294    """Run one episode of a task and return the result."""295    tool_names = [t["function"]["name"] for t in tools]296    max_steps = MAX_STEPS_PER_TASK.get(task_name, 160)297 298    obs = await env.reset(task=task_name)299    obs = obs.observation300    dashboard = obs.metadata.get("dashboard", "")301 302    log_start(task=task_name, model=MODEL_NAME)303    log.debug(304        "[%s] Episode started. Tools: %s max_steps=%d", task_name, tool_names, max_steps305    )306 307    chat_history: List[dict] = [308        {"role": "system", "content": SYSTEM_PROMPT},309        {310            "role": "user",311            "content": f"Your shift has started. Current dashboard:\n\n{dashboard}",312        },313    ]314 315    step_count = 0316    final_reward = 0.0317    done = obs.done318    consecutive_errors = 0319    rewards: List[float] = []320    past_shift_summaries: List[str] = []321    current_shift_messages: List[dict] = []322 323    # 429 rate-limit tracking324    active_model = MODEL_NAME325    rate_limit_times: List[float] = []326    backoff_count = 0327 328    episode_start = time.monotonic()329 330    while not done and step_count < max_steps:331        step_count += 1332        log.debug(333            "[%s] Step %d/%d — %d messages in context",334            task_name,335            step_count,336            MAX_STEPS_PER_TASK[task_name],337            len(chat_history),338        )339 340        try:341            response = client.chat.completions.create(342                model=active_model,343                messages=chat_history,344                tools=tools,345                tool_choice="required",346                max_completion_tokens=MAX_TOKENS,347                temperature=TEMPERATURE,348            )349            consecutive_errors = 0350            backoff_count = 0351        except RateLimitError as e:352            now = time.monotonic()353            rate_limit_times.append(now)354            # Purge timestamps outside the tracking window355            rate_limit_times[:] = [356                t for t in rate_limit_times if now - t <= _429_WINDOW357            ]358 359            backoff_count += 1360            backoff_secs = min(361                _429_BASE_BACKOFF * (2 ** (backoff_count - 1)), _429_MAX_BACKOFF362            )363            log.warning(364                "[%s] Step %d — 429 RateLimitError (count=%d in last %ds, backoff=%.1fs): %s",365                task_name,366                step_count,367                len(rate_limit_times),368                _429_WINDOW,369                backoff_secs,370                e,371            )372 373            if (374                len(rate_limit_times) >= _429_DOWNGRADE_THRESHOLD375                and active_model != SMALL_MODEL376            ):377                active_model = SMALL_MODEL378                log.warning(379                    "[%s] Step %d — Downgrading model to %s due to repeated 429 errors",380                    task_name,381                    step_count,382                    SMALL_MODEL,383                )384 385            await asyncio.sleep(backoff_secs)386            continue387        except BadRequestError as e:388            consecutive_errors += 1389            log.warning(390                "[%s] Step %d — BadRequestError (%d/%d): %s",391                task_name,392                step_count,393                consecutive_errors,394                MAX_CONSECUTIVE_ERRORS,395                e,396            )397            if consecutive_errors >= MAX_CONSECUTIVE_ERRORS:398                log.error(399                    "[%s] Aborting after %d consecutive errors",400                    task_name,401                    MAX_CONSECUTIVE_ERRORS,402                )403                break404 405            err_msg = (406                f"Your previous tool call was rejected with an error:\n{e}\n\n"407                "Please retry with a valid tool call. If your budget is exhausted, call end_shift()."408            )409            chat_history.append({"role": "user", "content": err_msg})410            current_shift_messages.append({"role": "user", "content": err_msg})411            continue412 413        message = response.choices[0].message414        log.debug(415            "[%s] Step %d — finish_reason=%s tool_calls=%d",416            task_name,417            step_count,418            response.choices[0].finish_reason,419            len(message.tool_calls) if message.tool_calls else 0,420        )421 422        if not message.tool_calls:423            log.warning(424                "[%s] Step %d — no tool_calls in response; falling back to end_shift",425                task_name,426                step_count,427            )428            tool_name = "end_shift"429            tool_args = {}430            tool_call_id = "fallback"431        else:432            tc = message.tool_calls[0]433            tool_name = tc.function.name434            tool_call_id = tc.id435            try:436                tool_args = json.loads(tc.function.arguments)437            except (json.JSONDecodeError, AttributeError):438                log.warning(439                    "[%s] Step %d — failed to parse tool arguments: %r",440                    task_name,441                    step_count,442                    tc.function.arguments,443                )444                tool_args = {}445 446        if tool_name not in tool_names:447            log.warning(448                "[%s] Step %d — unknown tool %r; falling back to end_shift",449                task_name,450                step_count,451                tool_name,452            )453            tool_name = "end_shift"454            tool_args = {}455 456        log.debug(457            "[%s] Step %d — calling %s(%s)", task_name, step_count, tool_name, tool_args458        )459 460        assistant_msg = {461            "role": "assistant",462            "content": None,463            "tool_calls": [464                {465                    "id": tool_call_id,466                    "type": "function",467                    "function": {468                        "name": tool_name,469                        "arguments": json.dumps(tool_args),470                    },471                }472            ],473        }474        chat_history.append(assistant_msg)475        current_shift_messages.append(assistant_msg)476 477        action = CallToolAction(tool_name=tool_name, arguments=tool_args)478        step_result = await env.step(action)479        obs = step_result.observation480        done = obs.done481 482        result_text = obs.metadata.get("tool_result", str(obs.metadata))483        step_reward = obs.reward or 0.0484        step_error: Optional[str] = None485 486        if "EPISODE COMPLETE" in (result_text or ""):487            log.info("[%s] Step %d — episode complete detected", task_name, step_count)488            done = True489 490        if obs.reward is not None and obs.reward > 0:491            final_reward = obs.reward492 493        rewards.append(step_reward)494        action_str = f"{tool_name}({json.dumps(tool_args)})"495        log_step(496            step=step_count,497            action=action_str,498            reward=step_reward,499            done=done,500            error=step_error,501        )502 503        tool_result_msg = {504            "role": "tool",505            "tool_call_id": tool_call_id,506            "content": result_text[:2700] if result_text else "OK",507        }508        chat_history.append(tool_result_msg)509        current_shift_messages.append(tool_result_msg)510 511        # Budget exhausted — inject directive and skip sleep512        if "Action budget exhausted" in (result_text or ""):513            log.info(514                "[%s] Step %d — budget exhausted; injecting end_shift directive",515                task_name,516                step_count,517            )518            directive = (519                "SYSTEM ALERT: Your action budget for this shift is fully exhausted. "520                "You MUST call end_shift() as your very next action. "521                "Every other tool call will fail until you do."522            )523            chat_history.append({"role": "user", "content": directive})524            current_shift_messages.append({"role": "user", "content": directive})525            continue526 527        await asyncio.sleep(SLEEP_BETWEEN_STEPS)528 529        # Shift ended — summarise and prune context, then set up next shift530        if (531            tool_name == "end_shift"532            and "END OF SHIFT" in (result_text or "")533            and not done534        ):535            shift_day = "?"536            for part in (result_text or "").split():537                if part.isdigit():538                    shift_day = part539                    break540 541            shift_summary = _make_shift_summary(shift_day, result_text or "")542            log.debug("[%s] Shift %s summary:\n%s", task_name, shift_day, shift_summary)543            past_shift_summaries.append(shift_summary)544            log.info(545                "[%s] Step %d — shift %s ended; pruning context (%d summaries)",546                task_name,547                step_count,548                shift_day,549                len(past_shift_summaries),550            )551 552            summaries_msg = {553                "role": "user",554                "content": "COMPLETED SHIFT SUMMARIES:\n\n"555                + "\n\n".join(past_shift_summaries),556            }557            trimmed = (558                current_shift_messages[-SHIFT_HISTORY_KEEP:]559                if len(current_shift_messages) > SHIFT_HISTORY_KEEP560                else list(current_shift_messages)561            )562            # Remove budget-exhausted directives so they don't bleed into the next shift563            trimmed = [564                m565                for m in trimmed566                if "Action budget exhausted" not in (m.get("content") or "")567            ]568            # Strip orphaned leading tool-response messages to avoid API errors569            while trimmed and trimmed[0].get("role") == "tool":570                log.debug(571                    "[%s] Dropping orphaned leading tool msg (tool_call_id=%s)",572                    task_name,573                    trimmed[0].get("tool_call_id"),574                )575                trimmed = trimmed[1:]576 577            chat_history = (578                [579                    {"role": "system", "content": SYSTEM_PROMPT},580                    summaries_msg,581                ]582                + trimmed583                + [584                    {585                        "role": "user",586                        "content": "Your next shift has begun. The dashboard is shown above in the last tool result. "587                        "Continue managing the supply chain.",588                    },589                ]590            )591            current_shift_messages = []592 593    episode_duration = time.monotonic() - episode_start594    log.info(595        "[%s] Episode finished. steps=%d done=%s final_reward=%.4f",596        task_name,597        step_count,598        done,599        final_reward,600    )601    log.debug("[%s] Episode duration: %.1fs", task_name, episode_duration)602    return {603        "task": task_name,604        "reward": final_reward,605        "steps": step_count,606        "done": done,607        "rewards": rewards,608        "duration": episode_duration,609    }610 611 612async def async_main() -> None:613    if not API_KEY:614        raise SystemExit("HF_TOKEN or API_KEY must be set.")615    if not MODEL_NAME:616        raise SystemExit("MODEL_NAME or MODEL must be set.")617 618    log.info("Starting. API_BASE_URL=%s MODEL_NAME=%s", API_BASE_URL, MODEL_NAME)619    log.info("Tasks: %s", TASKS)620 621    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)622    results = []623    script_start = time.monotonic()624 625    for task_name in TASKS:626        log.info("Launching task: %s", task_name)627        task_start = time.monotonic()628        env = await create_env(task_name)629        final_reward = 0.0630        success = False631        steps = 0632        step_rewards: List[float] = []633        try:634            mcp_tools = await env.list_tools()635            tools = _tools_to_openai_format(mcp_tools)636            log.info("[%s] %d tools discovered", task_name, len(tools))637 638            result = await run_task_episode(env, client, tools, task_name)639            results.append(result)640            final_reward = result["reward"]641            steps = result["steps"]642            success = result["done"]643            step_rewards = result["rewards"]644            log.info(645                "[%s] Task complete: reward=%.4f steps=%d",646                task_name,647                final_reward,648                steps,649            )650        except Exception as e:651            log.error("[%s] Task failed with exception: %s", task_name, e)652        finally:653            try:654                await env.close()655            except Exception as e:656                log.error("[%s] env.close() failed: %s", task_name, e)657            log_end(658                success=success, steps=steps, score=final_reward, rewards=step_rewards659            )660            log.debug(661                "[%s] Total task wall time: %.1fs",662                task_name,663                time.monotonic() - task_start,664            )665 666    total_duration = time.monotonic() - script_start667    if results:668        avg_reward = sum(r["reward"] for r in results) / len(results)669        log.info("All tasks complete. avg_reward=%.4f", avg_reward)670    log.debug("Overall script duration: %.1fs", total_duration)671 672 673def main() -> None:674    asyncio.run(async_main())675 676 677if __name__ == "__main__":678    main()679