diff --git a/src/openjarvis/optimize/__init__.py b/src/openjarvis/optimize/__init__.py new file mode 100644 index 00000000..3c8bb243 --- /dev/null +++ b/src/openjarvis/optimize/__init__.py @@ -0,0 +1,20 @@ +"""Optimization framework for OpenJarvis configuration tuning.""" + +from openjarvis.optimize.search_space import DEFAULT_SEARCH_SPACE, build_search_space +from openjarvis.optimize.types import ( + OptimizationRun, + SearchDimension, + SearchSpace, + TrialConfig, + TrialResult, +) + +__all__ = [ + "SearchDimension", + "SearchSpace", + "TrialConfig", + "TrialResult", + "OptimizationRun", + "build_search_space", + "DEFAULT_SEARCH_SPACE", +] diff --git a/src/openjarvis/optimize/search_space.py b/src/openjarvis/optimize/search_space.py new file mode 100644 index 00000000..dcb1029e --- /dev/null +++ b/src/openjarvis/optimize/search_space.py @@ -0,0 +1,189 @@ +"""Search space builder and default search space for configuration optimization.""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from openjarvis.optimize.types import SearchDimension, SearchSpace + + +def build_search_space(config: Dict[str, Any]) -> SearchSpace: + """Build a SearchSpace from a TOML-style config dict. + + Expected format:: + + { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + "values": ["orchestrator", "native_react"], + "description": "Agent architecture", + }, + { + "name": "intelligence.temperature", + "type": "continuous", + "low": 0.0, + "high": 1.0, + "description": "Generation temperature", + }, + ], + "fixed": {"engine": "ollama", "model": "qwen3:8b"}, + "constraints": { + "rules": ["SimpleAgent should only have max_turns = 1"], + }, + } + } + """ + opt = config.get("optimize", {}) + search_entries: List[Dict[str, Any]] = opt.get("search", []) + fixed: Dict[str, Any] = dict(opt.get("fixed", {})) + constraints_sec = opt.get("constraints", {}) + constraints: List[str] = list(constraints_sec.get("rules", [])) + + dimensions: List[SearchDimension] = [] + for entry in search_entries: + # Infer pillar from the first segment of the dotted name + name = entry.get("name", "") + pillar = name.split(".")[0] if "." in name else "" + + dimensions.append( + SearchDimension( + name=name, + dim_type=entry.get("type", "categorical"), + values=list(entry.get("values", [])), + low=entry.get("low"), + high=entry.get("high"), + description=entry.get("description", ""), + pillar=pillar, + ) + ) + + return SearchSpace( + dimensions=dimensions, + fixed=fixed, + constraints=constraints, + ) + + +# --------------------------------------------------------------------------- +# Default search space covering all 5 pillars +# --------------------------------------------------------------------------- + +DEFAULT_SEARCH_SPACE = SearchSpace( + dimensions=[ + # Intelligence pillar + SearchDimension( + name="intelligence.model", + dim_type="categorical", + values=[ + "qwen3:8b", + "qwen3:4b", + "qwen3:1.7b", + "llama3.1:8b", + "llama3.1:70b", + "gemma2:9b", + "mistral:7b", + "deepseek-r1:8b", + ], + description="The LLM model to use for generation", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + description="Generation temperature (0 = deterministic, 1 = creative)", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.max_tokens", + dim_type="integer", + low=256, + high=8192, + description="Maximum tokens to generate per response", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.top_p", + dim_type="continuous", + low=0.0, + high=1.0, + description="Nucleus sampling probability threshold", + pillar="intelligence", + ), + SearchDimension( + name="intelligence.system_prompt", + dim_type="text", + description="System prompt to guide model behavior", + pillar="intelligence", + ), + # Engine pillar + SearchDimension( + name="engine.backend", + dim_type="categorical", + values=["ollama", "vllm", "sglang", "llamacpp", "mlx", "lmstudio"], + description="Inference engine backend", + pillar="engine", + ), + # Agent pillar + SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator", "native_react", "native_openhands"], + description="Agent architecture to use", + pillar="agent", + ), + SearchDimension( + name="agent.max_turns", + dim_type="integer", + low=1, + high=30, + description="Maximum number of agent reasoning turns", + pillar="agent", + ), + # Tools pillar + SearchDimension( + name="tools.tool_set", + dim_type="subset", + values=[ + "calculator", + "think", + "file_read", + "file_write", + "web_search", + "code_interpreter", + "llm", + "shell_exec", + "apply_patch", + "http_request", + "database_query", + ], + description="Set of tools available to the agent", + pillar="tools", + ), + # Learning pillar + SearchDimension( + name="learning.routing_policy", + dim_type="categorical", + values=["heuristic", "grpo", "bandit", "learned"], + description="Router policy for model/agent selection", + pillar="learning", + ), + ], + fixed={}, + constraints=[ + "SimpleAgent (agent.type='simple') should only have max_turns = 1", + "agent.max_turns must be >= 1", + "intelligence.temperature and intelligence.top_p " + "should not both be at extreme values", + ], +) + + +__all__ = [ + "build_search_space", + "DEFAULT_SEARCH_SPACE", +] diff --git a/src/openjarvis/optimize/types.py b/src/openjarvis/optimize/types.py new file mode 100644 index 00000000..3ee61d1d --- /dev/null +++ b/src/openjarvis/optimize/types.py @@ -0,0 +1,150 @@ +"""Core data types for the optimization framework.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from openjarvis.evals.core.types import RunSummary +from openjarvis.recipes.loader import Recipe + + +@dataclass(slots=True) +class SearchDimension: + """One tunable dimension in the config space.""" + + name: str # e.g. "agent.type", "intelligence.temperature" + dim_type: str # "categorical", "continuous", "integer", "subset", "text" + # categorical/subset: explicit options + values: List[Any] = field(default_factory=list) + low: Optional[float] = None # continuous/integer lower bound + high: Optional[float] = None # continuous/integer upper bound + description: str = "" # human-readable explanation for the LLM optimizer + pillar: str = "" # intelligence | engine | agent | tools | learning + + +@dataclass(slots=True) +class SearchSpace: + """The full space of configs the optimizer can propose.""" + + dimensions: List[SearchDimension] = field(default_factory=list) + fixed: Dict[str, Any] = field(default_factory=dict) # params NOT being optimized + constraints: List[str] = field(default_factory=list) # natural language constraints + + def to_prompt_description(self) -> str: + """Render search space as structured text for the LLM optimizer.""" + lines: List[str] = [] + lines.append("# Search Space") + lines.append("") + + # Group dimensions by pillar + by_pillar: Dict[str, List[SearchDimension]] = {} + for dim in self.dimensions: + key = dim.pillar or "other" + by_pillar.setdefault(key, []).append(dim) + + for pillar, dims in sorted(by_pillar.items()): + lines.append(f"## {pillar.title()}") + for dim in dims: + lines.append(f"- **{dim.name}** ({dim.dim_type})") + if dim.description: + lines.append(f" Description: {dim.description}") + if dim.dim_type in ("categorical", "subset"): + lines.append(f" Options: {dim.values}") + elif dim.dim_type in ("continuous", "integer"): + lines.append(f" Range: [{dim.low}, {dim.high}]") + elif dim.dim_type == "text": + lines.append(" Free-form text") + lines.append("") + + if self.fixed: + lines.append("## Fixed Parameters") + for k, v in sorted(self.fixed.items()): + lines.append(f"- {k} = {v}") + lines.append("") + + if self.constraints: + lines.append("## Constraints") + for c in self.constraints: + lines.append(f"- {c}") + lines.append("") + + return "\n".join(lines) + + +# Mapping from dotted param names to Recipe constructor fields. +_PARAM_TO_RECIPE: Dict[str, str] = { + "intelligence.model": "model", + "intelligence.temperature": "temperature", + "intelligence.quantization": "quantization", + "engine.backend": "engine_key", + "agent.type": "agent_type", + "agent.max_turns": "max_turns", + "agent.system_prompt": "system_prompt", + "tools.tool_set": "tools", + "learning.routing_policy": "routing_policy", + "learning.agent_policy": "agent_policy", +} + + +@dataclass(slots=True) +class TrialConfig: + """A single candidate configuration proposed by the optimizer.""" + + trial_id: str + params: Dict[str, Any] = field(default_factory=dict) # dotted keys -> values + reasoning: str = "" # optimizer's explanation + + def to_recipe(self) -> Recipe: + """Map params back to Recipe fields.""" + kwargs: Dict[str, Any] = {} + for dotted_key, value in self.params.items(): + recipe_field = _PARAM_TO_RECIPE.get(dotted_key) + if recipe_field is not None: + kwargs[recipe_field] = value + + return Recipe( + name=f"trial-{self.trial_id}", + **kwargs, + ) + + +@dataclass(slots=True) +class TrialResult: + """Result of evaluating a trial, with both scalar and textual feedback.""" + + trial_id: str + config: TrialConfig + accuracy: float = 0.0 + mean_latency_seconds: float = 0.0 + total_cost_usd: float = 0.0 + total_energy_joules: float = 0.0 + total_tokens: int = 0 + samples_evaluated: int = 0 + analysis: str = "" + failure_modes: List[str] = field(default_factory=list) + per_sample_feedback: List[Dict[str, Any]] = field(default_factory=list) + summary: Optional[RunSummary] = None + + +@dataclass(slots=True) +class OptimizationRun: + """Complete optimization session.""" + + run_id: str + search_space: SearchSpace + trials: List[TrialResult] = field(default_factory=list) + best_trial: Optional[TrialResult] = None + best_recipe_path: Optional[str] = None + status: str = "running" # running | completed | failed + optimizer_model: str = "" + benchmark: str = "" + + +__all__ = [ + "SearchDimension", + "SearchSpace", + "TrialConfig", + "TrialResult", + "OptimizationRun", +] diff --git a/tests/test_optimize_types.py b/tests/test_optimize_types.py new file mode 100644 index 00000000..bbfe4697 --- /dev/null +++ b/tests/test_optimize_types.py @@ -0,0 +1,488 @@ +"""Tests for openjarvis.optimize.types module.""" + +from __future__ import annotations + +from openjarvis.optimize.types import ( + OptimizationRun, + SearchDimension, + SearchSpace, + TrialConfig, + TrialResult, +) +from openjarvis.recipes.loader import Recipe + +# --------------------------------------------------------------------------- +# SearchDimension +# --------------------------------------------------------------------------- + + +class TestSearchDimension: + """Tests for SearchDimension dataclass.""" + + def test_categorical_dimension(self) -> None: + dim = SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator", "native_react"], + description="Agent architecture", + pillar="agent", + ) + assert dim.name == "agent.type" + assert dim.dim_type == "categorical" + assert dim.values == [ + "simple", + "orchestrator", + "native_react", + ] + assert dim.description == "Agent architecture" + assert dim.pillar == "agent" + assert dim.low is None + assert dim.high is None + + def test_continuous_dimension(self) -> None: + dim = SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + description="Generation temperature", + pillar="intelligence", + ) + assert dim.dim_type == "continuous" + assert dim.low == 0.0 + assert dim.high == 1.0 + assert dim.values == [] + + def test_integer_dimension(self) -> None: + dim = SearchDimension( + name="agent.max_turns", + dim_type="integer", + low=1, + high=30, + pillar="agent", + ) + assert dim.dim_type == "integer" + assert dim.low == 1 + assert dim.high == 30 + + def test_subset_dimension(self) -> None: + dim = SearchDimension( + name="tools.tool_set", + dim_type="subset", + values=["calculator", "think", "web_search"], + pillar="tools", + ) + assert dim.dim_type == "subset" + assert len(dim.values) == 3 + + def test_text_dimension(self) -> None: + dim = SearchDimension( + name="intelligence.system_prompt", + dim_type="text", + description="System prompt to guide model behavior", + pillar="intelligence", + ) + assert dim.dim_type == "text" + assert dim.values == [] + assert dim.low is None + assert dim.high is None + + def test_defaults(self) -> None: + dim = SearchDimension(name="x", dim_type="categorical") + assert dim.values == [] + assert dim.low is None + assert dim.high is None + assert dim.description == "" + assert dim.pillar == "" + + def test_mutable_default_isolation(self) -> None: + """Ensure mutable defaults are independent.""" + dim1 = SearchDimension(name="a", dim_type="categorical") + dim2 = SearchDimension(name="b", dim_type="categorical") + dim1.values.append("x") + assert dim2.values == [] + + +# --------------------------------------------------------------------------- +# SearchSpace +# --------------------------------------------------------------------------- + + +class TestSearchSpace: + """Tests for SearchSpace dataclass.""" + + def test_empty_search_space(self) -> None: + space = SearchSpace() + assert space.dimensions == [] + assert space.fixed == {} + assert space.constraints == [] + + def test_search_space_with_dimensions(self) -> None: + dims = [ + SearchDimension( + name="a", + dim_type="categorical", + values=["x", "y"], + ), + SearchDimension( + name="b", + dim_type="continuous", + low=0.0, + high=1.0, + ), + ] + space = SearchSpace( + dimensions=dims, + fixed={"engine": "ollama"}, + constraints=["a must not be x when b > 0.5"], + ) + assert len(space.dimensions) == 2 + assert space.fixed == {"engine": "ollama"} + assert len(space.constraints) == 1 + + def test_to_prompt_description_has_header(self) -> None: + space = SearchSpace( + dimensions=[ + SearchDimension( + name="agent.type", + dim_type="categorical", + values=["simple", "orchestrator"], + description="Agent kind", + pillar="agent", + ), + ], + ) + desc = space.to_prompt_description() + assert "# Search Space" in desc + assert "## Agent" in desc + assert "agent.type" in desc + assert "categorical" in desc + assert "Agent kind" in desc + assert "simple" in desc + assert "orchestrator" in desc + + def test_to_prompt_description_continuous(self) -> None: + space = SearchSpace( + dimensions=[ + SearchDimension( + name="intelligence.temperature", + dim_type="continuous", + low=0.0, + high=1.0, + pillar="intelligence", + ), + ], + ) + desc = space.to_prompt_description() + assert "Range:" in desc + assert "0.0" in desc + assert "1.0" in desc + + def test_to_prompt_description_text(self) -> None: + space = SearchSpace( + dimensions=[ + SearchDimension( + name="intelligence.system_prompt", + dim_type="text", + pillar="intelligence", + ), + ], + ) + desc = space.to_prompt_description() + assert "Free-form text" in desc + + def test_to_prompt_description_fixed_params(self) -> None: + space = SearchSpace( + dimensions=[], + fixed={"engine": "ollama", "model": "qwen3:8b"}, + ) + desc = space.to_prompt_description() + assert "## Fixed Parameters" in desc + assert "engine = ollama" in desc + assert "model = qwen3:8b" in desc + + def test_to_prompt_description_constraints(self) -> None: + space = SearchSpace( + dimensions=[], + constraints=[ + "max_turns must be >= 1", + "temperature must be <= 1.0", + ], + ) + desc = space.to_prompt_description() + assert "## Constraints" in desc + assert "max_turns must be >= 1" in desc + assert "temperature must be <= 1.0" in desc + + def test_to_prompt_description_groups_by_pillar(self) -> None: + space = SearchSpace( + dimensions=[ + SearchDimension( + name="a.x", + dim_type="categorical", + values=["1"], + pillar="agent", + ), + SearchDimension( + name="i.y", + dim_type="continuous", + low=0, + high=1, + pillar="intelligence", + ), + SearchDimension( + name="a.z", + dim_type="integer", + low=1, + high=10, + pillar="agent", + ), + ], + ) + desc = space.to_prompt_description() + # Both agent dimensions under the Agent header + assert "## Agent" in desc + assert "## Intelligence" in desc + + def test_mutable_default_isolation(self) -> None: + s1 = SearchSpace() + s2 = SearchSpace() + s1.dimensions.append( + SearchDimension(name="x", dim_type="categorical"), + ) + s1.fixed["key"] = "val" + s1.constraints.append("rule") + assert s2.dimensions == [] + assert s2.fixed == {} + assert s2.constraints == [] + + +# --------------------------------------------------------------------------- +# TrialConfig +# --------------------------------------------------------------------------- + + +class TestTrialConfig: + """Tests for TrialConfig dataclass.""" + + def test_creation(self) -> None: + tc = TrialConfig( + trial_id="t1", + params={ + "agent.type": "orchestrator", + "intelligence.temperature": 0.7, + }, + reasoning="Higher temperature for creativity", + ) + assert tc.trial_id == "t1" + assert tc.params["agent.type"] == "orchestrator" + assert tc.reasoning == "Higher temperature for creativity" + + def test_defaults(self) -> None: + tc = TrialConfig(trial_id="t0") + assert tc.params == {} + assert tc.reasoning == "" + + def test_to_recipe_basic(self) -> None: + tc = TrialConfig( + trial_id="abc", + params={ + "intelligence.model": "qwen3:8b", + "intelligence.temperature": 0.5, + "engine.backend": "ollama", + "agent.type": "native_react", + "agent.max_turns": 10, + "tools.tool_set": ["calculator", "think"], + "learning.routing_policy": "grpo", + }, + ) + recipe = tc.to_recipe() + assert isinstance(recipe, Recipe) + assert recipe.name == "trial-abc" + assert recipe.model == "qwen3:8b" + assert recipe.temperature == 0.5 + assert recipe.engine_key == "ollama" + assert recipe.agent_type == "native_react" + assert recipe.max_turns == 10 + assert recipe.tools == ["calculator", "think"] + assert recipe.routing_policy == "grpo" + + def test_to_recipe_partial_params(self) -> None: + tc = TrialConfig( + trial_id="partial", + params={"intelligence.temperature": 0.3}, + ) + recipe = tc.to_recipe() + assert recipe.temperature == 0.3 + assert recipe.model is None + assert recipe.engine_key is None + assert recipe.agent_type is None + + def test_to_recipe_unknown_params_ignored(self) -> None: + tc = TrialConfig( + trial_id="unk", + params={"some.unknown.param": "value"}, + ) + recipe = tc.to_recipe() + assert recipe.name == "trial-unk" + # Unknown params should not cause an error + + def test_to_recipe_system_prompt(self) -> None: + tc = TrialConfig( + trial_id="sp", + params={ + "agent.system_prompt": "You are a helpful assistant.", + }, + ) + recipe = tc.to_recipe() + assert recipe.system_prompt == "You are a helpful assistant." + + def test_to_recipe_quantization(self) -> None: + tc = TrialConfig( + trial_id="q", + params={"intelligence.quantization": "q4_K_M"}, + ) + recipe = tc.to_recipe() + assert recipe.quantization == "q4_K_M" + + def test_to_recipe_agent_policy(self) -> None: + tc = TrialConfig( + trial_id="ap", + params={"learning.agent_policy": "icl_updater"}, + ) + recipe = tc.to_recipe() + assert recipe.agent_policy == "icl_updater" + + def test_mutable_default_isolation(self) -> None: + tc1 = TrialConfig(trial_id="a") + tc2 = TrialConfig(trial_id="b") + tc1.params["x"] = 1 + assert "x" not in tc2.params + + +# --------------------------------------------------------------------------- +# TrialResult +# --------------------------------------------------------------------------- + + +class TestTrialResult: + """Tests for TrialResult dataclass.""" + + def test_creation_with_defaults(self) -> None: + config = TrialConfig(trial_id="t1") + result = TrialResult(trial_id="t1", config=config) + assert result.trial_id == "t1" + assert result.accuracy == 0.0 + assert result.mean_latency_seconds == 0.0 + assert result.total_cost_usd == 0.0 + assert result.total_energy_joules == 0.0 + assert result.total_tokens == 0 + assert result.samples_evaluated == 0 + assert result.analysis == "" + assert result.failure_modes == [] + assert result.per_sample_feedback == [] + assert result.summary is None + + def test_creation_with_values(self) -> None: + config = TrialConfig( + trial_id="t2", + params={"agent.type": "orchestrator"}, + ) + result = TrialResult( + trial_id="t2", + config=config, + accuracy=0.85, + mean_latency_seconds=1.2, + total_cost_usd=0.05, + total_energy_joules=150.0, + total_tokens=5000, + samples_evaluated=100, + analysis="Good accuracy, moderate latency", + failure_modes=["timeout on long inputs"], + per_sample_feedback=[ + {"id": "s1", "correct": True}, + ], + ) + assert result.accuracy == 0.85 + assert result.mean_latency_seconds == 1.2 + assert result.total_cost_usd == 0.05 + assert result.total_energy_joules == 150.0 + assert result.total_tokens == 5000 + assert result.samples_evaluated == 100 + assert result.analysis == "Good accuracy, moderate latency" + assert result.failure_modes == ["timeout on long inputs"] + assert len(result.per_sample_feedback) == 1 + + def test_mutable_default_isolation(self) -> None: + c1 = TrialConfig(trial_id="a") + c2 = TrialConfig(trial_id="b") + r1 = TrialResult(trial_id="a", config=c1) + r2 = TrialResult(trial_id="b", config=c2) + r1.failure_modes.append("error") + r1.per_sample_feedback.append({"x": 1}) + assert r2.failure_modes == [] + assert r2.per_sample_feedback == [] + + +# --------------------------------------------------------------------------- +# OptimizationRun +# --------------------------------------------------------------------------- + + +class TestOptimizationRun: + """Tests for OptimizationRun dataclass.""" + + def test_creation_defaults(self) -> None: + space = SearchSpace() + run = OptimizationRun( + run_id="run-001", + search_space=space, + ) + assert run.run_id == "run-001" + assert run.search_space is space + assert run.trials == [] + assert run.best_trial is None + assert run.best_recipe_path is None + assert run.status == "running" + assert run.optimizer_model == "" + assert run.benchmark == "" + + def test_creation_with_values(self) -> None: + space = SearchSpace() + config = TrialConfig( + trial_id="t1", + params={"agent.type": "orchestrator"}, + ) + result = TrialResult( + trial_id="t1", + config=config, + accuracy=0.9, + ) + run = OptimizationRun( + run_id="run-002", + search_space=space, + trials=[result], + best_trial=result, + best_recipe_path="/tmp/best.toml", + status="completed", + optimizer_model="gpt-5-mini", + benchmark="supergpqa", + ) + assert len(run.trials) == 1 + assert run.best_trial is result + assert run.best_recipe_path == "/tmp/best.toml" + assert run.status == "completed" + assert run.optimizer_model == "gpt-5-mini" + assert run.benchmark == "supergpqa" + + def test_mutable_default_isolation(self) -> None: + space = SearchSpace() + r1 = OptimizationRun(run_id="a", search_space=space) + r2 = OptimizationRun(run_id="b", search_space=space) + r1.trials.append( + TrialResult( + trial_id="t", + config=TrialConfig(trial_id="t"), + ), + ) + assert r2.trials == [] diff --git a/tests/test_search_space.py b/tests/test_search_space.py new file mode 100644 index 00000000..903ddca3 --- /dev/null +++ b/tests/test_search_space.py @@ -0,0 +1,442 @@ +"""Tests for openjarvis.optimize.search_space module.""" + +from __future__ import annotations + +from openjarvis.optimize.search_space import ( + DEFAULT_SEARCH_SPACE, + build_search_space, +) +from openjarvis.optimize.types import SearchSpace + +# --------------------------------------------------------------------------- +# build_search_space +# --------------------------------------------------------------------------- + + +class TestBuildSearchSpace: + """Tests for the build_search_space() factory function.""" + + def test_basic_build(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + "values": [ + "orchestrator", + "native_react", + ], + "description": "Agent architecture", + }, + ], + "fixed": { + "engine": "ollama", + "model": "qwen3:8b", + }, + "constraints": { + "rules": [ + "SimpleAgent should only have " + "max_turns = 1", + ], + }, + }, + } + space = build_search_space(config) + assert isinstance(space, SearchSpace) + assert len(space.dimensions) == 1 + dim = space.dimensions[0] + assert dim.name == "agent.type" + assert dim.dim_type == "categorical" + assert dim.values == ["orchestrator", "native_react"] + assert dim.description == "Agent architecture" + assert dim.pillar == "agent" + + def test_fixed_params_preserved(self) -> None: + config = { + "optimize": { + "search": [], + "fixed": { + "engine": "ollama", + "model": "qwen3:8b", + }, + }, + } + space = build_search_space(config) + assert space.fixed == { + "engine": "ollama", + "model": "qwen3:8b", + } + + def test_constraints_parsed(self) -> None: + config = { + "optimize": { + "search": [], + "constraints": { + "rules": [ + "max_turns must be >= 1", + "temperature should be <= 1.0", + ], + }, + }, + } + space = build_search_space(config) + assert len(space.constraints) == 2 + assert "max_turns must be >= 1" in space.constraints + assert "temperature should be <= 1.0" in space.constraints + + def test_continuous_dimension_build(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "intelligence.temperature", + "type": "continuous", + "low": 0.0, + "high": 1.0, + "description": "Generation temperature", + }, + ], + }, + } + space = build_search_space(config) + dim = space.dimensions[0] + assert dim.dim_type == "continuous" + assert dim.low == 0.0 + assert dim.high == 1.0 + assert dim.pillar == "intelligence" + + def test_integer_dimension_build(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "agent.max_turns", + "type": "integer", + "low": 1, + "high": 30, + }, + ], + }, + } + space = build_search_space(config) + dim = space.dimensions[0] + assert dim.dim_type == "integer" + assert dim.low == 1 + assert dim.high == 30 + + def test_subset_dimension_build(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "tools.tool_set", + "type": "subset", + "values": [ + "calculator", + "think", + "web_search", + ], + }, + ], + }, + } + space = build_search_space(config) + dim = space.dimensions[0] + assert dim.dim_type == "subset" + assert dim.values == [ + "calculator", + "think", + "web_search", + ] + assert dim.pillar == "tools" + + def test_text_dimension_build(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "intelligence.system_prompt", + "type": "text", + "description": "System prompt", + }, + ], + }, + } + space = build_search_space(config) + dim = space.dimensions[0] + assert dim.dim_type == "text" + assert dim.values == [] + assert dim.pillar == "intelligence" + + def test_multiple_dimensions(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + "values": ["simple"], + }, + { + "name": "intelligence.temperature", + "type": "continuous", + "low": 0.0, + "high": 1.0, + }, + { + "name": "tools.tool_set", + "type": "subset", + "values": ["calculator"], + }, + ], + }, + } + space = build_search_space(config) + assert len(space.dimensions) == 3 + + def test_empty_config(self) -> None: + space = build_search_space({}) + assert space.dimensions == [] + assert space.fixed == {} + assert space.constraints == [] + + def test_empty_optimize_section(self) -> None: + space = build_search_space({"optimize": {}}) + assert space.dimensions == [] + assert space.fixed == {} + assert space.constraints == [] + + def test_pillar_inferred_from_name(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "learning.routing_policy", + "type": "categorical", + "values": ["grpo"], + }, + { + "name": "engine.backend", + "type": "categorical", + "values": ["ollama"], + }, + ], + }, + } + space = build_search_space(config) + assert space.dimensions[0].pillar == "learning" + assert space.dimensions[1].pillar == "engine" + + def test_no_dot_in_name_gives_empty_pillar(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "standalone", + "type": "categorical", + "values": ["a"], + }, + ], + }, + } + space = build_search_space(config) + assert space.dimensions[0].pillar == "" + + def test_missing_description_defaults_empty(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + "values": ["simple"], + }, + ], + }, + } + space = build_search_space(config) + assert space.dimensions[0].description == "" + + def test_missing_values_defaults_empty_list(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "agent.type", + "type": "categorical", + }, + ], + }, + } + space = build_search_space(config) + assert space.dimensions[0].values == [] + + def test_missing_constraints_section(self) -> None: + config = { + "optimize": { + "search": [ + { + "name": "a.b", + "type": "categorical", + "values": ["x"], + }, + ], + "fixed": {"k": "v"}, + }, + } + space = build_search_space(config) + assert space.constraints == [] + + +# --------------------------------------------------------------------------- +# DEFAULT_SEARCH_SPACE +# --------------------------------------------------------------------------- + + +_DIMS = DEFAULT_SEARCH_SPACE.dimensions + + +def _find_dim(name: str): + """Helper to find a dimension by name.""" + return next(d for d in _DIMS if d.name == name) + + +class TestDefaultSearchSpace: + """Tests for the DEFAULT_SEARCH_SPACE module-level constant.""" + + def test_is_search_space(self) -> None: + assert isinstance(DEFAULT_SEARCH_SPACE, SearchSpace) + + def test_has_all_five_pillars(self) -> None: + pillars = {dim.pillar for dim in _DIMS} + assert "intelligence" in pillars + assert "engine" in pillars + assert "agent" in pillars + assert "tools" in pillars + assert "learning" in pillars + + def test_intelligence_dimensions(self) -> None: + intel_dims = [d for d in _DIMS if d.pillar == "intelligence"] + intel_names = {d.name for d in intel_dims} + assert "intelligence.model" in intel_names + assert "intelligence.temperature" in intel_names + assert "intelligence.max_tokens" in intel_names + assert "intelligence.top_p" in intel_names + assert "intelligence.system_prompt" in intel_names + + def test_intelligence_model_is_categorical(self) -> None: + dim = _find_dim("intelligence.model") + assert dim.dim_type == "categorical" + assert len(dim.values) > 0 + + def test_intelligence_temperature_range(self) -> None: + dim = _find_dim("intelligence.temperature") + assert dim.dim_type == "continuous" + assert dim.low == 0.0 + assert dim.high == 1.0 + + def test_intelligence_max_tokens_range(self) -> None: + dim = _find_dim("intelligence.max_tokens") + assert dim.dim_type == "integer" + assert dim.low == 256 + assert dim.high == 8192 + + def test_intelligence_system_prompt_is_text(self) -> None: + dim = _find_dim("intelligence.system_prompt") + assert dim.dim_type == "text" + + def test_engine_backend_options(self) -> None: + dim = _find_dim("engine.backend") + assert dim.dim_type == "categorical" + expected = { + "ollama", "vllm", "sglang", + "llamacpp", "mlx", "lmstudio", + } + assert set(dim.values) == expected + + def test_agent_type_options(self) -> None: + dim = _find_dim("agent.type") + assert dim.dim_type == "categorical" + expected = { + "simple", "orchestrator", + "native_react", "native_openhands", + } + assert set(dim.values) == expected + + def test_agent_max_turns_range(self) -> None: + dim = _find_dim("agent.max_turns") + assert dim.dim_type == "integer" + assert dim.low == 1 + assert dim.high == 30 + + def test_tools_tool_set_is_subset(self) -> None: + dim = _find_dim("tools.tool_set") + assert dim.dim_type == "subset" + assert "calculator" in dim.values + assert "think" in dim.values + + def test_learning_routing_policy(self) -> None: + dim = _find_dim("learning.routing_policy") + assert dim.dim_type == "categorical" + expected = {"heuristic", "grpo", "bandit", "learned"} + assert set(dim.values) == expected + + def test_has_constraints(self) -> None: + assert len(DEFAULT_SEARCH_SPACE.constraints) > 0 + + def test_all_dimensions_have_descriptions(self) -> None: + for dim in _DIMS: + assert dim.description != "", ( + f"Dimension {dim.name} has no description" + ) + + def test_all_dimensions_have_pillars(self) -> None: + for dim in _DIMS: + assert dim.pillar != "", ( + f"Dimension {dim.name} has no pillar" + ) + + +# --------------------------------------------------------------------------- +# to_prompt_description rendering +# --------------------------------------------------------------------------- + + +class TestToPromptDescription: + """Tests for SearchSpace.to_prompt_description().""" + + def test_default_space_renders(self) -> None: + desc = DEFAULT_SEARCH_SPACE.to_prompt_description() + assert isinstance(desc, str) + assert len(desc) > 100 + + def test_all_dimensions_appear_in_description(self) -> None: + desc = DEFAULT_SEARCH_SPACE.to_prompt_description() + for dim in _DIMS: + assert dim.name in desc, ( + f"Dimension {dim.name} not in description" + ) + + def test_all_pillar_headers_in_description(self) -> None: + desc = DEFAULT_SEARCH_SPACE.to_prompt_description() + for pillar in ( + "Intelligence", "Engine", "Agent", + "Tools", "Learning", + ): + assert f"## {pillar}" in desc, ( + f"Pillar header {pillar} not in description" + ) + + def test_constraints_in_description(self) -> None: + desc = DEFAULT_SEARCH_SPACE.to_prompt_description() + assert "## Constraints" in desc + for constraint in DEFAULT_SEARCH_SPACE.constraints: + assert constraint in desc + + def test_empty_space_renders(self) -> None: + space = SearchSpace() + desc = space.to_prompt_description() + assert "# Search Space" in desc + assert "## Fixed Parameters" not in desc + assert "## Constraints" not in desc