source/study/space/app.py

"""Add an aggregate study view while preserving the pinned Trackio application."""

from contextlib import asynccontextmanager
import csv
import io
import json
import os
from pathlib import Path
import sys

from starlette.applications import Starlette
from starlette.responses import FileResponse, HTMLResponse, JSONResponse
from starlette.routing import Mount, Route


HERE = Path(__file__).resolve().parent
# In the Space the reviewed harness lives under source/; locally it lives above us.
STUDY = HERE / "source" / "study" if (HERE / "source" / "study").is_dir() else HERE.parent
sys.path.insert(0, str(STUDY))
from publish_results import AXES, FILES, OPTIONAL_FILES, PREFIX, prepare_payloads


PAGE = """<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>Beam + Pi · ProgramBench</title>
<style>
body{margin:0;background:#fff;color:#222;font:16px/1.5 system-ui,sans-serif}
main{max-width:1200px;margin:auto;padding:28px 24px}
h1{font-size:clamp(25px,4vw,36px);line-height:1.2;margin:0 0 20px}
#chart{height:640px}
@media(max-width:650px){main{padding:20px 12px}#chart{height:520px}}
</style><script src="/study/plotly.js"></script></head>
<body><main>
<h1>Beam + Pi · ProgramBench</h1>
<div id="chart" aria-label="Mean hidden-test fraction versus latency chart"></div>
</main><script>
// Responsive layout start
function fitFigure(figure,width){
 const fitted=JSON.parse(JSON.stringify(figure)), layout=fitted.layout;
 const narrow=width<650;
 const labels={single:'Single agent',peers:'Five-agent team',async:'Async subagents'};
 fitted.data=fitted.data.filter(trace=>trace.yaxis!=='y2');
 for(const trace of fitted.data){
  const suffix=trace.name.match(/ · (pending|clock unavailable)$/)?.[0]||'';
  if(labels[trace.legendgroup]) trace.name=labels[trace.legendgroup]+suffix;
  trace.hovertemplate='%{x:,.1f}s · mean score %{y:.3f}<extra>%{fullData.name}</extra>';
 }
 layout.title={text:''};layout.annotations=[];
 layout.height=narrow?520:640;
 layout.margin={t:narrow?100:65,b:70,l:narrow?60:75,r:20};
 layout.legend={orientation:narrow?'v':'h',x:0,y:1.04,yanchor:'bottom',font:{size:narrow?12:14}};
 layout.yaxis.domain=[0,1];layout.xaxis.anchor='y';delete layout.yaxis2;
 layout.xaxis.title={...layout.xaxis.title,font:{size:narrow?11:14}};
 layout.yaxis.title={...layout.yaxis.title,text:narrow?'Hidden tests passed':'Fraction of hidden tests passed',font:{size:narrow?12:14}};
 layout.xaxis.tickfont={size:narrow?11:12};layout.yaxis.tickfont={size:narrow?11:12};
 if(narrow){
  const stride=Math.max(1,Math.ceil(layout.xaxis.tickvals.length/4));
  layout.xaxis.ticktext=layout.xaxis.ticktext.filter((_,i)=>i%stride===0);
  layout.xaxis.tickvals=layout.xaxis.tickvals.filter((_,i)=>i%stride===0);
 }
 return {figure:fitted,config:{responsive:true,displaylogo:false,displayModeBar:false}};
}
// Responsive layout end
const chart=document.getElementById('chart');let publishedFigure=null,chartWidth=0,requestNumber=0;
async function renderFigure(){
 if(!publishedFigure) return;
 chartWidth=Math.round(chart.clientWidth);const view=fitFigure(publishedFigure,chartWidth);
 chart.style.height=view.figure.layout.height+'px';
 await Plotly.react(chart,view.figure.data,view.figure.layout,view.config);
}
new ResizeObserver(()=>{if(Math.round(chart.clientWidth)!==chartWidth) renderFigure();}).observe(chart);
async function refresh(){
 const current=++requestNumber;
 try{
  const response=await fetch('/study/data?axis=reference_seconds',{cache:'no-store'});
  if(!response.ok) throw new Error('unavailable');
  const payload=await response.json();if(current!==requestNumber) return;
  publishedFigure=payload.figure;await renderFigure();
 }catch(error){if(current===requestNumber && !publishedFigure) chart.textContent='Results unavailable';}
}
refresh();setInterval(refresh,60000);
</script></body></html>"""


def parse_curves(payload):
    rows = []
    for row in csv.DictReader(io.StringIO(payload.decode())):
        rows.append({**row, "time_seconds": float(row["time_seconds"]),
                     "mean_hidden_test_fraction": float(row["mean_hidden_test_fraction"]) if row["mean_hidden_test_fraction"] else None,
                     "clock_valid": row["clock_valid"] == "True",
                     **{key: int(row[key]) for key in ("tasks_total", "tasks_with_valid_grade", "episodes_configured")}})
    return rows


def study_snapshot(directory, axis):
    from plot_results import build_plotly

    directory = Path(directory)
    if not all((directory / name).is_file() for name in FILES):
        return {"available": False, "figure": build_plotly(None, [], axis=axis), "downloads": [],
                "message": "No published benchmark grades yet. Reference calibration and infrastructure checks are tracked separately."}
    payloads = prepare_payloads(directory)
    summary = json.loads(payloads["summary.json"])
    if "recovery-report.json" in payloads:
        summary["recovery_report"] = json.loads(payloads["recovery-report.json"])
    status = json.loads(payloads["study-status.json"])
    scope = summary.get("analysis_scope")
    included = [row for row in summary["episodes"] if not scope or row["task"] in scope["included_tasks"]]
    included_finals = sum(row["has_valid_final_grade"] for row in included)
    included_grades = sum(row["has_valid_grade"] for row in included)
    subset_ready = all(row["has_valid_final_grade"] and row["clock_valid"]
                       and row["status"] in {"completed", "budget_exhausted"} for row in included)
    grades = sum(row["has_valid_grade"] for row in summary["episodes"])
    final_grades = sum(row["has_valid_final_grade"] for row in summary["episodes"])
    invalid = sum(row["invalid_grade_count"] for row in summary["episodes"])
    current_invalid = sum(row["current_attempt_invalid_grade_count"] for row in summary["episodes"])
    recovered = (any(row["attempt_count"] > 1 for row in summary["episodes"])
                 or bool(summary.get("recovery_report", {}).get("entries")))
    message = (("Recovered pilot. " if recovered else "") +
               f"Study status: {status['status']}. Final artifacts graded: {final_grades}/15; any valid checkpoint: {grades}/15. "
               f"Invalid grades: {invalid - current_invalid} historical + {current_invalid} current. "
               f"Aggregates generated {summary['generated_at_utc']}.")
    message += " Historical counts refer to archived solver attempts."
    if scope:
        message = (f"Post-hoc four-task subset ({'complete' if subset_ready else 'provisional'}): jq, xsv, FIGlet, hexyl. "
                   f"Included final artifacts graded: {included_finals}/12; any valid checkpoint: {included_grades}/12. "
                   "Ripgrep is excluded from all three conditions due to slow saved-artifact grading; "
                   "all 15 original episode records and all token charges remain in the evidence and accounting. "
                   "Original 15-episode evidence: " + message)
    if summary.get("recovery_report", {}).get("entries"):
        message += f" Reviewed artifact corrections: {len(summary['recovery_report']['entries'])}; original hashes and scores are retained in recovery-report.json."
    if recovered:
        message += " Original budgets and all prior attempt solver/reference time remain counted."
    if any(row["kind"] == "candidate_collection_failed" for row in summary.get("recovery_report", {}).get("entries", [])):
        message += " Reviewed candidate collection omissions are counted as failed tests against the fixed official mask."
    clock = summary["reference_clock"]
    note = (f"Operational reference clock: decode {clock['decode_tokens_per_second']:g} tokens/s; "
            f"prefill {clock['prefill_tokens_per_second']:g} tokens/s. Primary prefill counts estimated newly entering context; "
            "the sensitivity view counts all API prompt tokens. Only completed operations contribute; causal handoffs are included. "
            "These assumed rates do not reproduce Anthropic's exact latency calculation.")
    from publish_results import FIGURE_FILES
    downloads = [{"label": name.rsplit(".", 1)[-1].upper(), "url": f"/study/download/{name}"}
                 for name in FIGURE_FILES if name.startswith(f"score-vs-latency-{axis}.") and (directory / name).is_file()]
    downloads.extend({"label": name + (" · original 5-task evidence" if scope else ""),
                      "url": f"/study/download/{name}"} for name in ("curves.csv", "summary.json"))
    downloads.extend({"label": name, "url": f"/study/download/{name}"} for name in OPTIONAL_FILES if name in payloads)
    return {"available": True, "figure": build_plotly(summary, parse_curves(payloads["curves.csv"]), axis=axis),
            "message": message, "clock_note": note, "downloads": downloads,
            "analysis_scope": scope, "tasks_total": len(scope["included_tasks"]) if scope else 5,
            "included_episodes": len(included), "excluded_episodes": len(summary["episodes"]) - len(included),
            "analysis_status": "complete" if subset_ready else "provisional",
            "episodes": [{**{key: row[key] for key in ("run_id", "status", "final_score", "has_valid_grade", "has_valid_final_grade", "attempt_count", "current_attempt_invalid_grade_count", "api_tokens_charged")},
                          "analysis_included": not scope or row["task"] in scope["included_tasks"],
                          "historical_invalid_grade_count": row["invalid_grade_count"] - row["current_attempt_invalid_grade_count"]}
                         for row in summary["episodes"]]}


def create_app(trackio_app, *, results_dir=None):
    """Wrap, rather than remount or replace, Trackio's existing paths and lifespan."""
    from publish_results import FIGURE_FILES

    directory = Path(results_dir) if results_dir is not None else HERE / PREFIX

    async def page(request):
        return HTMLResponse(PAGE)

    async def data(request):
        axis = request.query_params.get("axis", "reference_seconds")
        if axis not in AXES:
            return JSONResponse({"error": "Unknown latency measure"}, status_code=400)
        try:
            return JSONResponse(study_snapshot(directory, axis), headers={"Cache-Control": "no-store"})
        except Exception:
            # Repository contents are inputs; never leak paths or exception text.
            return JSONResponse({"error": "Published aggregates could not be validated"}, status_code=503)

    async def plotly_script(request):
        import plotly
        return FileResponse(Path(plotly.__file__).parent / "package_data" / "plotly.min.js",
                            media_type="application/javascript")

    async def download(request):
        name = request.path_params["name"]
        if name not in (*FILES, *OPTIONAL_FILES, *FIGURE_FILES):
            return JSONResponse({"error": "Unknown aggregate file"}, status_code=404)
        path = directory / name
        if path.is_symlink() or not path.is_file():
            return JSONResponse({"error": "Aggregate file unavailable"}, status_code=404)
        if name in (*FILES, *OPTIONAL_FILES):
            try:
                payload = prepare_payloads(directory)[name]
                from starlette.responses import Response
                return Response(payload, media_type="text/csv" if name.endswith(".csv") else "application/json",
                                headers={"Content-Disposition": f'attachment; filename="{name}"'})
            except Exception:
                return JSONResponse({"error": "Published aggregates could not be validated"}, status_code=503)
        return FileResponse(path, filename=name)

    @asynccontextmanager
    async def lifespan(app):
        async with trackio_app.router.lifespan_context(trackio_app):
            yield

    return Starlette(routes=[Route("/study", page), Route("/study/", page), Route("/study/data", data),
                             Route("/study/plotly.js", plotly_script), Route("/study/download/{name}", download),
                             Mount("/", app=trackio_app)], lifespan=lifespan)


if __name__ == "__main__":
    import uvicorn
    from trackio.server import build_starlette_app_only

    backend, _write_token = build_starlette_app_only(mcp_server=os.environ.get("GRADIO_MCP_SERVER") == "True")
    uvicorn.run(create_app(backend), host="0.0.0.0", port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")))