From d234471c03c82826d168b5aefe81da4f8ff8f5ea Mon Sep 17 00:00:00 2001 From: Jon Saad-Falcon Date: Wed, 4 Mar 2026 23:56:30 +0000 Subject: [PATCH] feat(optimize): add optimization engine, SQLite store, and config loader Add the orchestration layer that ties together the LLM optimizer, trial runner, and persistence into a propose-evaluate-analyze loop with early stopping and recipe export. 36 new tests pass. Co-Authored-By: Claude Opus 4.6 --- src/openjarvis/optimize/config.py | 34 ++ src/openjarvis/optimize/optimizer.py | 282 ++++++++++++ src/openjarvis/optimize/store.py | 301 +++++++++++++ tests/test_optimize_store.py | 358 +++++++++++++++ tests/test_optimizer_engine.py | 650 +++++++++++++++++++++++++++ 5 files changed, 1625 insertions(+) create mode 100644 src/openjarvis/optimize/config.py create mode 100644 src/openjarvis/optimize/optimizer.py create mode 100644 src/openjarvis/optimize/store.py create mode 100644 tests/test_optimize_store.py create mode 100644 tests/test_optimizer_engine.py diff --git a/src/openjarvis/optimize/config.py b/src/openjarvis/optimize/config.py new file mode 100644 index 00000000..1b94182e --- /dev/null +++ b/src/openjarvis/optimize/config.py @@ -0,0 +1,34 @@ +"""TOML config loader for optimization runs.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Union + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover + import tomli as tomllib # type: ignore[no-redef] + + +def load_optimize_config(path: Union[str, Path]) -> Dict[str, Any]: + """Load an optimization config TOML file. + + Returns the raw dict with keys such as ``optimize.max_trials``, + ``optimize.benchmark``, ``optimize.search``, ``optimize.fixed``, + ``optimize.constraints``, etc. + + Raises: + FileNotFoundError: If *path* does not exist. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Optimization config not found: {path}") + + with open(path, "rb") as fh: + data: Dict[str, Any] = tomllib.load(fh) + + return data + + +__all__ = ["load_optimize_config"] diff --git a/src/openjarvis/optimize/optimizer.py b/src/openjarvis/optimize/optimizer.py new file mode 100644 index 00000000..aaeff4d7 --- /dev/null +++ b/src/openjarvis/optimize/optimizer.py @@ -0,0 +1,282 @@ +"""OptimizationEngine -- orchestrates the optimize loop. + +Ties together the LLM optimizer, trial runner, and persistence store +into a single propose -> evaluate -> analyze -> repeat loop. +""" + +from __future__ import annotations + +import logging +import uuid +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +try: + import tomli_w +except ModuleNotFoundError: # pragma: no cover + tomli_w = None # type: ignore[assignment] + +from openjarvis.optimize.llm_optimizer import LLMOptimizer +from openjarvis.optimize.store import OptimizationStore +from openjarvis.optimize.trial_runner import TrialRunner +from openjarvis.optimize.types import ( + OptimizationRun, + SearchSpace, + TrialResult, +) + +LOGGER = logging.getLogger(__name__) + + +class OptimizationEngine: + """Orchestrates the optimize loop: propose -> evaluate -> analyze -> repeat.""" + + def __init__( + self, + search_space: SearchSpace, + llm_optimizer: LLMOptimizer, + trial_runner: TrialRunner, + store: Optional[OptimizationStore] = None, + max_trials: int = 20, + early_stop_patience: int = 5, + ) -> None: + self.search_space = search_space + self.llm_optimizer = llm_optimizer + self.trial_runner = trial_runner + self.store = store + self.max_trials = max_trials + self.early_stop_patience = early_stop_patience + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def run( + self, + progress_callback: Optional[Callable[[int, int], None]] = None, + ) -> OptimizationRun: + """Execute the full optimization loop. + + 1. Generate a run_id via uuid. + 2. ``llm_optimizer.propose_initial()`` -> first config. + 3. Loop up to ``max_trials``: + a. ``trial_runner.run_trial(config)`` -> TrialResult + b. ``llm_optimizer.analyze_trial(config, summary, traces)`` + c. Update TrialResult with analysis text + d. Append to history + e. If store, ``store.save_trial(result)`` + f. Update best_trial if accuracy improved + g. Check early stopping (no improvement for *patience* trials) + h. If not stopped, ``llm_optimizer.propose_next(history)`` + 4. Set run status to ``"completed"``. + 5. If store, ``store.save_run(optimization_run)``. + 6. Return the :class:`OptimizationRun`. + + Args: + progress_callback: Optional ``(trial_num, max_trials) -> None`` + called after each trial completes. + """ + run_id = uuid.uuid4().hex[:16] + optimization_run = OptimizationRun( + run_id=run_id, + search_space=self.search_space, + status="running", + optimizer_model=self.llm_optimizer.optimizer_model, + benchmark=getattr(self.trial_runner, "benchmark", ""), + ) + + history: List[TrialResult] = [] + best_accuracy = -1.0 + trials_without_improvement = 0 + + # First config + config = self.llm_optimizer.propose_initial() + + for trial_num in range(1, self.max_trials + 1): + LOGGER.info( + "Trial %d/%d (id=%s)", + trial_num, + self.max_trials, + config.trial_id, + ) + + # Evaluate + result = self.trial_runner.run_trial(config) + + # Analyze + if result.summary is not None: + analysis = self.llm_optimizer.analyze_trial( + config, + result.summary, + ) + else: + analysis = "" + result.analysis = analysis + + # Record + history.append(result) + optimization_run.trials.append(result) + + # Persist trial + if self.store is not None: + self.store.save_trial(run_id, result) + + # Track best + if result.accuracy > best_accuracy: + best_accuracy = result.accuracy + optimization_run.best_trial = result + trials_without_improvement = 0 + else: + trials_without_improvement += 1 + + # Progress callback + if progress_callback is not None: + progress_callback(trial_num, self.max_trials) + + # Early stopping + if trials_without_improvement >= self.early_stop_patience: + LOGGER.info( + "Early stopping after %d trials without improvement.", + self.early_stop_patience, + ) + break + + # Propose next (unless this was the last trial) + if trial_num < self.max_trials: + config = self.llm_optimizer.propose_next(history) + + optimization_run.status = "completed" + + if self.store is not None: + self.store.save_run(optimization_run) + + return optimization_run + + def export_best_recipe( + self, run: OptimizationRun, path: Path + ) -> Path: + """Export the best trial's config as a TOML recipe file. + + Args: + run: A completed :class:`OptimizationRun`. + path: Destination path for the TOML file. + + Returns: + The *path* written to. + + Raises: + ValueError: If there is no best trial in the run. + """ + if run.best_trial is None: + raise ValueError("No best trial to export.") + + recipe_data = self._trial_to_recipe_dict(run.best_trial) + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + if tomli_w is not None: + with open(path, "wb") as fh: + tomli_w.dump(recipe_data, fh) + else: + # Fallback: write TOML manually + self._write_toml_fallback(recipe_data, path) + + run.best_recipe_path = str(path) + return path + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _trial_to_recipe_dict(trial: TrialResult) -> Dict[str, Any]: + """Convert a TrialResult into a Recipe-style TOML dict.""" + params = trial.config.params + recipe: Dict[str, Any] = { + "recipe": { + "name": f"optimized-{trial.trial_id}", + "description": ( + f"Auto-optimized config (accuracy={trial.accuracy:.4f})" + ), + "version": "1.0.0", + }, + } + + # Intelligence section + intel: Dict[str, Any] = {} + if "intelligence.model" in params: + intel["model"] = params["intelligence.model"] + if "intelligence.temperature" in params: + intel["temperature"] = params["intelligence.temperature"] + if "intelligence.quantization" in params: + intel["quantization"] = params["intelligence.quantization"] + if "intelligence.system_prompt" in params: + intel["system_prompt"] = params["intelligence.system_prompt"] + if "intelligence.max_tokens" in params: + intel["max_tokens"] = params["intelligence.max_tokens"] + if "intelligence.top_p" in params: + intel["top_p"] = params["intelligence.top_p"] + if intel: + recipe["intelligence"] = intel + + # Engine section + engine: Dict[str, Any] = {} + if "engine.backend" in params: + engine["key"] = params["engine.backend"] + if engine: + recipe["engine"] = engine + + # Agent section + agent: Dict[str, Any] = {} + if "agent.type" in params: + agent["type"] = params["agent.type"] + if "agent.max_turns" in params: + agent["max_turns"] = params["agent.max_turns"] + if "agent.system_prompt" in params: + agent["system_prompt"] = params["agent.system_prompt"] + if "tools.tool_set" in params: + agent["tools"] = params["tools.tool_set"] + if agent: + recipe["agent"] = agent + + # Learning section + learning: Dict[str, Any] = {} + if "learning.routing_policy" in params: + learning["routing"] = params["learning.routing_policy"] + if "learning.agent_policy" in params: + learning["agent"] = params["learning.agent_policy"] + if learning: + recipe["learning"] = learning + + return recipe + + @staticmethod + def _write_toml_fallback( + data: Dict[str, Any], path: Path + ) -> None: + """Write a simple nested dict as TOML without tomli_w.""" + lines: List[str] = [] + for section, values in data.items(): + if not isinstance(values, dict): + continue + lines.append(f"[{section}]") + for key, val in values.items(): + if isinstance(val, str): + lines.append(f'{key} = "{val}"') + elif isinstance(val, bool): + lines.append(f"{key} = {'true' if val else 'false'}") + elif isinstance(val, (int, float)): + lines.append(f"{key} = {val}") + elif isinstance(val, list): + items = ", ".join( + f'"{v}"' if isinstance(v, str) else str(v) + for v in val + ) + lines.append(f"{key} = [{items}]") + else: + lines.append(f'{key} = "{val}"') + lines.append("") + path.write_text("\n".join(lines), encoding="utf-8") + + +__all__ = ["OptimizationEngine"] diff --git a/src/openjarvis/optimize/store.py b/src/openjarvis/optimize/store.py new file mode 100644 index 00000000..d27a8da2 --- /dev/null +++ b/src/openjarvis/optimize/store.py @@ -0,0 +1,301 @@ +"""SQLite-backed storage for optimization runs and trials.""" + +from __future__ import annotations + +import json +import sqlite3 +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from openjarvis.optimize.types import ( + OptimizationRun, + SearchSpace, + TrialConfig, + TrialResult, +) + +_CREATE_RUNS = """\ +CREATE TABLE IF NOT EXISTS optimization_runs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + run_id TEXT NOT NULL UNIQUE, + search_space TEXT NOT NULL DEFAULT '{}', + status TEXT NOT NULL DEFAULT 'running', + optimizer_model TEXT NOT NULL DEFAULT '', + benchmark TEXT NOT NULL DEFAULT '', + best_trial_id TEXT, + best_recipe_path TEXT, + created_at REAL NOT NULL DEFAULT 0.0, + updated_at REAL NOT NULL DEFAULT 0.0 +); +""" + +_CREATE_TRIALS = """\ +CREATE TABLE IF NOT EXISTS trial_results ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + trial_id TEXT NOT NULL, + run_id TEXT NOT NULL, + config TEXT NOT NULL DEFAULT '{}', + reasoning TEXT NOT NULL DEFAULT '', + accuracy REAL NOT NULL DEFAULT 0.0, + mean_latency_seconds REAL NOT NULL DEFAULT 0.0, + total_cost_usd REAL NOT NULL DEFAULT 0.0, + total_energy_joules REAL NOT NULL DEFAULT 0.0, + total_tokens INTEGER NOT NULL DEFAULT 0, + samples_evaluated INTEGER NOT NULL DEFAULT 0, + analysis TEXT NOT NULL DEFAULT '', + failure_modes TEXT NOT NULL DEFAULT '[]', + created_at REAL NOT NULL DEFAULT 0.0, + FOREIGN KEY (run_id) REFERENCES optimization_runs(run_id) +); +""" + +_INSERT_RUN = """\ +INSERT OR REPLACE INTO optimization_runs ( + run_id, search_space, status, optimizer_model, benchmark, + best_trial_id, best_recipe_path, created_at, updated_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) +""" + +_INSERT_TRIAL = """\ +INSERT OR REPLACE INTO trial_results ( + trial_id, run_id, config, reasoning, accuracy, + mean_latency_seconds, total_cost_usd, total_energy_joules, + total_tokens, samples_evaluated, analysis, failure_modes, + created_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +""" + + +class OptimizationStore: + """SQLite-backed storage for optimization runs and trials.""" + + def __init__(self, db_path: Union[str, Path]) -> None: + self._db_path = str(db_path) + self._conn = sqlite3.connect(self._db_path) + self._conn.execute("PRAGMA journal_mode=WAL") + self._conn.execute(_CREATE_RUNS) + self._conn.execute(_CREATE_TRIALS) + self._conn.commit() + + # ------------------------------------------------------------------ + # Runs + # ------------------------------------------------------------------ + + def save_run(self, run: OptimizationRun) -> None: + """Persist an optimization run (insert or update).""" + now = time.time() + search_space_json = self._search_space_to_json(run.search_space) + best_trial_id = run.best_trial.trial_id if run.best_trial else None + self._conn.execute( + _INSERT_RUN, + ( + run.run_id, + search_space_json, + run.status, + run.optimizer_model, + run.benchmark, + best_trial_id, + run.best_recipe_path, + now, + now, + ), + ) + self._conn.commit() + + def get_run(self, run_id: str) -> Optional[OptimizationRun]: + """Retrieve an optimization run by id, or ``None``.""" + row = self._conn.execute( + "SELECT * FROM optimization_runs WHERE run_id = ?", + (run_id,), + ).fetchone() + if row is None: + return None + return self._row_to_run(row) + + def list_runs(self, limit: int = 50) -> List[Dict[str, Any]]: + """Return summary dicts of recent optimization runs.""" + rows = self._conn.execute( + "SELECT * FROM optimization_runs ORDER BY created_at DESC LIMIT ?", + (limit,), + ).fetchall() + result: List[Dict[str, Any]] = [] + for row in rows: + result.append( + { + "run_id": row[1], + "status": row[3], + "optimizer_model": row[4], + "benchmark": row[5], + "best_trial_id": row[6], + "best_recipe_path": row[7], + "created_at": row[8], + "updated_at": row[9], + } + ) + return result + + # ------------------------------------------------------------------ + # Trials + # ------------------------------------------------------------------ + + def save_trial(self, run_id: str, trial: TrialResult) -> None: + """Persist a single trial result.""" + now = time.time() + self._conn.execute( + _INSERT_TRIAL, + ( + trial.trial_id, + run_id, + json.dumps(trial.config.params), + trial.config.reasoning, + trial.accuracy, + trial.mean_latency_seconds, + trial.total_cost_usd, + trial.total_energy_joules, + trial.total_tokens, + trial.samples_evaluated, + trial.analysis, + json.dumps(trial.failure_modes), + now, + ), + ) + self._conn.commit() + + def get_trials(self, run_id: str) -> List[TrialResult]: + """Retrieve all trial results for a given run.""" + rows = self._conn.execute( + "SELECT * FROM trial_results WHERE run_id = ? ORDER BY id", + (run_id,), + ).fetchall() + return [self._row_to_trial(r) for r in rows] + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def close(self) -> None: + """Close the underlying SQLite connection.""" + self._conn.close() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _search_space_to_json(space: SearchSpace) -> str: + """Serialize a SearchSpace to JSON.""" + dims = [] + for d in space.dimensions: + dims.append( + { + "name": d.name, + "dim_type": d.dim_type, + "values": d.values, + "low": d.low, + "high": d.high, + "description": d.description, + "pillar": d.pillar, + } + ) + return json.dumps( + { + "dimensions": dims, + "fixed": space.fixed, + "constraints": space.constraints, + } + ) + + @staticmethod + def _json_to_search_space(raw: str) -> SearchSpace: + """Deserialize a SearchSpace from JSON.""" + from openjarvis.optimize.types import SearchDimension + + data = json.loads(raw) + dims = [] + for d in data.get("dimensions", []): + dims.append( + SearchDimension( + name=d.get("name", ""), + dim_type=d.get("dim_type", "categorical"), + values=d.get("values", []), + low=d.get("low"), + high=d.get("high"), + description=d.get("description", ""), + pillar=d.get("pillar", ""), + ) + ) + return SearchSpace( + dimensions=dims, + fixed=data.get("fixed", {}), + constraints=data.get("constraints", []), + ) + + def _row_to_run(self, row: tuple) -> OptimizationRun: + """Convert a database row to an OptimizationRun.""" + run_id = row[1] + search_space = self._json_to_search_space(row[2]) + status = row[3] + optimizer_model = row[4] + benchmark = row[5] + best_trial_id = row[6] + best_recipe_path = row[7] + + # Load trials for this run + trials = self.get_trials(run_id) + + # Find the best trial + best_trial: Optional[TrialResult] = None + if best_trial_id: + for t in trials: + if t.trial_id == best_trial_id: + best_trial = t + break + + return OptimizationRun( + run_id=run_id, + search_space=search_space, + trials=trials, + best_trial=best_trial, + best_recipe_path=best_recipe_path, + status=status, + optimizer_model=optimizer_model, + benchmark=benchmark, + ) + + @staticmethod + def _row_to_trial(row: tuple) -> TrialResult: + """Convert a database row to a TrialResult.""" + trial_id = row[1] + # row[2] = run_id (not stored on TrialResult) + params = json.loads(row[3]) + reasoning = row[4] + accuracy = row[5] + mean_latency = row[6] + cost = row[7] + energy = row[8] + tokens = row[9] + samples = row[10] + analysis = row[11] + failure_modes = json.loads(row[12]) + + config = TrialConfig( + trial_id=trial_id, + params=params, + reasoning=reasoning, + ) + return TrialResult( + trial_id=trial_id, + config=config, + accuracy=accuracy, + mean_latency_seconds=mean_latency, + total_cost_usd=cost, + total_energy_joules=energy, + total_tokens=tokens, + samples_evaluated=samples, + analysis=analysis, + failure_modes=failure_modes, + ) + + +__all__ = ["OptimizationStore"] diff --git a/tests/test_optimize_store.py b/tests/test_optimize_store.py new file mode 100644 index 00000000..9a308a6c --- /dev/null +++ b/tests/test_optimize_store.py @@ -0,0 +1,358 @@ +"""Tests for openjarvis.optimize.store module.""" + +from __future__ import annotations + +import json + +from openjarvis.optimize.store import OptimizationStore +from openjarvis.optimize.types import ( + OptimizationRun, + SearchDimension, + SearchSpace, + TrialConfig, + TrialResult, +) + + +def _sample_search_space() -> SearchSpace: + return SearchSpace( + dimensions=[ + SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator"], + pillar="agent", + ), + SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + pillar="intelligence", + ), + ], + fixed={"engine": "ollama"}, + constraints=["max_turns >= 1"], + ) + + +def _sample_trial( + trial_id: str = "t1", + accuracy: float = 0.8, + params: dict | None = None, +) -> TrialResult: + if params is None: + params = {"agent.type": "orchestrator", "intelligence.temperature": 0.5} + config = TrialConfig( + trial_id=trial_id, + params=params, + reasoning="testing", + ) + return TrialResult( + trial_id=trial_id, + config=config, + accuracy=accuracy, + mean_latency_seconds=1.5, + total_cost_usd=0.02, + total_energy_joules=100.0, + total_tokens=3000, + samples_evaluated=50, + analysis="Solid performance", + failure_modes=["timeout on long inputs"], + ) + + +# --------------------------------------------------------------------------- +# OptimizationStore.__init__ +# --------------------------------------------------------------------------- + + +class TestOptimizationStoreInit: + """Tests for OptimizationStore initialization.""" + + def test_creates_tables(self, tmp_path) -> None: + db = tmp_path / "opt.db" + store = OptimizationStore(db) + # Verify tables exist by querying them + runs = store.list_runs() + trials = store.get_trials("nonexistent") + assert runs == [] + assert trials == [] + store.close() + + def test_creates_tables_string_path(self, tmp_path) -> None: + db = str(tmp_path / "opt.db") + store = OptimizationStore(db) + runs = store.list_runs() + assert runs == [] + store.close() + + def test_wal_mode(self, tmp_path) -> None: + db = tmp_path / "opt.db" + store = OptimizationStore(db) + row = store._conn.execute("PRAGMA journal_mode").fetchone() + assert row[0] == "wal" + store.close() + + +# --------------------------------------------------------------------------- +# Trial persistence +# --------------------------------------------------------------------------- + + +class TestTrialPersistence: + """Tests for save_trial + get_trials roundtrip.""" + + def test_save_and_get_trials(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + run_id = "run-001" + # Need to save a run first for foreign key + space = _sample_search_space() + run = OptimizationRun(run_id=run_id, search_space=space) + store.save_run(run) + + trial = _sample_trial("t1", accuracy=0.8) + store.save_trial(run_id, trial) + + trials = store.get_trials(run_id) + assert len(trials) == 1 + t = trials[0] + assert t.trial_id == "t1" + assert t.accuracy == 0.8 + assert t.mean_latency_seconds == 1.5 + assert t.total_cost_usd == 0.02 + assert t.total_energy_joules == 100.0 + assert t.total_tokens == 3000 + assert t.samples_evaluated == 50 + assert t.analysis == "Solid performance" + assert t.failure_modes == ["timeout on long inputs"] + assert t.config.params == { + "agent.type": "orchestrator", + "intelligence.temperature": 0.5, + } + assert t.config.reasoning == "testing" + store.close() + + def test_multiple_trials(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + run_id = "run-002" + space = _sample_search_space() + store.save_run(OptimizationRun(run_id=run_id, search_space=space)) + + store.save_trial(run_id, _sample_trial("t1", accuracy=0.7)) + store.save_trial(run_id, _sample_trial("t2", accuracy=0.85)) + store.save_trial(run_id, _sample_trial("t3", accuracy=0.9)) + + trials = store.get_trials(run_id) + assert len(trials) == 3 + assert [t.trial_id for t in trials] == ["t1", "t2", "t3"] + assert [t.accuracy for t in trials] == [0.7, 0.85, 0.9] + store.close() + + def test_get_trials_empty_run(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + trials = store.get_trials("nonexistent-run") + assert trials == [] + store.close() + + def test_trial_with_empty_failure_modes(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + run_id = "run-003" + space = _sample_search_space() + store.save_run(OptimizationRun(run_id=run_id, search_space=space)) + + config = TrialConfig(trial_id="t1", params={"x": 1}) + trial = TrialResult( + trial_id="t1", + config=config, + accuracy=0.5, + failure_modes=[], + ) + store.save_trial(run_id, trial) + + loaded = store.get_trials(run_id) + assert loaded[0].failure_modes == [] + store.close() + + +# --------------------------------------------------------------------------- +# Run persistence +# --------------------------------------------------------------------------- + + +class TestRunPersistence: + """Tests for save_run + get_run roundtrip.""" + + def test_save_and_get_run(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + best = _sample_trial("best", accuracy=0.95) + run = OptimizationRun( + run_id="run-abc", + search_space=space, + best_trial=best, + best_recipe_path="/tmp/best.toml", + status="completed", + optimizer_model="claude-sonnet-4-6", + benchmark="supergpqa", + ) + # Save the best trial to the DB so get_run can reconstruct it + store.save_run(run) + store.save_trial("run-abc", best) + + loaded = store.get_run("run-abc") + assert loaded is not None + assert loaded.run_id == "run-abc" + assert loaded.status == "completed" + assert loaded.optimizer_model == "claude-sonnet-4-6" + assert loaded.benchmark == "supergpqa" + assert loaded.best_recipe_path == "/tmp/best.toml" + assert loaded.best_trial is not None + assert loaded.best_trial.trial_id == "best" + assert loaded.best_trial.accuracy == 0.95 + store.close() + + def test_get_run_not_found(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + result = store.get_run("nonexistent") + assert result is None + store.close() + + def test_save_run_without_best_trial(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + run = OptimizationRun( + run_id="run-no-best", + search_space=space, + status="running", + ) + store.save_run(run) + + loaded = store.get_run("run-no-best") + assert loaded is not None + assert loaded.best_trial is None + assert loaded.status == "running" + store.close() + + def test_search_space_roundtrip(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + run = OptimizationRun(run_id="run-space", search_space=space) + store.save_run(run) + + loaded = store.get_run("run-space") + assert loaded is not None + assert len(loaded.search_space.dimensions) == 2 + assert loaded.search_space.dimensions[0].name == "agent.type" + assert loaded.search_space.dimensions[0].dim_type == "categorical" + assert loaded.search_space.dimensions[0].values == [ + "simple", + "orchestrator", + ] + assert loaded.search_space.dimensions[1].name == "intelligence.temperature" + assert loaded.search_space.dimensions[1].low == 0.0 + assert loaded.search_space.dimensions[1].high == 1.0 + assert loaded.search_space.fixed == {"engine": "ollama"} + assert loaded.search_space.constraints == ["max_turns >= 1"] + store.close() + + def test_run_with_trials_loaded(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + run = OptimizationRun(run_id="run-trials", search_space=space) + store.save_run(run) + + store.save_trial("run-trials", _sample_trial("t1", accuracy=0.7)) + store.save_trial("run-trials", _sample_trial("t2", accuracy=0.85)) + + loaded = store.get_run("run-trials") + assert loaded is not None + assert len(loaded.trials) == 2 + assert loaded.trials[0].trial_id == "t1" + assert loaded.trials[1].trial_id == "t2" + store.close() + + +# --------------------------------------------------------------------------- +# list_runs +# --------------------------------------------------------------------------- + + +class TestListRuns: + """Tests for list_runs.""" + + def test_list_runs_empty(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + runs = store.list_runs() + assert runs == [] + store.close() + + def test_list_runs_returns_summaries(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + + for i in range(3): + run = OptimizationRun( + run_id=f"run-{i}", + search_space=space, + status="completed", + optimizer_model="test-model", + benchmark="test-bench", + ) + store.save_run(run) + + runs = store.list_runs() + assert len(runs) == 3 + # Should have summary keys + for r in runs: + assert "run_id" in r + assert "status" in r + assert "optimizer_model" in r + assert "benchmark" in r + assert "created_at" in r + store.close() + + def test_list_runs_limit(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + space = _sample_search_space() + + for i in range(10): + run = OptimizationRun( + run_id=f"run-{i}", + search_space=space, + ) + store.save_run(run) + + runs = store.list_runs(limit=5) + assert len(runs) == 5 + store.close() + + +# --------------------------------------------------------------------------- +# close +# --------------------------------------------------------------------------- + + +class TestClose: + """Tests for close.""" + + def test_close(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + store.close() + # After close, operations should raise + try: + store.list_runs() + assert False, "Expected error after close" + except Exception: + pass + + def test_double_close(self, tmp_path) -> None: + """Double close should not raise.""" + store = OptimizationStore(tmp_path / "opt.db") + store.close() + # Second close may or may not raise depending on sqlite3 + # Just verify it doesn't crash hard + try: + store.close() + except Exception: + pass diff --git a/tests/test_optimizer_engine.py b/tests/test_optimizer_engine.py new file mode 100644 index 00000000..f18608ac --- /dev/null +++ b/tests/test_optimizer_engine.py @@ -0,0 +1,650 @@ +"""Tests for openjarvis.optimize.optimizer module.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +try: + import tomllib +except ModuleNotFoundError: + import tomli as tomllib # type: ignore[no-redef] + +from openjarvis.evals.core.types import RunSummary +from openjarvis.optimize.optimizer import OptimizationEngine +from openjarvis.optimize.store import OptimizationStore +from openjarvis.optimize.types import ( + OptimizationRun, + SearchDimension, + SearchSpace, + TrialConfig, + TrialResult, +) + + +def _sample_search_space() -> SearchSpace: + return SearchSpace( + dimensions=[ + SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator"], + pillar="agent", + ), + SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + pillar="intelligence", + ), + ], + fixed={"engine": "ollama"}, + ) + + +def _sample_summary(accuracy: float = 0.8) -> RunSummary: + return RunSummary( + benchmark="test", + category="reasoning", + backend="jarvis-direct", + model="test-model", + total_samples=50, + scored_samples=50, + correct=int(accuracy * 50), + accuracy=accuracy, + errors=0, + mean_latency_seconds=1.0, + total_cost_usd=0.01, + ) + + +def _sample_trial_result( + trial_id: str = "t1", + accuracy: float = 0.8, + summary: RunSummary | None = None, +) -> TrialResult: + config = TrialConfig( + trial_id=trial_id, + params={"agent.type": "orchestrator"}, + reasoning="test reasoning", + ) + if summary is None: + summary = _sample_summary(accuracy) + return TrialResult( + trial_id=trial_id, + config=config, + accuracy=accuracy, + mean_latency_seconds=1.0, + total_cost_usd=0.01, + samples_evaluated=50, + summary=summary, + ) + + +# --------------------------------------------------------------------------- +# __init__ +# --------------------------------------------------------------------------- + + +class TestOptimizationEngineInit: + """Tests for OptimizationEngine.__init__.""" + + def test_stores_all_params(self) -> None: + space = _sample_search_space() + optimizer = MagicMock() + runner = MagicMock() + store = MagicMock() + + engine = OptimizationEngine( + search_space=space, + llm_optimizer=optimizer, + trial_runner=runner, + store=store, + max_trials=10, + early_stop_patience=3, + ) + + assert engine.search_space is space + assert engine.llm_optimizer is optimizer + assert engine.trial_runner is runner + assert engine.store is store + assert engine.max_trials == 10 + assert engine.early_stop_patience == 3 + + def test_default_params(self) -> None: + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=MagicMock(), + trial_runner=MagicMock(), + ) + assert engine.store is None + assert engine.max_trials == 20 + assert engine.early_stop_patience == 5 + + +# --------------------------------------------------------------------------- +# run() with mocked dependencies +# --------------------------------------------------------------------------- + + +class TestOptimizationEngineRun: + """Tests for OptimizationEngine.run().""" + + def test_basic_run(self) -> None: + space = _sample_search_space() + optimizer = MagicMock() + runner = MagicMock() + + initial_config = TrialConfig( + trial_id="init", + params={"agent.type": "orchestrator"}, + ) + second_config = TrialConfig( + trial_id="next", + params={"agent.type": "simple"}, + ) + optimizer.propose_initial.return_value = initial_config + optimizer.propose_next.return_value = second_config + optimizer.analyze_trial.return_value = "analysis text" + optimizer.optimizer_model = "test-model" + + runner.run_trial.return_value = _sample_trial_result( + "init", accuracy=0.8, + ) + runner.benchmark = "supergpqa" + + engine = OptimizationEngine( + search_space=space, + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=2, + ) + + # Make second trial return different accuracy + runner.run_trial.side_effect = [ + _sample_trial_result("init", accuracy=0.8), + _sample_trial_result("next", accuracy=0.85), + ] + + result = engine.run() + + assert isinstance(result, OptimizationRun) + assert result.status == "completed" + assert len(result.trials) == 2 + assert result.best_trial is not None + assert result.best_trial.accuracy == 0.85 + assert result.optimizer_model == "test-model" + assert result.benchmark == "supergpqa" + + optimizer.propose_initial.assert_called_once() + assert runner.run_trial.call_count == 2 + assert optimizer.analyze_trial.call_count == 2 + + def test_single_trial(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + config = TrialConfig(trial_id="only", params={}) + optimizer.propose_initial.return_value = config + optimizer.analyze_trial.return_value = "good" + optimizer.optimizer_model = "m" + runner.run_trial.return_value = _sample_trial_result("only", 0.9) + runner.benchmark = "test" + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=1, + ) + result = engine.run() + + assert len(result.trials) == 1 + assert result.best_trial.accuracy == 0.9 + # propose_next should NOT be called when max_trials=1 + optimizer.propose_next.assert_not_called() + + def test_analysis_text_set_on_result(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t1", params={}, + ) + optimizer.analyze_trial.return_value = "detailed analysis" + optimizer.optimizer_model = "m" + runner.run_trial.return_value = _sample_trial_result("t1", 0.8) + runner.benchmark = "b" + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=1, + ) + result = engine.run() + + assert result.trials[0].analysis == "detailed analysis" + + def test_run_without_summary_skips_analysis(self) -> None: + """If trial result has no summary, analysis should be empty.""" + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t1", params={}, + ) + optimizer.optimizer_model = "m" + + # Result with no summary + result_no_summary = TrialResult( + trial_id="t1", + config=TrialConfig(trial_id="t1", params={}), + accuracy=0.5, + summary=None, + ) + runner.run_trial.return_value = result_no_summary + runner.benchmark = "b" + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=1, + ) + run = engine.run() + + assert run.trials[0].analysis == "" + optimizer.analyze_trial.assert_not_called() + + +# --------------------------------------------------------------------------- +# Early stopping +# --------------------------------------------------------------------------- + + +class TestEarlyStopping: + """Tests for early stopping behavior.""" + + def test_early_stop_after_patience(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.propose_next.side_effect = [ + TrialConfig(trial_id=f"t{i}", params={}) + for i in range(1, 20) + ] + optimizer.analyze_trial.return_value = "ok" + optimizer.optimizer_model = "m" + runner.benchmark = "b" + + # First trial is the best; all subsequent are worse + results = [_sample_trial_result("t0", accuracy=0.9)] + for i in range(1, 20): + results.append(_sample_trial_result(f"t{i}", accuracy=0.5)) + runner.run_trial.side_effect = results + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=20, + early_stop_patience=3, + ) + run = engine.run() + + # Should stop after 1 (best) + 3 (patience) = 4 trials + assert len(run.trials) == 4 + assert run.best_trial.trial_id == "t0" + assert run.status == "completed" + + def test_no_early_stop_when_improving(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.propose_next.side_effect = [ + TrialConfig(trial_id=f"t{i}", params={}) + for i in range(1, 5) + ] + optimizer.analyze_trial.return_value = "ok" + optimizer.optimizer_model = "m" + runner.benchmark = "b" + + # Accuracy keeps improving + results = [ + _sample_trial_result("t0", accuracy=0.5), + _sample_trial_result("t1", accuracy=0.6), + _sample_trial_result("t2", accuracy=0.7), + _sample_trial_result("t3", accuracy=0.8), + _sample_trial_result("t4", accuracy=0.9), + ] + runner.run_trial.side_effect = results + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=5, + early_stop_patience=3, + ) + run = engine.run() + + assert len(run.trials) == 5 + assert run.best_trial.accuracy == 0.9 + + +# --------------------------------------------------------------------------- +# Progress callback +# --------------------------------------------------------------------------- + + +class TestProgressCallback: + """Tests for progress_callback.""" + + def test_callback_called(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.propose_next.side_effect = [ + TrialConfig(trial_id=f"t{i}", params={}) + for i in range(1, 3) + ] + optimizer.analyze_trial.return_value = "ok" + optimizer.optimizer_model = "m" + runner.benchmark = "b" + + results = [ + _sample_trial_result(f"t{i}", accuracy=0.5 + i * 0.1) + for i in range(3) + ] + runner.run_trial.side_effect = results + + callback = MagicMock() + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=3, + ) + engine.run(progress_callback=callback) + + assert callback.call_count == 3 + callback.assert_any_call(1, 3) + callback.assert_any_call(2, 3) + callback.assert_any_call(3, 3) + + def test_no_callback(self) -> None: + """run() should work fine without a callback.""" + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.analyze_trial.return_value = "ok" + optimizer.optimizer_model = "m" + runner.run_trial.return_value = _sample_trial_result("t0", 0.8) + runner.benchmark = "b" + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + max_trials=1, + ) + run = engine.run() + assert run.status == "completed" + + +# --------------------------------------------------------------------------- +# export_best_recipe +# --------------------------------------------------------------------------- + + +class TestExportBestRecipe: + """Tests for export_best_recipe.""" + + def test_exports_valid_toml(self, tmp_path) -> None: + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=MagicMock(), + trial_runner=MagicMock(), + ) + + best = _sample_trial_result("best", accuracy=0.95) + best.config.params = { + "intelligence.model": "qwen3:8b", + "intelligence.temperature": 0.3, + "engine.backend": "ollama", + "agent.type": "native_react", + "agent.max_turns": 10, + "tools.tool_set": ["calculator", "think"], + "learning.routing_policy": "grpo", + } + + run = OptimizationRun( + run_id="run-export", + search_space=_sample_search_space(), + best_trial=best, + status="completed", + ) + + path = tmp_path / "best_recipe.toml" + result_path = engine.export_best_recipe(run, path) + + assert result_path == path + assert path.exists() + assert run.best_recipe_path == str(path) + + # Verify it's valid TOML + with open(path, "rb") as fh: + data = tomllib.load(fh) + + assert data["recipe"]["name"] == "optimized-best" + assert "0.9500" in data["recipe"]["description"] + assert data["intelligence"]["model"] == "qwen3:8b" + assert data["intelligence"]["temperature"] == 0.3 + assert data["engine"]["key"] == "ollama" + assert data["agent"]["type"] == "native_react" + assert data["agent"]["max_turns"] == 10 + assert data["agent"]["tools"] == ["calculator", "think"] + assert data["learning"]["routing"] == "grpo" + + def test_export_creates_parent_dirs(self, tmp_path) -> None: + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=MagicMock(), + trial_runner=MagicMock(), + ) + + best = _sample_trial_result("best", accuracy=0.9) + run = OptimizationRun( + run_id="run-dirs", + search_space=_sample_search_space(), + best_trial=best, + status="completed", + ) + + path = tmp_path / "nested" / "deep" / "recipe.toml" + result_path = engine.export_best_recipe(run, path) + assert result_path.exists() + + def test_export_no_best_trial_raises(self, tmp_path) -> None: + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=MagicMock(), + trial_runner=MagicMock(), + ) + + run = OptimizationRun( + run_id="run-no-best", + search_space=_sample_search_space(), + best_trial=None, + ) + + try: + engine.export_best_recipe(run, tmp_path / "out.toml") + assert False, "Expected ValueError" + except ValueError as e: + assert "No best trial" in str(e) + + def test_export_minimal_params(self, tmp_path) -> None: + """Export with minimal params should still produce valid TOML.""" + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=MagicMock(), + trial_runner=MagicMock(), + ) + + config = TrialConfig(trial_id="min", params={}) + best = TrialResult( + trial_id="min", config=config, accuracy=0.5, + ) + run = OptimizationRun( + run_id="run-min", + search_space=_sample_search_space(), + best_trial=best, + ) + + path = tmp_path / "minimal.toml" + engine.export_best_recipe(run, path) + + with open(path, "rb") as fh: + data = tomllib.load(fh) + + assert data["recipe"]["name"] == "optimized-min" + + +# --------------------------------------------------------------------------- +# run() with store +# --------------------------------------------------------------------------- + + +class TestRunWithStore: + """Tests for run() with a real OptimizationStore.""" + + def test_saves_trials_and_run(self, tmp_path) -> None: + store = OptimizationStore(tmp_path / "opt.db") + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.propose_next.return_value = TrialConfig( + trial_id="t1", params={}, + ) + optimizer.analyze_trial.return_value = "analysis" + optimizer.optimizer_model = "m" + runner.benchmark = "b" + + runner.run_trial.side_effect = [ + _sample_trial_result("t0", accuracy=0.7), + _sample_trial_result("t1", accuracy=0.8), + ] + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + store=store, + max_trials=2, + ) + run = engine.run() + + # Verify trials were saved + trials = store.get_trials(run.run_id) + assert len(trials) == 2 + + # Verify run was saved + loaded_run = store.get_run(run.run_id) + assert loaded_run is not None + assert loaded_run.status == "completed" + assert len(loaded_run.trials) == 2 + + store.close() + + def test_no_store_does_not_error(self) -> None: + optimizer = MagicMock() + runner = MagicMock() + + optimizer.propose_initial.return_value = TrialConfig( + trial_id="t0", params={}, + ) + optimizer.analyze_trial.return_value = "ok" + optimizer.optimizer_model = "m" + runner.run_trial.return_value = _sample_trial_result("t0", 0.8) + runner.benchmark = "b" + + engine = OptimizationEngine( + search_space=_sample_search_space(), + llm_optimizer=optimizer, + trial_runner=runner, + store=None, + max_trials=1, + ) + run = engine.run() + assert run.status == "completed" + + +# --------------------------------------------------------------------------- +# config.py +# --------------------------------------------------------------------------- + + +class TestLoadOptimizeConfig: + """Tests for load_optimize_config.""" + + def test_loads_toml_file(self, tmp_path) -> None: + from openjarvis.optimize.config import load_optimize_config + + toml_content = b""" +[optimize] +max_trials = 15 +benchmark = "supergpqa" + +[[optimize.search]] +name = "agent.type" +type = "categorical" +values = ["simple", "orchestrator"] + +[optimize.fixed] +engine = "ollama" +""" + path = tmp_path / "optimize.toml" + path.write_bytes(toml_content) + + config = load_optimize_config(path) + assert config["optimize"]["max_trials"] == 15 + assert config["optimize"]["benchmark"] == "supergpqa" + assert len(config["optimize"]["search"]) == 1 + assert config["optimize"]["fixed"]["engine"] == "ollama" + + def test_file_not_found(self, tmp_path) -> None: + from openjarvis.optimize.config import load_optimize_config + + try: + load_optimize_config(tmp_path / "nonexistent.toml") + assert False, "Expected FileNotFoundError" + except FileNotFoundError: + pass + + def test_loads_string_path(self, tmp_path) -> None: + from openjarvis.optimize.config import load_optimize_config + + path = tmp_path / "test.toml" + path.write_bytes(b'[optimize]\nmax_trials = 5\n') + + config = load_optimize_config(str(path)) + assert config["optimize"]["max_trials"] == 5