import asyncio import json import math import os import textwrap from typing import Any from openai import OpenAI from openenv.core.generic_client import GenericEnvClient API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") HF_TOKEN = os.getenv("HF_TOKEN") LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") OPENENV_BASE_URL = os.getenv("OPENENV_BASE_URL") TASK_NAME = os.getenv("TASK_NAME", "engineer-manager") BENCHMARK = os.getenv("BENCHMARK", "openenv") MAX_STEPS = int(os.getenv("MAX_STEPS", "32")) TEMPERATURE = float(os.getenv("TEMPERATURE", "0.1")) MAX_TOKENS = int(os.getenv("MAX_TOKENS", "120")) SYSTEM_PROMPT = textwrap.dedent( """ You are selecting actions for an environment that simulates an engineer-manager workday. Return exactly one compact JSON object with integer keys: {"target_slot": , "operation": } Operations: 0 = idle 1 = schedule work at target_slot 2 = reschedule a meeting at target_slot 3 = toggle mute comms Goals: - Maximize sustained deep work flow_score. - Avoid unnecessary social_debt and calendar_churn. - Prefer scheduling work into future empty slots. - Use reschedule_meeting only when it clearly helps. - Toggle mute comms early if distractions are high and it is currently off. Rules: - target_slot must be within the timeline bounds. - Return JSON only. No markdown. No explanation. """ ).strip() def _require_env(name: str, value: str | None) -> str: if value: return value raise RuntimeError(f"Missing required environment variable: {name}") def _sanitize_field(value: Any) -> str: text = str(value) return text.replace("\r", " ").replace("\n", " ").strip() def log_start(task: str, env: str, model: str) -> None: print( f"[START] task={_sanitize_field(task)} env={_sanitize_field(env)} model={_sanitize_field(model)}", flush=True, ) def log_step( step: int, action: str, reward: float, done: bool, error: str | None, ) -> None: error_text = "null" if error in (None, "") else _sanitize_field(error) print( f"[STEP] step={step} action={_sanitize_field(action)} reward={reward:.2f} " f"done={str(done).lower()} error={error_text}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None: rewards_text = ",".join(f"{reward:.2f}" for reward in rewards) print( f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_text}", flush=True, ) def estimate_max_flow_score(timeline: list[int]) -> float: slot_count = len(timeline) if slot_count <= 0: return 1.0 hours = slot_count * 0.5 return max(1.0, hours * hours) def normalize_score(total_reward: float, observation: dict[str, Any]) -> float: timeline = observation.get("timeline") or [] max_score = estimate_max_flow_score(timeline) normalized = total_reward / max_score return min(1.0, max(0.0, normalized)) def first_future_slot(observation: dict[str, Any], kind: int) -> int | None: timeline = observation.get("timeline") or [] current_slot = int(observation.get("current_slot", 0)) for index in range(current_slot, len(timeline)): if int(timeline[index]) == kind: return index return None def first_future_empty_slot(observation: dict[str, Any]) -> int | None: return first_future_slot(observation, 0) def build_user_prompt( step: int, observation: dict[str, Any], rewards: list[float], history: list[str], ) -> str: timeline = observation.get("timeline") or [] metadata = observation.get("metadata") or {} return textwrap.dedent( f""" step={step} current_slot={int(observation.get("current_slot", 0))} current_time={observation.get("current_time", "unknown")} mute_comms={bool(observation.get("mute_comms", False))} distraction_risk={float(observation.get("distraction_risk", 0.0))} flow_score={float(observation.get("flow_score", 0.0)):.2f} social_debt={float(observation.get("social_debt", 0.0)):.2f} calendar_churn={int(observation.get("calendar_churn", 0))} recovery_state={int(observation.get("recovery_state", 0))} timeline={timeline} task_buffer={json.dumps(observation.get("task_buffer", []), separators=(",", ":"))} last_rewards={",".join(f"{reward:.2f}" for reward in rewards[-5:]) or "none"} recent_history={json.dumps(history[-5:])} last_metadata={json.dumps(metadata, separators=(",", ":"))} Choose the single next action. """ ).strip() def choose_fallback_action(observation: dict[str, Any]) -> dict[str, int]: current_slot = int(observation.get("current_slot", 0)) distraction_risk = float(observation.get("distraction_risk", 0.0)) mute_comms = bool(observation.get("mute_comms", False)) if distraction_risk >= 0.2 and not mute_comms: return {"target_slot": current_slot, "operation": 3} empty_slot = first_future_empty_slot(observation) if empty_slot is not None and observation.get("task_buffer"): return {"target_slot": empty_slot, "operation": 1} meeting_slot = first_future_slot(observation, 2) if meeting_slot is not None and current_slot <= meeting_slot: return {"target_slot": meeting_slot, "operation": 2} return {"target_slot": current_slot, "operation": 0} def coerce_action(raw_text: str, observation: dict[str, Any]) -> dict[str, int]: timeline = observation.get("timeline") or [] max_slot = max(0, len(timeline) - 1) fallback = choose_fallback_action(observation) try: data = json.loads(raw_text) target_slot = int(data["target_slot"]) operation = int(data["operation"]) except Exception: return fallback if operation not in {0, 1, 2, 3}: return fallback target_slot = min(max(target_slot, 0), max_slot) return {"target_slot": target_slot, "operation": operation} def get_model_action( client: OpenAI, step: int, observation: dict[str, Any], rewards: list[float], history: list[str], ) -> dict[str, int]: user_prompt = build_user_prompt(step, observation, rewards, history) try: completion = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=TEMPERATURE, max_tokens=MAX_TOKENS, ) content = (completion.choices[0].message.content or "").strip() return coerce_action(content, observation) except Exception: return choose_fallback_action(observation) async def create_env() -> GenericEnvClient: if OPENENV_BASE_URL: env = GenericEnvClient(base_url=OPENENV_BASE_URL) await env.connect() return env image_name = _require_env("LOCAL_IMAGE_NAME", LOCAL_IMAGE_NAME) return await GenericEnvClient.from_docker_image(image_name) async def main() -> None: api_key = _require_env("HF_TOKEN", HF_TOKEN) client = OpenAI(base_url=API_BASE_URL, api_key=api_key) env = None rewards: list[float] = [] history: list[str] = [] steps_taken = 0 success = False score = 0.0 log_start(TASK_NAME, BENCHMARK, MODEL_NAME) try: env = await create_env() result = await env.reset() observation = dict(result.observation) for step in range(1, MAX_STEPS + 1): if result.done: break action = get_model_action(client, step, observation, rewards, history) result = await env.step(action) observation = dict(result.observation) reward = float(result.reward or 0.0) done = bool(result.done) metadata = observation.get("metadata") or {} error = metadata.get("last_action_error") rewards.append(reward) steps_taken = step action_text = ( f"target_slot={int(action['target_slot'])},operation={int(action['operation'])}" ) log_step(step, action_text, reward, done, error) history.append( f"step={step} action={action_text} reward={reward:.2f} " f"flow={float(observation.get('flow_score', 0.0)):.2f} " f"debt={float(observation.get('social_debt', 0.0)):.2f}" ) if done: break total_reward = math.fsum(rewards) score = normalize_score(total_reward, observation if "observation" in locals() else {}) score = round(score, 2) success = score > 0.0 finally: if env is not None: try: await env.close() except Exception: pass log_end(success=success, steps=steps_taken, score=score, rewards=rewards) if __name__ == "__main__": asyncio.run(main())