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
2026-03-04 23:56:30 +00:00
co-authored by Claude Opus 4.6
parent 44dc78ff7a
commit d234471c03
5 changed files with 1625 additions and 0 deletions
+34
View File
@@ -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"]
+282
View File
@@ -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"]
+301
View File
@@ -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"]