"""Environment logic for deterministic LLM release-readiness tasks.""" from __future__ import annotations import os import re from copy import deepcopy from typing import Any, Dict, Tuple from uuid import uuid4 try: from openenv.core.env_server.interfaces import Environment except ImportError: from openenv.core.env_server.interfaces import Environment try: from model_release_env.models import ( ModelReleaseAction, ModelReleaseObservation, ModelReleaseState, ) except ImportError: from models import ModelReleaseAction, ModelReleaseObservation, ModelReleaseState DEFAULT_TASK = os.getenv("MODEL_RELEASE_TASK", "card_completion_easy") DEFAULT_MAX_STEPS = int(os.getenv("MODEL_RELEASE_MAX_STEPS", "8")) def _normalize(value: str) -> str: return re.sub(r"\s+", " ", value.strip().lower()) def _check_rule(rule: Dict[str, Any], value: Any) -> bool: normalized_value = _normalize(str(value or "")) rule_type = rule["type"] if rule_type == "exact": return normalized_value == _normalize(rule["expected"]) if rule_type == "contains_all": return all(token in normalized_value for token in rule["tokens"]) return False TASKS: Dict[str, Dict[str, Any]] = { "card_completion_easy": { "difficulty": "easy", "goal": "Complete the missing release-card fields before shipping the model.", "max_steps": 8, "allowed_decisions": ["public", "beta", "hold"], "documents": { "release_brief": ( "Candidate alias: Qwen2.5-7B-Instruct. Context window: 32768 tokens. " "Training scope in the draft is acceptable and does not need edits." ), "eval_sheet": ( "Evaluation summary for the candidate release: gsm8k=0.78; math500=0.61; " "aime24=0.18. Use this exact compact format in the release package." ), "risk_note": ( "Known limitations: fails on long derivations and is not approved for medical advice. " "Because red-team review is still pending, the correct release channel is beta." ), }, "document_summaries": { "release_brief": "Base model identity and draft scope.", "eval_sheet": "Benchmark numbers approved for publication.", "risk_note": "Known limitations and launch recommendation.", }, "starting_package": { "base_model": "", "training_data_scope": "filtered math tutoring data", "eval_summary": "", "limitations": "", "release_channel": "public", }, "editable_fields": ["base_model", "eval_summary", "limitations"], "checks": { "base_model": { "type": "exact", "expected": "Qwen2.5-7B-Instruct", "weight": 0.25, }, "eval_summary": { "type": "exact", "expected": "gsm8k=0.78; math500=0.61; aime24=0.18", "weight": 0.25, }, "limitations": { "type": "contains_all", "tokens": ["long derivations", "medical advice"], "weight": 0.25, }, "release_channel": { "type": "exact", "expected": "beta", "weight": 0.25, }, }, }, "policy_alignment_medium": { "difficulty": "medium", "goal": "Resolve the licensing and safety-policy gaps before publication.", "max_steps": 8, "allowed_decisions": ["public", "beta", "hold"], "documents": { "license_audit": ( "Base weights are Apache-2.0. The published model card must keep the license field as apache-2.0." ), "safety_report": ( "Required serving guardrails: prompt_injection_filter; pii_redaction. " "Both must appear in the safety_guardrails field." ), "policy_note": ( "Blocked use cases: legal advice and medical advice. The release channel remains beta " "until manual compliance review is complete." ), }, "document_summaries": { "license_audit": "Approved license value.", "safety_report": "Mandatory runtime guardrails.", "policy_note": "Disallowed uses and launch channel.", }, "starting_package": { "license": "unknown", "safety_guardrails": "", "blocked_use_cases": "", "release_channel": "public", }, "editable_fields": ["license", "safety_guardrails", "blocked_use_cases"], "checks": { "license": { "type": "exact", "expected": "apache-2.0", "weight": 0.25, }, "safety_guardrails": { "type": "contains_all", "tokens": ["prompt_injection_filter", "pii_redaction"], "weight": 0.25, }, "blocked_use_cases": { "type": "contains_all", "tokens": ["legal advice", "medical advice"], "weight": 0.25, }, "release_channel": { "type": "exact", "expected": "beta", "weight": 0.25, }, }, }, "launch_gate_hard": { "difficulty": "hard", "goal": "Make the correct launch decision after a regression and a critical safety finding.", "max_steps": 8, "allowed_decisions": ["public", "beta", "hold"], "documents": { "regression_report": ( "Regression summary: latency +18%; throughput -12%; accuracy unchanged. " "Known regressions should list latency and throughput." ), "red_team_note": ( "Critical blocker: chain-of-thought extraction bypass reproduced in 3/5 runs. " "This must appear in the launch_blockers field." ), "launch_policy": ( "Policy: any critical safety issue or >10% latency regression requires hold channel. " "Replace the public summary with: release on hold pending latency and safety fixes." ), }, "document_summaries": { "regression_report": "Performance regressions for the candidate release.", "red_team_note": "Critical red-team finding.", "launch_policy": "Mandatory launch-gate rule.", }, "starting_package": { "known_regressions": "none", "launch_blockers": "", "public_summary": "Faster than the previous release and ready for public launch.", "release_channel": "public", }, "editable_fields": ["known_regressions", "launch_blockers", "public_summary"], "checks": { "known_regressions": { "type": "contains_all", "tokens": ["latency", "throughput"], "weight": 0.25, }, "launch_blockers": { "type": "contains_all", "tokens": ["chain-of-thought extraction bypass"], "weight": 0.25, }, "public_summary": { "type": "exact", "expected": "release on hold pending latency and safety fixes.", "weight": 0.25, }, "release_channel": { "type": "exact", "expected": "hold", "weight": 0.25, }, }, }, } class ModelReleaseEnvironment( Environment[ModelReleaseAction, ModelReleaseObservation, ModelReleaseState] ): """Release-readiness environment for deterministic LLM launch workflows.""" SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__(self, max_steps: int | None = None, default_task: str | None = None): super().__init__() self._configured_max_steps = max_steps or DEFAULT_MAX_STEPS self._default_task = default_task or DEFAULT_TASK self._task_name = self._default_task if self._default_task in TASKS else next(iter(TASKS)) self._task_spec: Dict[str, Any] = {} self._package: Dict[str, Any] = {} self._visible_documents: Dict[str, str] = {} self._inspected_documents: set[str] = set() self._last_action_error: str | None = None self._score_by_check: Dict[str, bool] = {} self._task_score = 0.0 self._state = ModelReleaseState( episode_id=str(uuid4()), step_count=0, task_name=self._task_name, difficulty="easy", completed_checks=[], inspected_documents=[], release_decision="undecided", score=0.0, ) self.reset(task_name=self._task_name) def reset( self, seed: int | None = None, episode_id: str | None = None, task_name: str | None = None, **kwargs: Any, ) -> ModelReleaseObservation: del seed, kwargs selected_task = task_name or self._default_task if selected_task not in TASKS: selected_task = next(iter(TASKS)) self._task_name = selected_task self._task_spec = deepcopy(TASKS[selected_task]) self._package = deepcopy(self._task_spec["starting_package"]) self._visible_documents = {} self._inspected_documents = set() self._last_action_error = None self._task_score, self._score_by_check = self._compute_score() self._state = ModelReleaseState( episode_id=episode_id or str(uuid4()), step_count=0, task_name=self._task_name, difficulty=self._task_spec["difficulty"], completed_checks=self._completed_checks(), inspected_documents=[], release_decision=self._package.get("release_channel", "undecided"), score=self._task_score, ) return self._build_observation(reward=0.0, done=False) def step(self, action: ModelReleaseAction, **kwargs: Any) -> ModelReleaseObservation: del kwargs self._state.step_count += 1 self._last_action_error = None reward = 0.0 score_before = self._task_score package_before = deepcopy(self._package) done = False metadata: Dict[str, Any] = {"action_type": action.action_type} if action.action_type == "inspect": reward = self._handle_inspect(action) elif action.action_type == "set_field": reward = self._handle_set_field(action, score_before, package_before) elif action.action_type == "set_decision": reward = self._handle_set_decision(action, score_before, package_before) elif action.action_type == "submit": self._task_score, self._score_by_check = self._compute_score() reward = round(self._task_score, 2) done = True metadata["submitted"] = True else: self._last_action_error = f"unsupported action_type: {action.action_type}" reward = -0.05 self._task_score, self._score_by_check = self._compute_score() if self._state.step_count >= self._task_spec["max_steps"]: done = True self._state.completed_checks = self._completed_checks() self._state.inspected_documents = sorted(self._inspected_documents) self._state.release_decision = self._package.get("release_channel", "undecided") self._state.score = self._task_score observation = self._build_observation(reward=reward, done=done) observation.metadata["score_before"] = round(score_before, 2) observation.metadata["score_after"] = round(self._task_score, 2) observation.metadata["package_changed"] = package_before != self._package observation.metadata.update(metadata) return observation @property def state(self) -> ModelReleaseState: return self._state def close(self) -> None: return None def _handle_inspect(self, action: ModelReleaseAction) -> float: document_name = action.target.strip() documents = self._task_spec["documents"] if document_name not in documents: self._last_action_error = f"unknown document: {document_name}" return -0.05 if document_name in self._inspected_documents: return 0.0 self._inspected_documents.add(document_name) self._visible_documents[document_name] = documents[document_name] return 0.04 def _handle_set_field( self, action: ModelReleaseAction, score_before: float, package_before: Dict[str, Any], ) -> float: target = action.target.strip() value = action.value.strip() if target not in self._task_spec["editable_fields"]: self._last_action_error = f"field is not editable in this task: {target}" return -0.05 if not value: self._last_action_error = f"empty value for field: {target}" return -0.05 self._package[target] = value new_score, _ = self._compute_score() if _normalize(str(package_before.get(target, ""))) == _normalize(value): return 0.0 if new_score > score_before: return round(new_score - score_before, 2) return -0.05 def _handle_set_decision( self, action: ModelReleaseAction, score_before: float, package_before: Dict[str, Any], ) -> float: decision = action.value.strip().lower() if decision not in self._task_spec["allowed_decisions"]: self._last_action_error = f"invalid decision: {decision}" return -0.05 self._package["release_channel"] = decision new_score, _ = self._compute_score() if _normalize(str(package_before.get("release_channel", ""))) == _normalize(decision): return 0.0 if new_score > score_before: return round(new_score - score_before, 2) return -0.05 def _compute_score(self) -> Tuple[float, Dict[str, bool]]: matched: Dict[str, bool] = {} total = 0.0 for name, rule in self._task_spec["checks"].items(): value = self._package.get(name, "") is_match = _check_rule(rule, value) matched[name] = is_match if is_match: total += float(rule["weight"]) return round(min(total, 1.0), 2), matched def _completed_checks(self) -> list[str]: return [name for name, passed in self._score_by_check.items() if passed] def _build_observation(self, reward: float, done: bool) -> ModelReleaseObservation: critical_gaps = [ name for name, satisfied in self._score_by_check.items() if not satisfied ] return ModelReleaseObservation( task_name=self._task_name, difficulty=self._task_spec["difficulty"], goal=self._task_spec["goal"], document_index=deepcopy(self._task_spec["document_summaries"]), visible_documents=deepcopy(self._visible_documents), package_snapshot=deepcopy(self._package), checklist_status=deepcopy(self._score_by_check), critical_gaps=critical_gaps, available_fields=list(self._task_spec["editable_fields"]), available_decisions=list(self._task_spec["allowed_decisions"]), inspected_documents=sorted(self._inspected_documents), remaining_steps=max(self._task_spec["max_steps"] - self._state.step_count, 0), last_action_error=self._last_action_error, score=self._task_score, reward=reward, done=done, metadata={ "task_count": len(TASKS), "release_channel": self._package.get("release_channel", "undecided"), }, )