Spaces:
Running
Running
| """Teacher-trace generation and SFT export helpers.""" | |
| from __future__ import annotations | |
| import json | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Callable, Iterable | |
| from models import SDKAction, SDKObservation | |
| from server.environment import SDKSovereignEnvironment | |
| from server.policy_runtime import load_teacher_agents | |
| from server.prompts import SYSTEM_AUDITOR, SYSTEM_LEAD | |
| PolicyMap = dict[str, Callable[[SDKObservation], SDKAction]] | |
| class TraceStep: | |
| """Serializable single-turn trace record.""" | |
| episode_id: str | |
| repo_id: str | |
| role: str | |
| observation: dict[str, Any] | |
| action: dict[str, Any] | |
| reward: float | |
| done: bool | |
| success: bool | |
| policy_mode: str | |
| def serialize_model(value: Any) -> dict[str, Any]: | |
| """Convert dataclass-style values into plain dictionaries.""" | |
| if hasattr(value, "model_dump"): | |
| return value.model_dump() | |
| if hasattr(value, "__dict__"): | |
| return dict(value.__dict__) | |
| raise TypeError(f"Cannot serialize object of type {type(value)!r}") | |
| def run_episode( | |
| env: SDKSovereignEnvironment, | |
| policies: PolicyMap, | |
| *, | |
| policy_mode: str, | |
| initial_observation: SDKObservation, | |
| ) -> list[TraceStep]: | |
| """Run one full episode and capture all transitions.""" | |
| observation = initial_observation | |
| trace: list[TraceStep] = [] | |
| while True: | |
| acting_observation = observation | |
| action = policies[acting_observation.current_role](acting_observation) | |
| observation = env.step(action) | |
| success = bool(observation.done and env.state.test_results and all(env.state.test_results.values())) | |
| trace.append( | |
| TraceStep( | |
| episode_id=env.state.episode_id, | |
| repo_id=env.state.repo_id, | |
| role=action.role, | |
| observation=serialize_model(acting_observation), | |
| action=serialize_model(action), | |
| reward=observation.reward, | |
| done=observation.done, | |
| success=success, | |
| policy_mode=policy_mode, | |
| ) | |
| ) | |
| if observation.done: | |
| return trace | |
| def generate_teacher_traces( | |
| *, | |
| episodes_per_repo: int = 10, | |
| repos_root: Path | None = None, | |
| logs_root: Path | None = None, | |
| seed: int = 0, | |
| ) -> list[TraceStep]: | |
| """Generate successful teacher traces across all repos.""" | |
| env = SDKSovereignEnvironment(repos_root=repos_root, logs_root=logs_root, seed=seed) | |
| policies = load_teacher_agents() | |
| traces: list[TraceStep] = [] | |
| for repo_id in sorted(env.repos.keys()): | |
| for _ in range(episodes_per_repo): | |
| observation = env.reset(repo_id=repo_id) | |
| traces.extend(run_episode(env, policies, policy_mode="teacher", initial_observation=observation)) | |
| return traces | |
| def write_jsonl(records: Iterable[dict[str, Any]], output_path: Path) -> None: | |
| """Write records as JSONL.""" | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| with output_path.open("w", encoding="utf-8") as handle: | |
| for record in records: | |
| handle.write(json.dumps(record) + "\n") | |
| def export_traces_jsonl(traces: Iterable[TraceStep], output_path: Path) -> None: | |
| """Persist trace records for later SFT / RL preprocessing.""" | |
| write_jsonl((trace.__dict__ for trace in traces), output_path) | |
| def build_sft_messages(role: str, observation: dict[str, Any], action: dict[str, Any]) -> list[dict[str, str]]: | |
| """Convert one trace item into a chat-format SFT sample.""" | |
| system_prompt = SYSTEM_AUDITOR if role == "auditor" else SYSTEM_LEAD | |
| user_payload = json.dumps(observation, ensure_ascii=True) | |
| assistant_payload = json.dumps(action, ensure_ascii=True) | |
| return [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_payload}, | |
| {"role": "assistant", "content": assistant_payload}, | |
| ] | |
| def export_sft_jsonl( | |
| traces: Iterable[TraceStep], | |
| output_path: Path, | |
| *, | |
| success_only: bool = True, | |
| ) -> None: | |
| """Export successful teacher traces as chat-format SFT rows.""" | |
| rows = [] | |
| for trace in traces: | |
| if success_only and not trace.success: | |
| continue | |
| rows.append( | |
| { | |
| "messages": build_sft_messages(trace.role, trace.observation, trace.action), | |
| "role": trace.role, | |
| "repo_id": trace.repo_id, | |
| "episode_id": trace.episode_id, | |
| "reward": trace.reward, | |
| "success": trace.success, | |
| "policy_mode": trace.policy_mode, | |
| } | |
| ) | |
| write_jsonl(rows, output_path) | |