"""
Export a self-contained offline replay: JSON + HTML with embedded run data.
Uses the deterministic fallback one step per task (five scrubber steps).
Usage:
python scripts/export_replay.py
Writes:
runs/demo_fallback/replay.json
runs/demo_fallback/replay.html
"""
from __future__ import annotations
import base64
import json
import os
import sys
from datetime import datetime, timezone
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from baseline_runner import FALLBACK_SOLUTIONS, TASK_IDS # noqa: E402
from env import SQLOptimEnv # noqa: E402
from models import Action # noqa: E402
def _build_payload() -> dict:
env = SQLOptimEnv()
steps: list[dict] = []
for i, task_id in enumerate(TASK_IDS):
obs = env.reset(task_id=task_id)
sol = FALLBACK_SOLUTIONS[task_id]
action = Action(
suggestions=sol["suggestions"],
optimized_query=sol["optimized_query"],
summary=sol["summary"],
estimated_improvement=sol["estimated_improvement"],
approved=sol["approved"],
)
result = env.step(action)
ex = result.info.get("execution") or {}
steps.append(
{
"index": i,
"task_id": task_id,
"task_name": obs.task_name,
"difficulty": obs.difficulty,
"reward": round(result.reward.score, 4),
"breakdown": dict(result.reward.breakdown),
"original_sql": obs.sql_query,
"optimized_sql": action.optimized_query,
"last_execution": ex,
}
)
run_id = datetime.now(timezone.utc).strftime("demo_fallback_%Y%m%dT%H%M%SZ")
return {
"run_id": run_id,
"environment": "sql-optim-env",
"policy": "deterministic_fallback",
"steps": steps,
}
HTML_TEMPLATE = """
SQL Optim Env — Episode replay
SQL Query Optimization — replay
Reward
Last execution (DuckDB)
"""
def main() -> None:
out_dir = os.path.join(ROOT, "runs", "demo_fallback")
os.makedirs(out_dir, exist_ok=True)
payload = _build_payload()
json_path = os.path.join(out_dir, "replay.json")
with open(json_path, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2)
raw = json.dumps(payload, ensure_ascii=False).encode("utf-8")
b64 = base64.b64encode(raw).decode("ascii")
html = HTML_TEMPLATE.format(b64=b64)
html_path = os.path.join(out_dir, "replay.html")
with open(html_path, "w", encoding="utf-8") as f:
f.write(html)
print(f"Wrote {json_path}")
print(f"Wrote {html_path}")
if __name__ == "__main__":
main()