From e35451330cdd3e3ca05ca4b025609d6bc344945b Mon Sep 17 00:00:00 2001 From: Masoud Masoumi Date: Sun, 6 Sep 2026 21:06:39 -0400 Subject: [PATCH] v0.3.0: warning-level quality gates + run_comparison, bug fixes Bug fixes: - simulate() options render Modelica booleans (true/false) instead of Python literals - summarize_for_llm() suppression count now counts only actually-suppressed records - check_model() no longer duplicates diagnostics from the return text - diagnostic parser accepts omc single-position locations (no end range) New features: - Warning-level quality gates: warning_gate_complaints() matches warnings (un/over-specified initial conditions, unit inconsistencies, over/under-determined systems); AgentLoop(warning_gate=...) reports stage 'quality' attempts and feeds complaints into the fix prompt; run_ladder threads the gate through the benchmark (opt-in) - run_comparison(): multi-run variance + cross-model comparison with per-model/per-task pass rates, attempts mean/stdev, transcripts under transcripts//rep/, aggregated in comparison.json Tests: 101 unit + 4 integration (was 88 + 4). Version bumped to 0.3.0. --- README.md | 56 ++++++++++++-- omagent/__init__.py | 14 ++-- omagent/errors.py | 65 ++++++++++++++-- omagent/loop.py | 30 ++++++-- omagent/runner.py | 95 ++++++++++++++++++++++- omagent/session.py | 17 ++++- pyproject.toml | 2 +- tests/test_quality.py | 118 ++++++++++++++++++++++++++++ tests/test_runner_comparison.py | 131 ++++++++++++++++++++++++++++++++ 9 files changed, 498 insertions(+), 30 deletions(-) create mode 100644 tests/test_quality.py create mode 100644 tests/test_runner_comparison.py diff --git a/README.md b/README.md index 7e87572..8e58b34 100644 --- a/README.md +++ b/README.md @@ -32,10 +32,12 @@ Key capabilities: transcript capture, in the format the OpenModelica benchmark discussion ([OpenModelica#15385](https://github.com/OpenModelica/OpenModelica/issues/15385)) calls for +- **Multi-run benchmarking** — repeat runs for variance measurement and + side-by-side cross-model comparison, aggregated into a single report - **LLM-backend-agnostic** — the loop depends on a one-method protocol; adapters ship for Anthropic and for any OpenAI-compatible endpoint (Ollama, LM Studio, llama.cpp, vLLM — local open-weight models included) -- **Tested** — 83 unit tests run without OpenModelica installed; 4 +- **Tested** — 101 unit tests run without OpenModelica installed; 4 integration tests validate against a live omc ## Installation @@ -133,6 +135,53 @@ is not, (5) multi-domain electro-mechanical. Per-task JSON transcripts (attempt history, diagnostics, code, LLM rounds) land in `transcripts/`, with `summary.json` aggregating results. +### Compare models with repeated runs + +```python +from omagent import OMSession, run_comparison +from omagent.llm import ClaudeLLM, OpenAICompatLLM + +comparison = run_comparison( + session_factory=OMSession, # fresh omc session per run + llm_factories={ + "claude-sonnet": lambda: ClaudeLLM(model="claude-sonnet-4-6"), + "local-qwen": lambda: OpenAICompatLLM(model="qwen2.5-coder:14b"), + }, + repeats=3, # runs per model for variance +) +``` + +Each (model, task, run) triple runs in isolation, so variance across repeats +reflects LLM/omc nondeterminism rather than state contamination. Transcripts +land in `transcripts//rep/`; `comparison.json` aggregates pass +rates per model and per task, plus mean/spread of attempts and wall time. +Use `--tasks`/`--max-tier` equivalents via `task_ids`/`max_tier`, and +`verbose=True` for progress and a final table. + +### Warning-level quality gates + +Some omc diagnostics come as warnings yet mean the model is sloppy — +under/over-specified initial conditions, inconsistent units, over-determined +systems. Quality gates turn those into verifier-style complaints that feed +the fix loop, without outright failing the operation: + +```python +from omagent import AgentLoop, OMSession, warning_gate_complaints + +loop = AgentLoop( + OMSession(), ClaudeLLM(), max_attempts=4, + verifier=my_verifier, + warning_gate=warning_gate_complaints, # opt-in; None by default +) +``` + +Gated attempts report stage `"quality"` and the gate complaint is appended +to the fix prompt's structured feedback. `run_ladder(..., warning_gate=...)` +threads the gate through the benchmark so scores can be produced under +either strictness. Custom gates are just callables over +`list[Diagnostic] -> Optional[str]`; `WARNING_GATE_PATTERNS` is the default +rule table you can extend. + ### Use pieces standalone ```python @@ -159,7 +208,7 @@ omagent/ tasks.py # benchmark task ladder definitions runner.py # ladder execution + transcript persistence examples/ # first_run.py, run_ladder.py -tests/ # 83 unit + 4 integration tests +tests/ # 101 unit + 4 integration tests ``` ## Design notes @@ -176,9 +225,6 @@ tests/ # 83 unit + 4 integration tests ## Roadmap -- Warning-level quality gates (e.g. treat "initial conditions over - specified" as a verifier complaint) -- Multi-run variance measurement and cross-model comparison in the runner - Optional MCP tool surface, composing with OMEdit's built-in MCP server - More ladder tiers targeting thermal/fluid domains and third-party libraries diff --git a/omagent/__init__.py b/omagent/__init__.py index dd9b7cf..f59527d 100644 --- a/omagent/__init__.py +++ b/omagent/__init__.py @@ -2,20 +2,24 @@ from .errors import ( Diagnostic, Kind, Severity, - classify, parse_error_string, parse_ompython_exception, parse_simulation_messages, summarize_for_llm, + classify, parse_error_string, parse_ompython_exception, parse_simulation_messages, + summarize_for_llm, warning_gate_complaints, WARNING_GATE_PATTERNS, ) from .session import Backend, OMSession, OpResult -from .loop import AgentLoop, Attempt, LLM, LoopResult, Verifier, extract_code, extract_model_name +from .loop import (AgentLoop, Attempt, LLM, LoopResult, Verifier, WarningGate, + extract_code, extract_model_name) +from .runner import run_ladder, run_comparison from .results import (SimulationResult, all_of, expect_bounds, expect_final, expect_value_at, load_result) -__version__ = "0.2.0" +__version__ = "0.3.0" __all__ = [ "Diagnostic", "Kind", "Severity", "classify", "parse_error_string", "parse_ompython_exception", "parse_simulation_messages", "summarize_for_llm", + "warning_gate_complaints", "WARNING_GATE_PATTERNS", "Backend", "OMSession", "OpResult", - "AgentLoop", "Attempt", "LLM", "LoopResult", "Verifier", - "extract_code", "extract_model_name", + "AgentLoop", "Attempt", "LLM", "LoopResult", "Verifier", "WarningGate", + "extract_code", "extract_model_name", "run_ladder", "run_comparison", "SimulationResult", "load_result", "expect_final", "expect_value_at", "expect_bounds", "all_of", ] diff --git a/omagent/errors.py b/omagent/errors.py index 02913fa..7415203 100644 --- a/omagent/errors.py +++ b/omagent/errors.py @@ -33,8 +33,10 @@ class Kind(str, Enum): # [/path/File.mo:12:3-14:20:writable] Error: message... +# omc also emits single positions without an end range: [...:12:3:writable] _LOCATED = re.compile( - r"^\[(?P[^\]]*?):(?P\d+):(?P\d+)-(?P\d+):(?P\d+):[^\]]*\]\s*" + r"^\[(?P[^\]]*?):(?P\d+):(?P\d+)" + r"(?:-(?P\d+):(?P\d+))?:[^\]]*\]\s*" r"(?PError|Warning|Notification):\s*(?P.*)$", re.DOTALL, ) @@ -153,8 +155,8 @@ def parse_error_string(raw: str) -> list[Diagnostic]: file=m.group("file") or None, line_start=int(m.group("l1")), col_start=int(m.group("c1")), - line_end=int(m.group("l2")), - col_end=int(m.group("c2")), + line_end=int(m.group("l2")) if m.group("l2") else None, + col_end=int(m.group("c2")) if m.group("c2") else None, )) continue m = _BARE.match(rec) @@ -214,15 +216,64 @@ def summarize_for_llm(diags: list[Diagnostic], limit: int = 20) -> str: lines: list[str] = [] for d in errors + warnings: b = d.brief() - if b not in seen: - seen.add(b) - lines.append(b) + if b in seen: + continue + seen.add(b) if len(lines) >= limit: - lines.append(f"... ({len(errors) + len(warnings) - limit} more suppressed)") break + lines.append(b) + suppressed = len(errors) + len(warnings) - len(lines) + if suppressed > 0: + lines.append(f"... ({suppressed} more suppressed)") return "\n".join(lines) if lines else "No errors or warnings." +# --------------------------------------------------------- quality gates -- +# Warning-level quality gates: diagnostics omc emits as mere warnings but +# which indicate real model-quality problems. Matched warnings are surfaced +# as verifier-style complaints and fed back into the fix loop instead of +# being ignored for success. Opt-in at the AgentLoop level so benchmark +# scores stay comparable across versions. + +WARNING_GATE_PATTERNS: list[tuple[str, re.Pattern]] = [ + ("initial conditions not fully specified", re.compile( + r"initial conditions .*not (?:fully )?specified", re.IGNORECASE)), + ("over- or inconsistently specified initial conditions", re.compile( + r"initial conditions .*(?:over-?specified|conflicting|inconsistent)" + r"|conflicting initial conditions", re.IGNORECASE)), + ("over/under-determined system", re.compile( + r"model (?:is )?(?:over-|under-)determined" + r"|(?:over-|under-)determined system", re.IGNORECASE)), + ("inconsistent units", re.compile( + r"units? (?:mismatch|inconsisten|are not equivalent)" + r"|unit (?:mismatch|inconsisten[ct])" + r"|units? .*(?:mismatch|inconsisten|equivalen)", re.IGNORECASE)), +] + +DEFAULT_WARNING_GATE = WARNING_GATE_PATTERNS + + +def warning_gate_complaints( + diags: list[Diagnostic], + rules: list[tuple[str, re.Pattern]] = DEFAULT_WARNING_GATE, +) -> Optional[str]: + """Return a label+message complaint for every gated warning, or None. + + Callables with this shape satisfy the ``omagent.loop.WarningGate`` + protocol: given the diagnostics of an attempt, return None when the + model passes the gate, else a human/LLM-readable complaint. + """ + out: list[str] = [] + for d in diags: + if d.severity != Severity.WARNING: + continue + for label, pat in rules: + if pat.search(d.message): + out.append(f"quality: {label}: {d.message.strip()}") + break + return "; ".join(out) if out else None + + # Newer OMPython raises OMCSessionException whose message embeds omc's log as # "[OMC log for 'sendExpression(...)']: [kind:level:id] message" _OMPY_EXC = re.compile( diff --git a/omagent/loop.py b/omagent/loop.py index d135be1..2ad8e93 100644 --- a/omagent/loop.py +++ b/omagent/loop.py @@ -36,6 +36,10 @@ def propose( # else a human/LLM-readable complaint that is fed into the next fix round. Verifier = Callable[[OpResult], Optional[str]] +# A quality gate inspects a whole attempt's diagnostics (warnings included); +# returns None when the model passes, else a complaint that is fed back. +WarningGate = Callable[[list[Diagnostic]], Optional[str]] + _MODEL_NAME = re.compile(r"^\s*(?:model|block|package)\s+([A-Za-z_][A-Za-z0-9_]*)", re.MULTILINE) @@ -57,7 +61,7 @@ def extract_model_name(code: str) -> Optional[str]: class Attempt: n: int code: str - stage: str # "load" | "check" | "simulate" | "verify" | "ok" + stage: str # "load" | "check" | "simulate" | "quality" | "verify" | "ok" diagnostics: list[Diagnostic] = field(default_factory=list) complaint: Optional[str] = None # verifier feedback, if any @@ -86,6 +90,7 @@ def __init__( max_attempts: int = 4, simulate_options: Optional[dict] = None, verifier: Optional[Verifier] = None, + warning_gate: Optional[WarningGate] = None, ): if max_attempts < 1: raise ValueError("max_attempts must be >= 1") @@ -94,6 +99,7 @@ def __init__( self.max_attempts = max_attempts self.simulate_options = simulate_options or {} self.verifier = verifier + self.warning_gate = warning_gate def run(self, task: str, model_name: Optional[str] = None) -> LoopResult: attempts: list[Attempt] = [] @@ -123,24 +129,36 @@ def run(self, task: str, model_name: Optional[str] = None) -> LoopResult: # -- internals -------------------------------------------------------- def _try_once(self, n: int, code: str, name: str): + all_diags: list[Diagnostic] = [] res = self.session.load_string(code) + all_diags += res.diagnostics if not res.success: - return Attempt(n, code, "load", res.diagnostics), None + return Attempt(n, code, "load", all_diags), None res = self.session.check_model(name) + all_diags += res.diagnostics if not res.success: - return Attempt(n, code, "check", res.diagnostics), None + return Attempt(n, code, "check", all_diags), None sim = self.session.simulate(name, **self.simulate_options) + all_diags += sim.diagnostics if not sim.success: - return Attempt(n, code, "simulate", sim.diagnostics), None + return Attempt(n, code, "simulate", all_diags), None + + # Quality gate first: warnings that indicate sloppy models must be + # fixed even when the sim was fine; complaints feed the fix prompt. + if self.warning_gate is not None: + complaint = self.warning_gate(all_diags) + if complaint: + att = Attempt(n, code, "quality", all_diags, complaint) + return att, sim if self.verifier is not None: complaint = self.verifier(sim) if complaint: - return Attempt(n, code, "verify", sim.diagnostics, complaint), sim + return Attempt(n, code, "verify", all_diags, complaint), sim - return Attempt(n, code, "ok", sim.diagnostics), sim + return Attempt(n, code, "ok", all_diags), sim def _feedback(self, att: Attempt) -> str: parts = [f"Attempt failed at stage '{att.stage}'."] diff --git a/omagent/runner.py b/omagent/runner.py index 980e9a4..e9dc33a 100644 --- a/omagent/runner.py +++ b/omagent/runner.py @@ -8,10 +8,11 @@ import json import pathlib +import statistics import time from typing import Callable, Optional -from .loop import AgentLoop, LLM +from .loop import AgentLoop, LLM, WarningGate from .session import OMSession from .tasks import get_tasks @@ -23,6 +24,7 @@ def run_ladder( max_tier: Optional[int] = None, out_dir: str = "transcripts", max_attempts: int = 4, + warning_gate: Optional[WarningGate] = None, verbose: bool = False, ) -> dict: out = pathlib.Path(out_dir) @@ -50,7 +52,8 @@ def run_ladder( loop = AgentLoop( session, llm, max_attempts=max_attempts, simulate_options=dict(task.simulate_options), - verifier=task.verifier) + verifier=task.verifier, + warning_gate=warning_gate) res = loop.run(task.prompt) elapsed = time.time() - t0 @@ -98,3 +101,91 @@ def run_ladder( "total": len(results)} (out / "summary.json").write_text(json.dumps(report, indent=2)) return report + + +def run_comparison( + session_factory: Callable[[], OMSession], + llm_factories: dict[str, Callable[[], LLM]], + task_ids: Optional[list[str]] = None, + max_tier: Optional[int] = None, + repeats: int = 1, + max_attempts: int = 4, + out_dir: str = "transcripts", + warning_gate: Optional[WarningGate] = None, + verbose: bool = False, +) -> dict: + """Run the ladder several times per LLM and compare models. + + Each (model, task, run) triple gets a fresh session and LLM, so runs are + fully independent; variance across repeats captures LLM / omc + nondeterminism rather than state contamination. + + Transcripts land in ``{out_dir}/{model}/rep{k}/{task}.json``; the + aggregated ``comparison.json`` includes per-model pass rates, per-task + success counts, and mean/variance of attempts and wall time. + """ + if repeats < 1: + raise ValueError("repeats must be >= 1") + if not llm_factories: + raise ValueError("llm_factories must not be empty") + + tasks = get_tasks(task_ids, max_tier) + root = pathlib.Path(out_dir) + models: dict[str, dict] = {} + + for name, llm_factory in llm_factories.items(): + per_task: dict[str, list[dict]] = {t.id: [] for t in tasks} + for k in range(1, repeats + 1): + run_dir = str(root / name / f"rep{k}") + report = run_ladder( + session_factory, llm_factory, + task_ids=[t.id for t in tasks], + out_dir=run_dir, max_attempts=max_attempts, + warning_gate=warning_gate, + verbose=False) + for row in report["results"]: + per_task[row["task"]].append(row) + if verbose: + passed = sum(r["success"] for r in report["results"]) + print(f"[{name} rep {k}/{repeats}] {passed}/" + f"{len(report['results'])} passed", flush=True) + + task_stats = {} + for t in tasks: + runs = per_task[t.id] + attempts = [r["attempts"] for r in runs] + elapsed = [r["elapsed_s"] for r in runs] + n = len(runs) + task_stats[t.id] = { + "tier": t.tier, + "runs": n, + "passed": sum(r["success"] for r in runs), + "pass_rate": sum(r["success"] for r in runs) / n if n else 0.0, + "attempts_mean": statistics.fmean(attempts) if n else 0.0, + "attempts_stdev": statistics.pstdev(attempts) if n > 1 else 0.0, + "elapsed_mean_s": statistics.fmean(elapsed) if n else 0.0, + "final_stages": sorted({r["final_stage"] for r in runs}), + } + + total = sum(len(runs) for runs in per_task.values()) + passed = sum(s["passed"] for s in task_stats.values()) + models[name] = { + "total": total, + "passed": passed, + "pass_rate": passed / total if total else 0.0, + "per_task": task_stats, + } + + comparison = {"repeats": repeats, "models": models} + root.mkdir(parents=True, exist_ok=True) + (root / "comparison.json").write_text(json.dumps(comparison, indent=2)) + + if verbose: + for name, m in models.items(): + print(f"{name}: {m['passed']}/{m['total']} " + f"({m['pass_rate']:.0%})") + for tid, s in m["per_task"].items(): + print(f" {tid:<24} {s['passed']}/{s['runs']}" + f" attempts mean {s['attempts_mean']:.2f}" + f" {s['elapsed_mean_s']:.1f}s") + return comparison diff --git a/omagent/session.py b/omagent/session.py index 0d113c2..56c2475 100644 --- a/omagent/session.py +++ b/omagent/session.py @@ -39,6 +39,15 @@ def _quote(s: str) -> str: return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"' +def _option_value(v: Any) -> str: + """Render a simulate() option as a Modelica literal.""" + if isinstance(v, bool): + return "true" if v else "false" + if isinstance(v, str): + return _quote(v) + return str(v) + + class OMSession: """High-level operations over an omc backend.""" @@ -112,13 +121,13 @@ def check_model(self, model: str) -> OpResult: ok = "completed successfully" in text and not any( d.severity == Severity.ERROR for d in diags) if not ok and text and "completed successfully" not in text: - diags = diags + parse_error_string(text) + extra = parse_error_string(text) + known = {d.message for d in diags} + diags = diags + [d for d in extra if d.message not in known] return OpResult("check_model", ok, text, diags) def simulate(self, model: str, **options: Any) -> OpResult: - opts = "".join( - f", {k}={_quote(v) if isinstance(v, str) else v}" - for k, v in options.items()) + opts = "".join(f", {k}={_option_value(v)}" for k, v in options.items()) val, exc_diags = self._safe_send(f"simulate({model}{opts})") diags = exc_diags + self._drain_diagnostics() messages, result_file = "", "" diff --git a/pyproject.toml b/pyproject.toml index 94ed89a..9600176 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "omagent" -version = "0.2.0" +version = "0.3.0" description = "LLM-assisted OpenModelica modeling: a tested agentic generate-compile-simulate-verify loop with structured omc diagnostics, quantitative trajectory verification, and a benchmark task ladder." readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_quality.py b/tests/test_quality.py new file mode 100644 index 0000000..538c956 --- /dev/null +++ b/tests/test_quality.py @@ -0,0 +1,118 @@ +"""Tests for warning-level quality gates (opt-in, loop + gate function).""" + +import inspect + +from omagent.errors import (Diagnostic, Severity, warning_gate_complaints) +from omagent.session import OMSession +from omagent.loop import AgentLoop +from omagent.runner import run_ladder + +from tests.fakes import FakeOMC +from tests.test_loop import GOOD, CHECK_OK, SIM_OK, ScriptedLLM + +INIT_WARN = ('"[:1:1-1:1:writable] Warning: The initial ' + 'conditions are not fully specified. For model T: the following ' + 'variables have no particular value: v\n"') + + +def _warning_diag(msg: str) -> Diagnostic: + return Diagnostic(Severity.WARNING, msg) + + +def test_gate_matches_initialization_warning(): + complaint = warning_gate_complaints( + [_warning_diag("The initial conditions are not fully specified.")]) + assert complaint is not None + assert "initial conditions not fully specified" in complaint + + +def test_gate_ignores_clean_and_benign_warnings(): + assert warning_gate_complaints([]) is None + assert warning_gate_complaints( + [_warning_diag("some unrelated warning")]) is None + + +def test_gate_ignores_error_severity(): + # Errors are already fatal for ops; the gate is warning-level only. + d = Diagnostic(Severity.ERROR, "The initial conditions are not fully specified.") + assert warning_gate_complaints([d]) is None + + +def test_gate_labels_each_category(): + complaint = warning_gate_complaints([ + _warning_diag("The initial conditions are over-specified. " + "2 conflicting start values."), + _warning_diag("The units of the expressions have to be equivalent."), + ]) or "" + assert "over- or inconsistently specified initial conditions" in complaint + assert "inconsistent units" in complaint + + +def test_loop_without_gate_ignores_warnings(): + fake = FakeOMC({ + "loadString(": [True], + "checkModel(": [CHECK_OK], + "simulate(": [dict(SIM_OK)], + "getErrorString()": [INIT_WARN, '""', '""'], + }) + res = AgentLoop(OMSession(fake), ScriptedLLM([GOOD]), max_attempts=1).run("task") + assert res.success + assert res.attempts[0].stage == "ok" + + +def test_loop_with_gate_sends_warning_back(tmp_path): + csv = tmp_path / "r.csv" + csv.write_text("time,x,v\n0,0,0\n1,0,0\n") + sim = dict(SIM_OK, resultFile=str(csv)) + fake = FakeOMC({ + "loadString(": [True] * 2, + "checkModel(": [CHECK_OK] * 2, + "simulate(": [dict(SIM_OK), dict(sim)], + # drain order per attempt: attempt 1 -> warn, '', '' ; attempt 2 clea n + "getErrorString()": [INIT_WARN, '""', '""', '""', '""', '""', '""'], + }) + loop = AgentLoop( + OMSession(fake), ScriptedLLM([GOOD, GOOD]), max_attempts=2, + warning_gate=warning_gate_complaints) + res = loop.run("task") + + assert len(res.attempts) == 2 + first = res.attempts[0] + assert first.stage == "quality" and first.failed + assert "initial conditions not fully specified" in (first.complaint or "") + + # The gated warning must reach the fix prompt as structured feedback. + calls = loop.llm.calls + assert calls[0][0] == "task" + assert calls[1][1] == GOOD # previous code attached + assert "initial conditions" in calls[1][2] + + assert res.attempts[-1].stage == "ok" + assert res.success + + +def test_ladder_with_gate_reports_quality_final_stage(tmp_path): + fake = FakeOMC({ + "loadString(": [True] * 2, + "checkModel(": [CHECK_OK] * 2, + "simulate(": [dict(SIM_OK)] * 2, + "getErrorString()": [INIT_WARN, '""', '""', '""', '""'] * 2, + }) + report = run_ladder( + session_factory=lambda: OMSession(fake), + llm_factory=lambda: ScriptedLLM([GOOD, GOOD]), + task_ids=["msd_equations"], + out_dir=str(tmp_path / "t"), + max_attempts=2, + warning_gate=warning_gate_complaints, + ) + r = report["results"][0] + assert r["success"] is False + assert r["final_stage"] == "quality" + assert r["attempts"] == 2 + + +def test_ladder_jit_accepts_gate_parameter(): + sig = inspect.signature(run_ladder) + assert "warning_gate" in sig.parameters + assert "verbose" in sig.parameters diff --git a/tests/test_runner_comparison.py b/tests/test_runner_comparison.py new file mode 100644 index 0000000..62851a0 --- /dev/null +++ b/tests/test_runner_comparison.py @@ -0,0 +1,131 @@ +"""Tests for run_comparison: multi-run variance + cross-model aggregation.""" + +import pytest + +from omagent.session import OMSession +from omagent.runner import run_comparison + +from tests.fakes import FakeOMC +from tests.test_loop import GOOD, BAD, CHECK_OK, ScriptedLLM + +SIM_OK = {"resultFile": "RESULT", "messages": "LOG_SUCCESS | info | The simulation finished successfully.\n"} +SIM_FAIL = {"resultFile": "", "messages": "Simulation execution failed for model: T\n"} + + +@pytest.fixture() +def comparison_session(tmp_path): + csv = tmp_path / "res.csv" + csv.write_text( + "time,x,v\n" + "".join( + f"{t/10:.1f},{0.1 if t <= 5 else 0.0:.3f},0\n" + for t in range(0, 101))) + ok = dict(SIM_OK, resultFile=str(csv)) + # One shared fake: strong consumes sim-ok replies, weak consumes sim-fail. + fake = FakeOMC({ + "loadString(": [True] * 20, + "checkModel(": [CHECK_OK] * 20, + "simulate(": [ok] * 2 + [dict(SIM_FAIL)] * 18, + "getErrorString()": ['""'] * 100, + }) + return lambda: OMSession(fake) + + +def test_comparison_arguments_validated(): + with pytest.raises(ValueError): + run_comparison(lambda: None, {}) + with pytest.raises(ValueError): + run_comparison(lambda: None, {"m": lambda: None}, repeats=0) + + +def test_comparison_two_models_two_repeats(comparison_session, tmp_path): + comparison = run_comparison( + session_factory=comparison_session, + llm_factories={ + "strong": lambda: ScriptedLLM([GOOD]), + "weak": lambda: ScriptedLLM([BAD, BAD, BAD]), + }, + task_ids=["msd_equations"], + repeats=2, + out_dir=str(tmp_path / "tr"), + max_attempts=3, + ) + + assert comparison["repeats"] == 2 + strong = comparison["models"]["strong"] + weak = comparison["models"]["weak"] + + assert strong["total"] == weak["total"] == 2 + assert strong["passed"] == 2 and strong["pass_rate"] == 1.0 + assert weak["passed"] == 0 and weak["pass_rate"] == 0.0 + + ts = strong["per_task"]["msd_equations"] + assert ts["runs"] == 2 and ts["passed"] == 2 and ts["pass_rate"] == 1.0 + assert ts["attempts_mean"] == 1.0 and ts["attempts_stdev"] == 0.0 + assert ts["final_stages"] == ["ok"] + assert ts["tier"] == 1 + assert ts["elapsed_mean_s"] >= 0.0 + + ws = weak["per_task"]["msd_equations"] + assert ws["passed"] == 0 and ws["final_stages"] == ["simulate"] + assert ws["attempts_mean"] == 3.0 and ws["attempts_stdev"] == 0.0 + + +def test_comparison_transcripts_per_model_rep(comparison_session, tmp_path): + run_comparison( + session_factory=comparison_session, + llm_factories={"m1": lambda: ScriptedLLM([GOOD])}, + task_ids=["msd_equations"], + repeats=2, + out_dir=str(tmp_path / "tr"), + max_attempts=2, + ) + root = tmp_path / "tr" + assert (root / "comparison.json").exists() + for rep in ("rep1", "rep2"): + assert (root / "m1" / rep / "msd_equations.json").exists() + assert (root / "m1" / rep / "summary.json").exists() + + +def test_comparison_verbose_output(comparison_session, tmp_path, capsys): + run_comparison( + session_factory=comparison_session, + llm_factories={"m1": lambda: ScriptedLLM([GOOD])}, + task_ids=["msd_equations"], + repeats=1, + out_dir=str(tmp_path / "tr"), + max_attempts=2, + verbose=True, + ) + out = capsys.readouterr().out + assert "m1 rep 1/1" in out + assert "1/1 passed" in out + assert "m1: 1/1 (100%)" in out + + +def test_comparison_attempt_variance_tracked(tmp_path): + csv = tmp_path / "res.csv" + csv.write_text( + "time,x,v\n" + "".join( + f"{t/10:.1f},{0.1 if t <= 5 else 0.0:.3f},0\n" + for t in range(0, 101))) + ok = dict(SIM_OK, resultFile=str(csv)) + # BAD attempt 1 fails at load; GOOD attempt 2 passes. The shared fake + # serves both repeats: error / ok-ok-ok / error / ok-ok-ok drains. + fake = FakeOMC({ + "loadString(": [True] * 10, + "checkModel(": [CHECK_OK] * 10, + "simulate(": [ok] * 10, + "getErrorString()": (['"Error: Class FooX not found in scope M.\n"'] + ['""'] * 3) * 6, + }) + comparison = run_comparison( + session_factory=lambda: OMSession(fake), + llm_factories={"var": lambda: ScriptedLLM([BAD, GOOD])}, + task_ids=["msd_equations"], + repeats=2, + out_dir=str(tmp_path / "tr"), + max_attempts=2, + ) + ts = comparison["models"]["var"]["per_task"]["msd_equations"] + assert ts["passed"] == 2 + assert ts["attempts_mean"] == 2.0 + assert ts["attempts_stdev"] == 0.0