From 9ef339f53195fe0dd9dd71effd3ddc9fe824014c Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon Date: Wed, 4 Mar 2026 23:49:23 +0000 Subject: [PATCH] feat(optimize): add trial runner and feedback system Add TrialRunner to bridge TrialConfig to the eval framework (EvalRunner), TraceJudge for LLM-as-judge scoring of agent traces, and FeedbackCollector for aggregating explicit, thumbs, and judge-driven feedback signals. 68 tests pass covering all three modules. Co-Authored-By: Claude Opus 4.6 --- src/openjarvis/optimize/feedback/__init__.py | 6 + src/openjarvis/optimize/feedback/collector.py | 115 ++++++ src/openjarvis/optimize/feedback/judge.py | 118 ++++++ src/openjarvis/optimize/trial_runner.py | 138 +++++++ tests/test_feedback_collector.py | 239 ++++++++++++ tests/test_trace_judge.py | 209 +++++++++++ tests/test_trial_runner.py | 340 ++++++++++++++++++ 7 files changed, 1165 insertions(+) create mode 100644 src/openjarvis/optimize/feedback/__init__.py create mode 100644 src/openjarvis/optimize/feedback/collector.py create mode 100644 src/openjarvis/optimize/feedback/judge.py create mode 100644 src/openjarvis/optimize/trial_runner.py create mode 100644 tests/test_feedback_collector.py create mode 100644 tests/test_trace_judge.py create mode 100644 tests/test_trial_runner.py diff --git a/src/openjarvis/optimize/feedback/__init__.py b/src/openjarvis/optimize/feedback/__init__.py new file mode 100644 index 00000000..e9292954 --- /dev/null +++ b/src/openjarvis/optimize/feedback/__init__.py @@ -0,0 +1,6 @@ +"""Feedback subsystem: LLM-as-judge scoring and signal aggregation.""" + +from openjarvis.optimize.feedback.collector import FeedbackCollector +from openjarvis.optimize.feedback.judge import TraceJudge + +__all__ = ["TraceJudge", "FeedbackCollector"] diff --git a/src/openjarvis/optimize/feedback/collector.py b/src/openjarvis/optimize/feedback/collector.py new file mode 100644 index 00000000..f62301cb --- /dev/null +++ b/src/openjarvis/optimize/feedback/collector.py @@ -0,0 +1,115 @@ +"""FeedbackCollector -- aggregates feedback from multiple sources.""" + +from __future__ import annotations + +import time +from typing import Any, Dict, List, Optional + +from openjarvis.core.types import Trace +from openjarvis.optimize.feedback.judge import TraceJudge + + +class FeedbackCollector: + """Collects feedback signals: explicit user scores + LLM judge evaluations. + + Signals are stored in-memory as dicts with at least ``trace_id``, + ``score``, ``source``, and ``timestamp`` keys. + """ + + def __init__(self) -> None: + self._records: List[Dict[str, Any]] = [] + + # ------------------------------------------------------------------ + # Recording helpers + # ------------------------------------------------------------------ + + def record_explicit( + self, + trace_id: str, + score: float, + source: str = "api", + ) -> None: + """Record an explicit numeric score (0-1) for a trace.""" + self._records.append({ + "trace_id": trace_id, + "score": min(max(score, 0.0), 1.0), + "source": source, + "timestamp": time.time(), + }) + + def record_thumbs(self, trace_id: str, thumbs_up: bool) -> None: + """Record a thumbs-up / thumbs-down signal (converted to 1.0/0.0).""" + self._records.append({ + "trace_id": trace_id, + "score": 1.0 if thumbs_up else 0.0, + "source": "thumbs", + "timestamp": time.time(), + }) + + # ------------------------------------------------------------------ + # Judge-driven evaluation + # ------------------------------------------------------------------ + + def evaluate_traces( + self, + traces: List[Trace], + judge: TraceJudge, + ) -> List[Dict[str, Any]]: + """Score *traces* via the LLM judge and record the results. + + Returns the list of newly created feedback records. + """ + new_records: List[Dict[str, Any]] = [] + for trace in traces: + score, feedback = judge.score_trace(trace) + record: Dict[str, Any] = { + "trace_id": trace.trace_id, + "score": score, + "source": "judge", + "feedback": feedback, + "timestamp": time.time(), + } + self._records.append(record) + new_records.append(record) + return new_records + + # ------------------------------------------------------------------ + # Queries + # ------------------------------------------------------------------ + + def get_records( + self, trace_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Return stored records, optionally filtered by *trace_id*.""" + if trace_id is None: + return list(self._records) + return [r for r in self._records if r["trace_id"] == trace_id] + + def stats(self) -> Dict[str, Any]: + """Return aggregate statistics over all recorded feedback. + + Returns a dict with ``count``, ``mean_score``, and a simple + ``distribution`` bucket (low / medium / high). + """ + count = len(self._records) + if count == 0: + return { + "count": 0, + "mean_score": 0.0, + "distribution": {"low": 0, "medium": 0, "high": 0}, + } + + scores = [r["score"] for r in self._records] + mean_score = sum(scores) / count + low = sum(1 for s in scores if s < 0.4) + medium = sum(1 for s in scores if 0.4 <= s < 0.7) + high = sum(1 for s in scores if s >= 0.7) + + return { + "count": count, + "mean_score": mean_score, + "distribution": {"low": low, "medium": medium, "high": high}, + } + + +__all__ = ["FeedbackCollector"] diff --git a/src/openjarvis/optimize/feedback/judge.py b/src/openjarvis/optimize/feedback/judge.py new file mode 100644 index 00000000..24012766 --- /dev/null +++ b/src/openjarvis/optimize/feedback/judge.py @@ -0,0 +1,118 @@ +"""TraceJudge -- LLM-as-judge scoring for agent traces.""" + +from __future__ import annotations + +import logging +import re +from typing import List, Tuple + +from openjarvis.core.types import Trace +from openjarvis.evals.core.backend import InferenceBackend + +LOGGER = logging.getLogger(__name__) + +_SYSTEM_PROMPT = ( + "You are an expert evaluator of AI assistant traces. " + "You will be shown a user query, the steps the assistant took, " + "and its final result. Rate the overall quality of the response " + "on a scale from 0.0 (completely wrong / unhelpful) to 1.0 " + "(perfect). Provide your score on the first line as a decimal " + "number, then explain your reasoning." +) + +_SCORE_RE = re.compile( + r"(?:Score|Rating|Quality)?\s*[:=]?\s*(\d+(?:\.\d+)?)" + r"(?:\s*/\s*(\d+(?:\.\d+)?))?", + re.IGNORECASE, +) + + +def _format_trace(trace: Trace) -> str: + """Render a Trace into a textual prompt for the judge.""" + lines: List[str] = [] + lines.append(f"## Query\n{trace.query}") + if trace.steps: + lines.append("\n## Steps") + for i, step in enumerate(trace.steps, 1): + step_input = step.input.get("content", str(step.input)) + step_output = step.output.get("content", str(step.output)) + lines.append( + f"{i}. [{step.step_type.value}] " + f"input={step_input!r} output={step_output!r} " + f"({step.duration_seconds:.3f}s)", + ) + lines.append(f"\n## Final Result\n{trace.result}") + return "\n".join(lines) + + +def _parse_score(text: str) -> float: + """Extract a 0-1 score from the judge response. + + Handles formats like ``0.85``, ``Score: 0.85``, ``Rating: 7/10``. + Falls back to 0.5 when parsing fails. + """ + match = _SCORE_RE.search(text) + if match is None: + LOGGER.warning("Could not parse score from judge response; defaulting to 0.5") + return 0.5 + + numerator = float(match.group(1)) + denominator_str = match.group(2) + + if denominator_str is not None: + denominator = float(denominator_str) + if denominator > 0: + return min(max(numerator / denominator, 0.0), 1.0) + return 0.5 + + # If the number is > 1.0 assume it is on a 0-10 scale + if numerator > 1.0: + return min(numerator / 10.0, 1.0) + + return min(max(numerator, 0.0), 1.0) + + +class TraceJudge: + """LLM-as-judge for scoring traces when no ground truth exists. + + Given a :class:`Trace`, the judge constructs a prompt showing the + query, agent steps, and final result, then asks an LLM to rate the + quality on a 0-1 scale. + """ + + def __init__(self, backend: InferenceBackend, model: str) -> None: + self._backend = backend + self._model = model + + def score_trace(self, trace: Trace) -> Tuple[float, str]: + """Score a single trace. + + Returns: + ``(score, feedback)`` where *score* is in [0, 1] and + *feedback* is the judge's textual reasoning. + """ + prompt = _format_trace(trace) + response = self._backend.generate( + prompt, + model=self._model, + system=_SYSTEM_PROMPT, + temperature=0.0, + max_tokens=1024, + ) + score = _parse_score(response) + return score, response + + def batch_evaluate( + self, traces: List[Trace], + ) -> List[Tuple[float, str]]: + """Evaluate multiple traces sequentially. + + Returns a list of ``(score, feedback)`` tuples, one per trace. + """ + results: List[Tuple[float, str]] = [] + for trace in traces: + results.append(self.score_trace(trace)) + return results + + +__all__ = ["TraceJudge"] diff --git a/src/openjarvis/optimize/trial_runner.py b/src/openjarvis/optimize/trial_runner.py new file mode 100644 index 00000000..2698183a --- /dev/null +++ b/src/openjarvis/optimize/trial_runner.py @@ -0,0 +1,138 @@ +"""TrialRunner -- evaluates a proposed config against a benchmark.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any, List + +from openjarvis.evals.core.types import RunConfig, RunSummary +from openjarvis.optimize.types import TrialConfig, TrialResult + +LOGGER = logging.getLogger(__name__) + + +class TrialRunner: + """Evaluates a proposed config against a benchmark. + + Bridges the optimization types (:class:`TrialConfig`) to the eval + framework (:class:`EvalRunner`) so the optimizer can score candidate + configurations end-to-end. + """ + + def __init__( + self, + benchmark: str, + max_samples: int = 50, + judge_model: str = "gpt-5-mini-2025-08-07", + output_dir: str = "results/optimize/", + ) -> None: + self.benchmark = benchmark + self.max_samples = max_samples + self.judge_model = judge_model + self.output_dir = output_dir + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def run_trial(self, trial: TrialConfig) -> TrialResult: + """Run *trial* against the configured benchmark and return a result. + + Steps: + 1. Convert ``trial`` to a :class:`Recipe` and extract params. + 2. Build a :class:`RunConfig` from recipe + benchmark settings. + 3. Lazily import eval-framework registries to resolve the + benchmark -> dataset + scorer, and build the backend. + 4. Execute via ``EvalRunner.run()`` -> :class:`RunSummary`. + 5. Map the summary into a :class:`TrialResult`. + """ + recipe = trial.to_recipe() + run_config = self._build_run_config(trial, recipe) + + # Lazy imports so the optimize package stays lightweight + from openjarvis.evals.cli import ( + _build_backend, + _build_dataset, + _build_judge_backend, + _build_scorer, + ) + from openjarvis.evals.core.runner import EvalRunner + + dataset = _build_dataset(self.benchmark) + backend = _build_backend( + run_config.backend, + run_config.engine_key, + run_config.agent_name or "orchestrator", + run_config.tools, + ) + judge_backend = _build_judge_backend(run_config.judge_model) + scorer = _build_scorer( + self.benchmark, judge_backend, run_config.judge_model, + ) + + try: + summary: RunSummary = EvalRunner( + run_config, dataset, backend, scorer, + ).run() + finally: + backend.close() + judge_backend.close() + + return self._summary_to_result(trial, summary) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_run_config(self, trial: TrialConfig, recipe: Any) -> RunConfig: + """Map recipe fields into a :class:`RunConfig`.""" + model = recipe.model or "default" + backend_name = "jarvis-direct" + if recipe.agent_type is not None: + backend_name = "jarvis-agent" + + model_slug = model.replace("/", "-").replace(":", "-") + output_path = str( + Path(self.output_dir) / f"{trial.trial_id}_{model_slug}.jsonl", + ) + + return RunConfig( + benchmark=self.benchmark, + backend=backend_name, + model=model, + max_samples=self.max_samples, + temperature=recipe.temperature if recipe.temperature is not None else 0.0, + judge_model=self.judge_model, + engine_key=recipe.engine_key, + agent_name=recipe.agent_type, + tools=list(recipe.tools) if recipe.tools else [], + output_path=output_path, + ) + + @staticmethod + def _summary_to_result( + trial: TrialConfig, summary: RunSummary, + ) -> TrialResult: + """Convert a :class:`RunSummary` to a :class:`TrialResult`.""" + total_tokens = summary.total_input_tokens + summary.total_output_tokens + + failure_modes: List[str] = [] + if summary.errors > 0: + failure_modes.append(f"{summary.errors} evaluation errors") + + return TrialResult( + trial_id=trial.trial_id, + config=trial, + accuracy=summary.accuracy, + mean_latency_seconds=summary.mean_latency_seconds, + total_cost_usd=summary.total_cost_usd, + total_energy_joules=summary.total_energy_joules, + total_tokens=total_tokens, + samples_evaluated=summary.total_samples, + failure_modes=failure_modes, + summary=summary, + ) + + +__all__ = ["TrialRunner"] diff --git a/tests/test_feedback_collector.py b/tests/test_feedback_collector.py new file mode 100644 index 00000000..93b2d50b --- /dev/null +++ b/tests/test_feedback_collector.py @@ -0,0 +1,239 @@ +"""Tests for openjarvis.optimize.feedback.collector module.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from openjarvis.core.types import Trace +from openjarvis.optimize.feedback.collector import FeedbackCollector +from openjarvis.optimize.feedback.judge import TraceJudge + + +def _make_trace(trace_id: str = "trace-001") -> Trace: + return Trace(trace_id=trace_id, query="Hello", result="Hi there") + + +# --------------------------------------------------------------------------- +# record_explicit +# --------------------------------------------------------------------------- + + +class TestRecordExplicit: + """FeedbackCollector.record_explicit stores records.""" + + def test_stores_record(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.8) + records = fc.get_records() + assert len(records) == 1 + assert records[0]["trace_id"] == "t1" + assert records[0]["score"] == 0.8 + assert records[0]["source"] == "api" + + def test_custom_source(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5, source="human") + assert fc.get_records()[0]["source"] == "human" + + def test_clamps_score_above_one(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 1.5) + assert fc.get_records()[0]["score"] == 1.0 + + def test_clamps_score_below_zero(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", -0.5) + assert fc.get_records()[0]["score"] == 0.0 + + def test_has_timestamp(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.7) + assert "timestamp" in fc.get_records()[0] + assert fc.get_records()[0]["timestamp"] > 0 + + def test_multiple_records(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + fc.record_explicit("t2", 0.9) + assert len(fc.get_records()) == 2 + + +# --------------------------------------------------------------------------- +# record_thumbs +# --------------------------------------------------------------------------- + + +class TestRecordThumbs: + """FeedbackCollector.record_thumbs converts boolean to score.""" + + def test_thumbs_up_is_one(self) -> None: + fc = FeedbackCollector() + fc.record_thumbs("t1", thumbs_up=True) + assert fc.get_records()[0]["score"] == 1.0 + + def test_thumbs_down_is_zero(self) -> None: + fc = FeedbackCollector() + fc.record_thumbs("t1", thumbs_up=False) + assert fc.get_records()[0]["score"] == 0.0 + + def test_source_is_thumbs(self) -> None: + fc = FeedbackCollector() + fc.record_thumbs("t1", thumbs_up=True) + assert fc.get_records()[0]["source"] == "thumbs" + + +# --------------------------------------------------------------------------- +# evaluate_traces +# --------------------------------------------------------------------------- + + +class TestEvaluateTraces: + """FeedbackCollector.evaluate_traces uses the judge.""" + + def test_returns_new_records(self) -> None: + backend = MagicMock() + backend.generate.return_value = "Score: 0.85\nGood" + judge = TraceJudge(backend=backend, model="m") + + fc = FeedbackCollector() + traces = [_make_trace("t1"), _make_trace("t2")] + new = fc.evaluate_traces(traces, judge) + + assert len(new) == 2 + assert new[0]["trace_id"] == "t1" + assert new[1]["trace_id"] == "t2" + + def test_records_stored_internally(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.7\nOk" + judge = TraceJudge(backend=backend, model="m") + + fc = FeedbackCollector() + fc.evaluate_traces([_make_trace("t1")], judge) + + assert len(fc.get_records()) == 1 + assert fc.get_records()[0]["source"] == "judge" + + def test_record_has_feedback_text(self) -> None: + backend = MagicMock() + backend.generate.return_value = "Score: 0.6\nNeeds improvement" + judge = TraceJudge(backend=backend, model="m") + + fc = FeedbackCollector() + fc.evaluate_traces([_make_trace()], judge) + + record = fc.get_records()[0] + assert "feedback" in record + assert "Needs improvement" in record["feedback"] + + def test_empty_traces_returns_empty(self) -> None: + judge = MagicMock(spec=TraceJudge) + fc = FeedbackCollector() + result = fc.evaluate_traces([], judge) + assert result == [] + + +# --------------------------------------------------------------------------- +# get_records +# --------------------------------------------------------------------------- + + +class TestGetRecords: + """FeedbackCollector.get_records filters by trace_id.""" + + def test_all_records_when_no_filter(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + fc.record_explicit("t2", 0.8) + assert len(fc.get_records()) == 2 + + def test_filter_by_trace_id(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + fc.record_explicit("t2", 0.8) + fc.record_explicit("t1", 0.9) + + records = fc.get_records(trace_id="t1") + assert len(records) == 2 + assert all(r["trace_id"] == "t1" for r in records) + + def test_filter_returns_empty_for_unknown_id(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + assert fc.get_records(trace_id="unknown") == [] + + def test_returns_copies(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + records = fc.get_records() + records.clear() + assert len(fc.get_records()) == 1 + + +# --------------------------------------------------------------------------- +# stats +# --------------------------------------------------------------------------- + + +class TestStats: + """FeedbackCollector.stats returns aggregate statistics.""" + + def test_empty_stats(self) -> None: + fc = FeedbackCollector() + s = fc.stats() + assert s["count"] == 0 + assert s["mean_score"] == 0.0 + assert s["distribution"] == {"low": 0, "medium": 0, "high": 0} + + def test_count(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + fc.record_explicit("t2", 0.9) + assert fc.stats()["count"] == 2 + + def test_mean_score(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.4) + fc.record_explicit("t2", 0.8) + assert abs(fc.stats()["mean_score"] - 0.6) < 1e-6 + + def test_distribution_low(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.1) + fc.record_explicit("t2", 0.2) + s = fc.stats() + assert s["distribution"]["low"] == 2 + assert s["distribution"]["medium"] == 0 + assert s["distribution"]["high"] == 0 + + def test_distribution_medium(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.5) + fc.record_explicit("t2", 0.6) + s = fc.stats() + assert s["distribution"]["medium"] == 2 + + def test_distribution_high(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("t1", 0.8) + fc.record_explicit("t2", 1.0) + s = fc.stats() + assert s["distribution"]["high"] == 2 + + def test_distribution_mixed(self) -> None: + fc = FeedbackCollector() + fc.record_explicit("a", 0.1) # low + fc.record_explicit("b", 0.5) # medium + fc.record_explicit("c", 0.9) # high + s = fc.stats() + assert s["distribution"]["low"] == 1 + assert s["distribution"]["medium"] == 1 + assert s["distribution"]["high"] == 1 + + def test_stats_with_thumbs(self) -> None: + fc = FeedbackCollector() + fc.record_thumbs("t1", thumbs_up=True) + fc.record_thumbs("t2", thumbs_up=False) + s = fc.stats() + assert s["count"] == 2 + assert abs(s["mean_score"] - 0.5) < 1e-6 diff --git a/tests/test_trace_judge.py b/tests/test_trace_judge.py new file mode 100644 index 00000000..d807def4 --- /dev/null +++ b/tests/test_trace_judge.py @@ -0,0 +1,209 @@ +"""Tests for openjarvis.optimize.feedback.judge module.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from openjarvis.core.types import StepType, Trace, TraceStep +from openjarvis.optimize.feedback.judge import TraceJudge, _parse_score + + +# --------------------------------------------------------------------------- +# _parse_score unit tests +# --------------------------------------------------------------------------- + + +class TestParseScore: + """Tests for the internal _parse_score helper.""" + + def test_plain_decimal(self) -> None: + assert _parse_score("0.85") == 0.85 + + def test_score_prefix(self) -> None: + assert _parse_score("Score: 0.90") == 0.90 + + def test_rating_prefix(self) -> None: + assert _parse_score("Rating: 0.75\nGood work.") == 0.75 + + def test_fraction_format(self) -> None: + score = _parse_score("Rating: 7/10\nSolid answer") + assert abs(score - 0.7) < 1e-6 + + def test_fraction_eight_over_ten(self) -> None: + score = _parse_score("Score: 8/10") + assert abs(score - 0.8) < 1e-6 + + def test_whole_number_over_one_treated_as_ten_scale(self) -> None: + score = _parse_score("Quality: 9\nExcellent") + assert abs(score - 0.9) < 1e-6 + + def test_zero_score(self) -> None: + assert _parse_score("Score: 0.0\nTerrible") == 0.0 + + def test_perfect_score(self) -> None: + assert _parse_score("1.0\nPerfect") == 1.0 + + def test_no_score_defaults_to_half(self) -> None: + assert _parse_score("This is just a comment.") == 0.5 + + def test_negative_clamped_to_zero(self) -> None: + # Edge case: if somehow negative appeared + score = _parse_score("Score: 0.0") + assert score == 0.0 + + def test_over_one_on_ten_scale_capped(self) -> None: + score = _parse_score("Score: 11") + assert score <= 1.0 + + def test_quality_prefix(self) -> None: + assert _parse_score("Quality: 0.65") == 0.65 + + +# --------------------------------------------------------------------------- +# TraceJudge +# --------------------------------------------------------------------------- + + +def _make_trace(**overrides) -> Trace: + """Create a minimal Trace for testing.""" + defaults = dict( + trace_id="trace-001", + query="What is 2+2?", + agent="orchestrator", + model="qwen3:8b", + result="The answer is 4.", + ) + defaults.update(overrides) + return Trace(**defaults) + + +class TestTraceJudgeInit: + """TraceJudge constructor stores backend and model.""" + + def test_stores_backend_and_model(self) -> None: + backend = MagicMock() + judge = TraceJudge(backend=backend, model="judge-model") + assert judge._backend is backend + assert judge._model == "judge-model" + + +class TestScoreTrace: + """TraceJudge.score_trace calls backend and parses result.""" + + def test_returns_score_and_feedback(self) -> None: + backend = MagicMock() + backend.generate.return_value = "Score: 0.85\nGood reasoning and correct answer." + judge = TraceJudge(backend=backend, model="judge-model") + + trace = _make_trace() + score, feedback = judge.score_trace(trace) + + assert score == 0.85 + assert "Good reasoning" in feedback + backend.generate.assert_called_once() + + def test_prompt_includes_query(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.9\nGreat" + judge = TraceJudge(backend=backend, model="m") + + trace = _make_trace(query="Explain gravity") + judge.score_trace(trace) + + call_args = backend.generate.call_args + prompt = call_args[0][0] + assert "Explain gravity" in prompt + + def test_prompt_includes_result(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.7\nOk" + judge = TraceJudge(backend=backend, model="m") + + trace = _make_trace(result="Gravity is a fundamental force.") + judge.score_trace(trace) + + call_args = backend.generate.call_args + prompt = call_args[0][0] + assert "Gravity is a fundamental force." in prompt + + def test_prompt_includes_steps(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.8\nGood" + judge = TraceJudge(backend=backend, model="m") + + step = TraceStep( + step_type=StepType.GENERATE, + timestamp=1.0, + duration_seconds=0.5, + input={"content": "some input"}, + output={"content": "some output"}, + ) + trace = _make_trace(steps=[step]) + judge.score_trace(trace) + + call_args = backend.generate.call_args + prompt = call_args[0][0] + assert "generate" in prompt + assert "some input" in prompt + + def test_uses_system_prompt(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.5" + judge = TraceJudge(backend=backend, model="m") + + judge.score_trace(_make_trace()) + + call_args = backend.generate.call_args + assert call_args[1]["system"] != "" + + def test_fraction_score_parsing(self) -> None: + backend = MagicMock() + backend.generate.return_value = "Rating: 7/10\nDecent answer" + judge = TraceJudge(backend=backend, model="m") + + score, _ = judge.score_trace(_make_trace()) + assert abs(score - 0.7) < 1e-6 + + +class TestBatchEvaluate: + """TraceJudge.batch_evaluate processes multiple traces.""" + + def test_returns_one_result_per_trace(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.8\nGood" + judge = TraceJudge(backend=backend, model="m") + + traces = [_make_trace(trace_id=f"t{i}") for i in range(3)] + results = judge.batch_evaluate(traces) + + assert len(results) == 3 + assert all(isinstance(r, tuple) and len(r) == 2 for r in results) + + def test_calls_score_trace_for_each(self) -> None: + backend = MagicMock() + backend.generate.return_value = "0.6\nAverage" + judge = TraceJudge(backend=backend, model="m") + + traces = [_make_trace(trace_id=f"t{i}") for i in range(5)] + judge.batch_evaluate(traces) + + assert backend.generate.call_count == 5 + + def test_empty_list_returns_empty(self) -> None: + backend = MagicMock() + judge = TraceJudge(backend=backend, model="m") + assert judge.batch_evaluate([]) == [] + + def test_different_scores_per_trace(self) -> None: + backend = MagicMock() + backend.generate.side_effect = [ + "Score: 0.9\nExcellent", + "Score: 0.3\nPoor", + ] + judge = TraceJudge(backend=backend, model="m") + + traces = [_make_trace(trace_id="a"), _make_trace(trace_id="b")] + results = judge.batch_evaluate(traces) + + assert results[0][0] == 0.9 + assert results[1][0] == 0.3 diff --git a/tests/test_trial_runner.py b/tests/test_trial_runner.py new file mode 100644 index 00000000..9e36d439 --- /dev/null +++ b/tests/test_trial_runner.py @@ -0,0 +1,340 @@ +"""Tests for openjarvis.optimize.trial_runner module.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from openjarvis.evals.core.types import RunConfig, RunSummary +from openjarvis.optimize.trial_runner import TrialRunner +from openjarvis.optimize.types import TrialConfig, TrialResult + + +class TestTrialRunnerInit: + """TrialRunner.__init__ stores parameters correctly.""" + + def test_default_params(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + assert runner.benchmark == "supergpqa" + assert runner.max_samples == 50 + assert runner.judge_model == "gpt-5-mini-2025-08-07" + assert runner.output_dir == "results/optimize/" + + def test_custom_params(self) -> None: + runner = TrialRunner( + benchmark="gaia", + max_samples=100, + judge_model="custom-judge", + output_dir="/tmp/results/", + ) + assert runner.benchmark == "gaia" + assert runner.max_samples == 100 + assert runner.judge_model == "custom-judge" + assert runner.output_dir == "/tmp/results/" + + +class TestBuildRunConfig: + """TrialRunner._build_run_config maps recipe fields correctly.""" + + def test_model_mapping(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig( + trial_id="t1", + params={"intelligence.model": "qwen3:8b"}, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert isinstance(cfg, RunConfig) + assert cfg.model == "qwen3:8b" + assert cfg.benchmark == "supergpqa" + + def test_agent_maps_to_agent_backend(self) -> None: + runner = TrialRunner(benchmark="gaia") + trial = TrialConfig( + trial_id="t2", + params={ + "intelligence.model": "llama3.1:8b", + "agent.type": "native_react", + }, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.backend == "jarvis-agent" + assert cfg.agent_name == "native_react" + + def test_no_agent_maps_to_direct_backend(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig( + trial_id="t3", + params={"intelligence.model": "qwen3:8b"}, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.backend == "jarvis-direct" + assert cfg.agent_name is None + + def test_tools_mapping(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig( + trial_id="t4", + params={ + "agent.type": "orchestrator", + "tools.tool_set": ["calculator", "think"], + }, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.tools == ["calculator", "think"] + + def test_temperature_mapping(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig( + trial_id="t5", + params={"intelligence.temperature": 0.7}, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.temperature == 0.7 + + def test_engine_key_mapping(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig( + trial_id="t6", + params={"engine.backend": "vllm"}, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.engine_key == "vllm" + + def test_max_samples_from_runner(self) -> None: + runner = TrialRunner(benchmark="supergpqa", max_samples=25) + trial = TrialConfig(trial_id="t7") + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.max_samples == 25 + + def test_judge_model_from_runner(self) -> None: + runner = TrialRunner(benchmark="supergpqa", judge_model="my-judge") + trial = TrialConfig(trial_id="t8") + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.judge_model == "my-judge" + + def test_output_path_contains_trial_id(self) -> None: + runner = TrialRunner(benchmark="supergpqa", output_dir="out/") + trial = TrialConfig( + trial_id="trial-abc", + params={"intelligence.model": "qwen3:8b"}, + ) + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert "trial-abc" in cfg.output_path + assert cfg.output_path.startswith("out/") + + def test_default_model_fallback(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t9") + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.model == "default" + + def test_default_temperature_fallback(self) -> None: + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t10") + recipe = trial.to_recipe() + cfg = runner._build_run_config(trial, recipe) + + assert cfg.temperature == 0.0 + + +class TestRunTrial: + """TrialRunner.run_trial integration (all eval deps mocked).""" + + def _make_summary(self, **overrides) -> RunSummary: + defaults = dict( + benchmark="supergpqa", + category="reasoning", + backend="jarvis-direct", + model="qwen3:8b", + total_samples=50, + scored_samples=48, + correct=40, + accuracy=0.8333, + errors=2, + mean_latency_seconds=1.5, + total_cost_usd=0.10, + total_energy_joules=500.0, + total_input_tokens=10000, + total_output_tokens=5000, + ) + defaults.update(overrides) + return RunSummary(**defaults) + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_returns_trial_result( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary() + mock_runner_instance = MagicMock() + mock_runner_instance.run.return_value = summary + mock_runner_cls.return_value = mock_runner_instance + + mock_backend = MagicMock() + mock_build_backend.return_value = mock_backend + mock_judge = MagicMock() + mock_build_judge.return_value = mock_judge + + runner = TrialRunner(benchmark="supergpqa", max_samples=50) + trial = TrialConfig( + trial_id="t-run", + params={"intelligence.model": "qwen3:8b"}, + ) + + result = runner.run_trial(trial) + + assert isinstance(result, TrialResult) + assert result.trial_id == "t-run" + assert result.config is trial + mock_runner_cls.assert_called_once() + mock_runner_instance.run.assert_called_once() + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_accuracy_from_summary( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary(accuracy=0.92) + mock_runner_cls.return_value.run.return_value = summary + mock_build_backend.return_value = MagicMock() + mock_build_judge.return_value = MagicMock() + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-acc", params={}) + result = runner.run_trial(trial) + + assert result.accuracy == 0.92 + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_tokens_summed( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary( + total_input_tokens=3000, + total_output_tokens=2000, + ) + mock_runner_cls.return_value.run.return_value = summary + mock_build_backend.return_value = MagicMock() + mock_build_judge.return_value = MagicMock() + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-tok", params={}) + result = runner.run_trial(trial) + + assert result.total_tokens == 5000 + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_summary_attached( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary() + mock_runner_cls.return_value.run.return_value = summary + mock_build_backend.return_value = MagicMock() + mock_build_judge.return_value = MagicMock() + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-sum", params={}) + result = runner.run_trial(trial) + + assert result.summary is summary + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_failure_modes_on_errors( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary(errors=5) + mock_runner_cls.return_value.run.return_value = summary + mock_build_backend.return_value = MagicMock() + mock_build_judge.return_value = MagicMock() + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-err", params={}) + result = runner.run_trial(trial) + + assert len(result.failure_modes) == 1 + assert "5" in result.failure_modes[0] + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_no_failure_modes_when_clean( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary(errors=0) + mock_runner_cls.return_value.run.return_value = summary + mock_build_backend.return_value = MagicMock() + mock_build_judge.return_value = MagicMock() + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-ok", params={}) + result = runner.run_trial(trial) + + assert result.failure_modes == [] + + @patch("openjarvis.evals.cli._build_scorer") + @patch("openjarvis.evals.cli._build_judge_backend") + @patch("openjarvis.evals.cli._build_dataset") + @patch("openjarvis.evals.cli._build_backend") + @patch("openjarvis.evals.core.runner.EvalRunner") + def test_run_trial_closes_backends( + self, mock_runner_cls, mock_build_backend, mock_build_dataset, + mock_build_judge, mock_build_scorer, + ) -> None: + summary = self._make_summary() + mock_runner_cls.return_value.run.return_value = summary + mock_backend = MagicMock() + mock_build_backend.return_value = mock_backend + mock_judge = MagicMock() + mock_build_judge.return_value = mock_judge + + runner = TrialRunner(benchmark="supergpqa") + trial = TrialConfig(trial_id="t-close", params={}) + runner.run_trial(trial) + + mock_backend.close.assert_called_once() + mock_judge.close.assert_called_once()