CoolFace
Apppublic

dhruv-punia-bits/memory-compaction-openenv

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
inference.py634 linesDownload Raw Back to root
1from __future__ import annotations2 3import json4import os5import re6import subprocess7import sys8import time9from typing import Any10 11import requests12from openai import OpenAI13 14 15BENCHMARK = "memory-compaction-openenv"16ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://127.0.0.1:7860")17API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")18MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")19OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") or "missing-api-key"20LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME", "memory-compaction-openenv")21MODEL_REQUEST_DELAY_SECONDS = float(os.getenv("MODEL_REQUEST_DELAY_SECONDS", "8"))22MODEL_MAX_RETRIES = int(os.getenv("MODEL_MAX_RETRIES", "3"))23MODEL_RETRY_BACKOFF_SECONDS = float(os.getenv("MODEL_RETRY_BACKOFF_SECONDS", "8"))24 25CANONICAL_TASKS = {26    "easy": 1101,27    "medium": 2201,28    "hard": 3301,29}30 31_LAST_MODEL_REQUEST_TS = 0.032STRICT_SCORE_EPSILON = 0.00133 34SYSTEM_PROMPT = """You are the planning layer for a memory-compaction system.35 36Your job is to read the current environment observation and produce exactly ONE valid JSON object.37 38You must decide the next action for a memory manager. The action must follow this exact schema:39{40  "operation": "append_memory" | "update_memory" | "delete_memory" | "replace_summary" | "noop",41  "memory_items": [42    {43      "memory_id": "string",44      "type": "preference" | "fact" | "task" | "constraint" | "plan" | "correction",45      "subject": "string",46      "predicate": "string",47      "object": "string",48      "confidence": 0.0,49      "source_turn_ids": [1],50      "source_text": "string",51      "status": "active" | "superseded" | "uncertain",52      "updated_from_memory_id": "string or null",53      "expires_at": "string or null",54      "importance": 0.0,55      "task_relevance": 0.0,56      "requires_confirmation": false57    }58  ],59  "summary_text": "short summary",60  "rationale": "one sentence"61}62 63Output requirements:641. Output ONLY JSON. No prose. No markdown. No code fences. No explanation before or after JSON.652. The top-level value must be a JSON object, not an array or string.663. Always include all 4 top-level keys: operation, memory_items, summary_text, rationale.674. memory_items must always be an array. Use [] when there is nothing durable to write.685. summary_text must always be a string. rationale must always be a string.696. Use null, true, false exactly as valid JSON values when needed.707. Do not invent extra top-level fields.718. If unsure, prefer a conservative valid JSON action such as replace_summary or noop over malformed output.72 73Decision policy:74- Prefer precision over recall for durable memory.75- Store durable memory only when the user provides a concrete fact, standing preference, task detail, correction, or named relationship likely to matter later.76- Use update_memory when a new value supersedes an older one.77- Use status="uncertain" and requires_confirmation=true when the user sounds tentative or explicitly unconfirmed.78- Do not store assistant chatter or generic small talk as durable memory.79- If a turn is contextual but not durable, use replace_summary with an empty memory_items list.80- If a turn should be ignored, use noop with an empty memory_items list.81 82Important:83- The current_turn, recent_turns, working_summary, and durable_memory in the observation tell you what is already known.84- Preserve corrections cleanly.85- Never output comments or trailing commas.86- Never wrap JSON in backticks.87 88Example 1: direct durable fact89Observation snippet:90{"current_turn":{"turn_id":1,"speaker":"user","text":"My name is Ava and I am traveling to Kyoto next month."}}91Valid output:92{"operation":"append_memory","memory_items":[{"memory_id":"name","type":"fact","subject":"Ava","predicate":"is","object":"traveler","confidence":0.97,"source_turn_ids":[1],"source_text":"My name is Ava and I am traveling to Kyoto next month.","status":"active","updated_from_memory_id":null,"expires_at":null,"importance":0.7,"task_relevance":0.55,"requires_confirmation":false},{"memory_id":"trip","type":"fact","subject":"Ava","predicate":"travel_city","object":"Kyoto","confidence":0.97,"source_turn_ids":[1],"source_text":"My name is Ava and I am traveling to Kyoto next month.","status":"active","updated_from_memory_id":null,"expires_at":null,"importance":0.95,"task_relevance":0.95,"requires_confirmation":false}],"summary_text":"User introduced herself and upcoming travel.","rationale":"Directly stated durable facts should be stored."}93 94Example 2: contextual but not durable95Observation snippet:96{"current_turn":{"turn_id":5,"speaker":"assistant","text":"That sounds exciting, I can help organize it."}}97Valid output:98{"operation":"noop","memory_items":[],"summary_text":"assistant: That sounds exciting, I can help organize it.","rationale":"Assistant chatter should not become durable memory."}99 100Example 3: uncertain future plan101Observation snippet:102{"current_turn":{"turn_id":10,"speaker":"user","text":"I might present this work next week, but that is not confirmed yet."}}103Valid output:104{"operation":"append_memory","memory_items":[{"memory_id":"presentation_plan_uncertain","type":"plan","subject":"project","predicate":"may_present","object":"present this work next week","confidence":0.4,"source_turn_ids":[10],"source_text":"I might present this work next week, but that is not confirmed yet.","status":"uncertain","updated_from_memory_id":null,"expires_at":null,"importance":0.25,"task_relevance":0.25,"requires_confirmation":true}],"summary_text":"Possible presentation next week, not confirmed.","rationale":"Tentative plans should be stored only as uncertain memory."}105 106Now read the actual observation and return the single valid JSON action object only.107"""108 109 110def compact_text(text: str, max_words: int = 28) -> str:111    return " ".join(text.split()[:max_words])112 113 114def memory_value(observation: dict[str, Any], memory_id: str, field: str) -> str | None:115    for item in observation.get("durable_memory", []):116        if item.get("memory_id") == memory_id:117            return item.get(field)118    return None119 120 121def extract_between(text: str, start_marker: str, end_marker: str | None = None) -> str | None:122    lowered = text.lower()123    start_index = lowered.find(start_marker.lower())124    if start_index == -1:125        return None126    start_index += len(start_marker)127    remainder = text[start_index:]128    if end_marker is None:129        return remainder.strip(" .,:;")130    end_index = remainder.lower().find(end_marker.lower())131    if end_index == -1:132        return remainder.strip(" .,:;")133    return remainder[:end_index].strip(" .,:;")134 135 136def extract_sentence_preference(text: str) -> str | None:137    match = re.search(r"\bprefer\b\s+(.+?)(?:\.|,|$)", text, flags=re.IGNORECASE)138    return match.group(1).strip() if match else None139 140 141def extract_project_and_teammate(text: str) -> tuple[str | None, str | None]:142    match = re.search(r"building a\s+(.+?)\s+with\s+([A-Za-z][A-Za-z\-']+)", text, flags=re.IGNORECASE)143    if not match:144        return None, None145    return match.group(1).strip(" .,"), match.group(2).strip(" .,")146 147 148def extract_deadline_and_tool(text: str) -> tuple[str | None, str | None]:149    deadline_match = re.search(r"deadline is\s+(.+?)(?:,|\.)", text, flags=re.IGNORECASE)150    tool_match = re.search(r"track work in\s+(.+?)(?:\.|,|$)", text, flags=re.IGNORECASE)151    return (152        deadline_match.group(1).strip() if deadline_match else None,153        tool_match.group(1).strip() if tool_match else None,154    )155 156 157def extract_named_role_update(text: str) -> tuple[str | None, str | None]:158    match = re.search(r"\b(mentor|owner|reviewer)\b.*?\bis\s+([A-Za-z][A-Za-z\-']+)", text, flags=re.IGNORECASE)159    if not match:160        return None, None161    return match.group(1).lower(), match.group(2).strip(" .,")162 163 164def is_explicit_commit_turn(text: str) -> bool:165    lowered = text.lower()166    return any(167        phrase in lowered168        for phrase in [169            "my name is",170            "prefer",171            "i am building a",172            "deadline is",173            "deadline moved to",174            "correction:",175            "mentor",176            "owner",177            "reviewer",178            "actually i switched from",179        ]180    )181 182 183def is_ambiguous_turn(text: str) -> bool:184    lowered = text.lower()185    return "might " in lowered or "not confirmed yet" in lowered186 187 188def memory_template(189    memory_id: str,190    memory_type: str,191    subject: str,192    predicate: str,193    obj: str,194    turn_id: int,195    source_text: str,196    *,197    confidence: float,198    status: str = "active",199    updated_from_memory_id: str | None = None,200    importance: float = 0.8,201    task_relevance: float = 0.8,202    requires_confirmation: bool = False,203) -> dict[str, Any]:204    return {205        "memory_id": memory_id,206        "type": memory_type,207        "subject": subject,208        "predicate": predicate,209        "object": obj,210        "confidence": confidence,211        "source_turn_ids": [turn_id],212        "source_text": source_text,213        "status": status,214        "updated_from_memory_id": updated_from_memory_id,215        "expires_at": None,216        "importance": importance,217        "task_relevance": task_relevance,218        "requires_confirmation": requires_confirmation,219    }220 221 222def heuristic_turn_plan(observation: dict[str, Any]) -> dict[str, Any]:223    current_turn = observation.get("current_turn") or {}224    text = current_turn.get("text", "")225    lowered = text.lower()226    summary_text = compact_text(f"{current_turn.get('speaker', 'turn')}: {text}".strip())227 228    if current_turn.get("speaker") != "user":229        return {230            "store_level": "ignore",231            "importance": 0.1,232            "task_relevance": 0.1,233            "requires_confirmation": False,234            "summary_text": summary_text,235            "reason": "Assistant turns are not durable memory by default.",236        }237    if "ignore generic small talk" in lowered or "finished reading" in lowered:238        return {239            "store_level": "ignore",240            "importance": 0.05,241            "task_relevance": 0.05,242            "requires_confirmation": False,243            "summary_text": summary_text,244            "reason": "Low-value or explicitly ignorable detail.",245        }246    if is_ambiguous_turn(text):247        return {248            "store_level": "summary_only",249            "importance": 0.25,250            "task_relevance": 0.25,251            "requires_confirmation": True,252            "summary_text": summary_text,253            "reason": "Ambiguous statement should stay uncertain.",254        }255    if is_explicit_commit_turn(text):256        return {257            "store_level": "durable",258            "importance": 0.85,259            "task_relevance": 0.9,260            "requires_confirmation": False,261            "summary_text": summary_text,262            "reason": "Explicit durable fact or update.",263        }264    return {265        "store_level": "summary_only",266        "importance": 0.35,267        "task_relevance": 0.35,268        "requires_confirmation": False,269        "summary_text": summary_text,270        "reason": "Contextual turn worth summarizing but not storing durably.",271    }272 273 274def normalize_plan(plan: dict[str, Any], observation: dict[str, Any]) -> dict[str, Any]:275    fallback = heuristic_turn_plan(observation)276    store_level = plan.get("store_level", fallback["store_level"])277    if store_level not in {"durable", "summary_only", "ignore"}:278        store_level = fallback["store_level"]279    return {280        "store_level": store_level,281        "importance": clamp_float(plan.get("importance", fallback["importance"])),282        "task_relevance": clamp_float(plan.get("task_relevance", fallback["task_relevance"])),283        "requires_confirmation": bool(plan.get("requires_confirmation", fallback["requires_confirmation"])),284        "summary_text": compact_text(str(plan.get("summary_text", fallback["summary_text"])), max_words=28),285        "reason": str(plan.get("reason", fallback["reason"])),286    }287 288 289def clamp_float(value: Any) -> float:290    try:291        return max(0.0, min(1.0, float(value)))292    except (TypeError, ValueError):293        return 0.5294 295 296def heuristic_action(observation: dict[str, Any]) -> tuple[dict[str, Any], str]:297    plan = heuristic_turn_plan(observation)298    return build_hybrid_action(observation, plan), "heuristic"299 300 301def wait_for_model_budget() -> None:302    global _LAST_MODEL_REQUEST_TS303    if MODEL_REQUEST_DELAY_SECONDS <= 0:304        _LAST_MODEL_REQUEST_TS = time.time()305        return306    now = time.time()307    wait_seconds = MODEL_REQUEST_DELAY_SECONDS - (now - _LAST_MODEL_REQUEST_TS)308    if wait_seconds > 0:309        time.sleep(wait_seconds)310    _LAST_MODEL_REQUEST_TS = time.time()311 312 313def should_retry_model_exception(exc: Exception) -> bool:314    status_code = getattr(exc, "status_code", None)315    if status_code in {408, 409, 429, 500, 502, 503, 504}:316        return True317    message = str(exc).lower()318    return isinstance(exc, json.JSONDecodeError) or any(319        token in message for token in ["rate limit", "too many requests", "timeout", "temporar", "unavailable", "overloaded", "server error"]320    )321 322 323def format_model_exception(exc: Exception) -> str:324    status_code = getattr(exc, "status_code", None)325    body = getattr(exc, "body", None)326    response = getattr(exc, "response", None)327    message = str(exc).replace("\n", " ").strip()328 329    parts: list[str] = [f"type={exc.__class__.__name__}"]330    if status_code is not None:331        parts.append(f"status={status_code}")332    if message:333        parts.append(f"message={message}")334    if body:335        parts.append(f"body={body}")336    elif response is not None:337        response_text = getattr(response, "text", None)338        if response_text:339            parts.append(f"response={response_text}")340    return " | ".join(parts)341 342 343def model_plan(client: OpenAI, observation: dict[str, Any]) -> dict[str, Any]:344    last_exc: Exception | None = None345    for attempt in range(1, MODEL_MAX_RETRIES + 1):346        try:347            wait_for_model_budget()348            completion = client.chat.completions.create(349                model=MODEL_NAME,350                messages=[351                    {"role": "system", "content": SYSTEM_PROMPT},352                    {"role": "user", "content": json.dumps(observation, indent=2)},353                ],354                temperature=0.0,355                response_format={"type": "json_object"},356            )357            content = completion.choices[0].message.content or "{}"358            return json.loads(content)359        except Exception as exc:360            last_exc = exc361            if attempt >= MODEL_MAX_RETRIES or not should_retry_model_exception(exc):362                raise363            time.sleep(MODEL_RETRY_BACKOFF_SECONDS * attempt)364    if last_exc is not None:365        raise last_exc366    raise RuntimeError("model_plan failed without exception")367 368 369def start_local_env_if_needed() -> subprocess.Popen[str] | None:370    try:371        response = requests.get(f"{ENV_BASE_URL}/health", timeout=3)372        if response.status_code == 200:373            return None374    except requests.RequestException:375        pass376 377    env = os.environ.copy()378    process = subprocess.Popen(379        [sys.executable, "-m", "uvicorn", "server.app:app", "--host", "127.0.0.1", "--port", "7860"],380        stdout=subprocess.DEVNULL,381        stderr=subprocess.DEVNULL,382        env=env,383    )384 385    deadline = time.time() + 20386    while time.time() < deadline:387        try:388            response = requests.get(f"{ENV_BASE_URL}/health", timeout=2)389            if response.status_code == 200:390                return process391        except requests.RequestException:392            time.sleep(0.5)393 394    process.terminate()395    raise RuntimeError(f"Could not start local environment for {LOCAL_IMAGE_NAME} at {ENV_BASE_URL}.")396 397 398def extract_durable_memories(observation: dict[str, Any], plan: dict[str, Any]) -> list[dict[str, Any]]:399    current_turn = observation.get("current_turn") or {}400    if current_turn.get("speaker") != "user":401        return []402    text = current_turn.get("text", "")403    lowered = text.lower()404    turn_id = current_turn.get("turn_id", 0)405    name = memory_value(observation, "name", "subject") or "user"406    project = memory_value(observation, "project", "object") or "project"407    importance = plan["importance"]408    task_relevance = plan["task_relevance"]409 410    if "my name is" in lowered and "traveling to" in lowered:411        parsed_name = extract_between(text, "my name is", ", and") or extract_between(text, "my name is", " and")412        city = extract_between(text, "traveling to", " next month")413        if not parsed_name or not city:414            return []415        return [416            memory_template("name", "fact", parsed_name, "is", "traveler", turn_id, text, confidence=0.97, importance=0.7, task_relevance=0.55),417            memory_template("trip", "fact", parsed_name, "travel_city", city, turn_id, text, confidence=0.97, importance=0.95, task_relevance=0.95),418        ]419    if "actually i switched from" in lowered and " to " in lowered:420        switched_match = re.search(r"switched from\s+.+?\s+to\s+(.+?)(?:\.|,|$)", text, flags=re.IGNORECASE)421        latest_pref = switched_match.group(1).strip() if switched_match else None422        if not latest_pref:423            return []424        return [425            memory_template("diet", "preference", name, "prefers", latest_pref, turn_id, text, confidence=0.98, status="superseded", updated_from_memory_id="diet_v2", importance=0.9, task_relevance=1.0),426            memory_template("diet_v2", "correction", name, "prefers", latest_pref, turn_id, text, confidence=0.98, updated_from_memory_id="diet", importance=0.98, task_relevance=1.0),427        ]428    if "prefer" in lowered:429        preference = extract_sentence_preference(text)430        if not preference:431            return []432        return [memory_template("diet", "preference", name, "prefers", preference, turn_id, text, confidence=0.93, importance=importance, task_relevance=task_relevance)]433    if "i am building a" in lowered and "with" in lowered:434        built_project, teammate = extract_project_and_teammate(text)435        if not built_project or not teammate:436            return []437        return [438            memory_template("project", "plan", name, "building", built_project, turn_id, text, confidence=0.94, importance=0.95, task_relevance=0.95),439            memory_template("teammate", "fact", built_project, "teammate", teammate, turn_id, text, confidence=0.92, importance=0.85, task_relevance=0.9),440        ]441    if "deadline is" in lowered and "track work in" in lowered:442        deadline, tool = extract_deadline_and_tool(text)443        if not deadline or not tool:444            return []445        return [446            memory_template("deadline", "task", project, "deadline", deadline, turn_id, text, confidence=0.94, importance=0.95, task_relevance=1.0),447            memory_template("tool", "preference", name, "tracks_work_in", tool, turn_id, text, confidence=0.88, importance=0.7, task_relevance=0.8),448        ]449    if "deadline moved to" in lowered:450        new_deadline = extract_between(text, "deadline moved to")451        if not new_deadline:452            return []453        return [454            memory_template("deadline", "task", project, "deadline", new_deadline, turn_id, text, confidence=0.98, status="superseded", updated_from_memory_id="deadline_v2", importance=0.9, task_relevance=1.0),455            memory_template("deadline_v2", "correction", project, "deadline", new_deadline, turn_id, text, confidence=0.98, updated_from_memory_id="deadline", importance=0.98, task_relevance=1.0),456        ]457    role_name, person = extract_named_role_update(text)458    if role_name and person:459        memory_id = "mentor" if role_name == "mentor" else f"{role_name}_contact"460        return [memory_template(memory_id, "fact", project, role_name, person, turn_id, text, confidence=0.86, importance=0.65, task_relevance=0.75)]461    if is_ambiguous_turn(text):462        return [463            memory_template(464                "presentation_plan_uncertain",465                "plan",466                project,467                "may_present",468                "present this work next week",469                turn_id,470                text,471                confidence=0.4,472                status="uncertain",473                importance=0.25,474                task_relevance=0.25,475                requires_confirmation=True,476            )477        ]478    return []479 480 481def build_hybrid_action(observation: dict[str, Any], plan: dict[str, Any]) -> dict[str, Any]:482    current_turn = observation.get("current_turn") or {}483    turn_text = current_turn.get("text", "")484    memories = extract_durable_memories(observation, plan)485 486    if current_turn.get("speaker") == "user" and is_explicit_commit_turn(turn_text) and memories:487        operation = "update_memory" if any(item["updated_from_memory_id"] for item in memories) else "append_memory"488        return {489            "operation": operation,490            "memory_items": memories,491            "summary_text": plan["summary_text"],492            "rationale": f"Deterministic explicit commit. {plan['reason']}",493        }494    if plan["store_level"] == "ignore" and not memories:495        return {496            "operation": "noop",497            "memory_items": [],498            "summary_text": observation.get("working_summary", ""),499            "rationale": plan["reason"],500        }501    if is_ambiguous_turn(turn_text):502        if plan["requires_confirmation"] and memories:503            return {504                "operation": "append_memory",505                "memory_items": memories,506                "summary_text": plan["summary_text"],507                "rationale": f"Ambiguous memory kept as uncertain. {plan['reason']}",508            }509        return {510            "operation": "replace_summary",511            "memory_items": [],512            "summary_text": plan["summary_text"],513            "rationale": plan["reason"],514        }515    if plan["store_level"] == "durable" and memories:516        operation = "update_memory" if any(item["updated_from_memory_id"] for item in memories) else "append_memory"517        return {518            "operation": operation,519            "memory_items": memories,520            "summary_text": plan["summary_text"],521            "rationale": f"Borderline durable decision. {plan['reason']}",522        }523    return {524        "operation": "replace_summary",525        "memory_items": [],526        "summary_text": plan["summary_text"],527        "rationale": plan["reason"],528    }529 530 531def select_baseline_tasks(tasks: list[dict[str, Any]]) -> list[dict[str, Any]]:532    tasks_by_seed = {task["seed"]: task for task in tasks}533    return [tasks_by_seed[seed] for seed in CANONICAL_TASKS.values() if seed in tasks_by_seed]534 535 536def format_action(action: dict[str, Any]) -> str:537    if not action.get("memory_items"):538        return action.get("operation", "noop")539    ids = ",".join(item.get("memory_id", "?") for item in action["memory_items"])540    return f"{action.get('operation', 'noop')}[{ids}]"541 542 543def log_start(task: str, env: str, model: str) -> None:544    print(f"[START] task={task} env={env} model={model}", flush=True)545 546 547def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None:548    error_value = error if error else "null"549    print(550        f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_value}",551        flush=True,552    )553 554 555def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None:556    rewards_str = ",".join(f"{reward:.2f}" for reward in rewards)557    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)558 559 560def strict_score(value: float) -> float:561    return max(STRICT_SCORE_EPSILON, min(1.0 - STRICT_SCORE_EPSILON, value))562 563 564def post_json(path: str, payload: dict[str, Any]) -> dict[str, Any]:565    response = requests.post(f"{ENV_BASE_URL}{path}", json=payload, timeout=60)566    response.raise_for_status()567    return response.json()568 569 570def run_task(client: OpenAI, task: dict[str, Any]) -> float:571    task_name = f"{task['difficulty']}-{task['seed']}"572    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)573 574    rewards: list[float] = []575    steps_taken = 0576    final_score = 0.0577    success = False578    try:579        observation = post_json("/reset", {"difficulty": task["difficulty"], "seed": task["seed"]})["observation"]580        done = False581        while not done:582            error: str | None = None583            try:584                raw_plan = model_plan(client, observation)585                plan = normalize_plan(raw_plan, observation)586            except Exception as exc:587                plan = heuristic_turn_plan(observation)588                error = "model_request_failed"589                print(f"[LLM_ERROR] {format_model_exception(exc)}", flush=True)590 591            action = build_hybrid_action(observation, plan)592            response = post_json("/step", action)593            reward = float(response["reward"])594            done = bool(response["done"])595            steps_taken += 1596            rewards.append(reward)597            log_step(598                step=steps_taken,599                action=format_action(action),600                reward=reward,601                done=done,602                error=error,603            )604            observation = response["observation"]605 606        final_score = strict_score(rewards[-1] if rewards else STRICT_SCORE_EPSILON)607        success = final_score >= 0.5608    finally:609        log_end(success=success, steps=steps_taken, score=final_score, rewards=rewards)610    return final_score611 612 613def main() -> None:614    local_env_process = start_local_env_if_needed()615    client = OpenAI(base_url=API_BASE_URL, api_key=OPENAI_API_KEY)616    try:617        tasks = requests.get(f"{ENV_BASE_URL}/tasks", timeout=30).json()618        baseline_tasks = select_baseline_tasks(tasks)619        if len(baseline_tasks) != 3:620            raise RuntimeError("Expected canonical easy/medium/hard tasks to be available from /tasks.")621        for task in baseline_tasks:622            run_task(client, task)623    finally:624        if local_env_process is not None:625            local_env_process.terminate()626            try:627                local_env_process.wait(timeout=5)628            except subprocess.TimeoutExpired:629                local_env_process.kill()630 631 632if __name__ == "__main__":633    main()634