mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 10:52:15 +00:00
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 <noreply@anthropic.com>
This commit is contained in:
Jon Saad-Falcon
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
44dc78ff7a
commit
d234471c03
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user