From 18135d97024ca33e7c81f24ec4aec43fd9d962d5 Mon Sep 17 00:00:00 2001 From: Ousama Ben Younes Date: Mon, 27 Jul 2026 19:12:05 +0000 Subject: [PATCH] feat(config): add per-request token budget --- docs/advanced/configuration.mdx | 5 +++++ strix/config/settings.py | 1 + strix/core/inputs.py | 2 ++ strix/core/runner.py | 1 + tests/test_config_loader.py | 7 ++++++- tests/test_inputs.py | 10 ++++++++++ tests/test_runner_rate_limit.py | 1 + tests/test_runner_root_prompt.py | 25 ++++++++++++++++++++++++- 8 files changed, 50 insertions(+), 2 deletions(-) diff --git a/docs/advanced/configuration.mdx b/docs/advanced/configuration.mdx index f1542b75..5dbfb9b4 100644 --- a/docs/advanced/configuration.mdx +++ b/docs/advanced/configuration.mdx @@ -31,6 +31,11 @@ Configure Strix using environment variables or a config file. Request timeout in seconds for LLM calls. + + Optional maximum output tokens for each agent LLM request. Leave unset to use + the provider/model default. + + Maximum number of retries for LLM API calls on transient failures. diff --git a/strix/config/settings.py b/strix/config/settings.py index 42a2c97e..cfe156e8 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -56,6 +56,7 @@ class LlmSettings(BaseSettings): default=False, alias="LLM_DISABLE_STREAMING", ) + max_tokens: int | None = Field(default=None, gt=0, alias="STRIX_LLM_MAX_TOKENS") timeout: int = Field(default=300, alias="LLM_TIMEOUT") stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT") max_tool_calls_per_turn: int = Field( diff --git a/strix/core/inputs.py b/strix/core/inputs.py index ea72abb7..54b21182 100644 --- a/strix/core/inputs.py +++ b/strix/core/inputs.py @@ -235,12 +235,14 @@ def make_model_settings( prompt_cache: bool = True, extra_headers: dict[str, str] | None = None, has_tools: bool = True, + max_tokens: int | None = None, ) -> ModelSettings: headers = _request_headers(model_name, extra_headers) model_settings = ModelSettings( parallel_tool_calls=False if has_tools else None, retry=DEFAULT_MODEL_RETRY, include_usage=True, + max_tokens=max_tokens, extra_args=request_timeout_extra_args(request_timeout), extra_headers=headers, ) diff --git a/strix/core/runner.py b/strix/core/runner.py index b4afdfaf..22ea4ea5 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -267,6 +267,7 @@ async def run_strix_scan( request_timeout=settings.llm.timeout, prompt_cache=settings.llm.prompt_cache, extra_headers=settings.llm.extra_headers, + max_tokens=settings.llm.max_tokens, ) run_config = RunConfig( model=resolved_model, diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index e83ab119..e0d324f4 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -10,7 +10,7 @@ from pydantic import AliasChoices, Field, ValidationError from pydantic.fields import FieldInfo from strix.config import loader -from strix.config.settings import ContextSettings +from strix.config.settings import ContextSettings, LlmSettings if TYPE_CHECKING: @@ -28,6 +28,7 @@ _LLM_ENV_KEYS = [ "OLLAMA_API_BASE", "STRIX_REASONING_EFFORT", "STRIX_FORCE_REQUIRED_TOOL_CHOICE", + "STRIX_LLM_MAX_TOKENS", "LLM_TIMEOUT", "PERPLEXITY_API_KEY", # RuntimeSettings @@ -129,6 +130,10 @@ def test_tool_output_max_bytes_accepts_floor() -> None: assert ContextSettings(STRIX_TOOL_OUTPUT_MAX_BYTES=1024).tool_output_max_bytes == 1024 +def test_llm_max_tokens_env_alias() -> None: + assert LlmSettings(STRIX_LLM_MAX_TOKENS=12_000).max_tokens == 12_000 + + # --------------------------------------------------------------------------- # # _aliases_for # --------------------------------------------------------------------------- # diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 2ff9a603..03e1231d 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -320,6 +320,16 @@ def test_make_model_settings_sets_request_timeout() -> None: assert settings.extra_args["timeout"] == 300.0 +def test_make_model_settings_sets_configured_token_budget() -> None: + settings = make_model_settings( + "none", + model_name="gpt-4o", + max_tokens=12_000, + ) + + assert settings.max_tokens == 12_000 + + def test_make_model_settings_omits_timeout_when_unset() -> None: settings = make_model_settings("none", model_name="gpt-4o") diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 3110ae2c..ae73b85f 100644 --- a/tests/test_runner_rate_limit.py +++ b/tests/test_runner_rate_limit.py @@ -42,6 +42,7 @@ async def test_persistent_rate_limit_stops_gracefully( timeout=300, prompt_cache=True, extra_headers=None, + max_tokens=None, ), runtime=types.SimpleNamespace(max_context_images=3), ) diff --git a/tests/test_runner_root_prompt.py b/tests/test_runner_root_prompt.py index 2c346203..60d9e206 100644 --- a/tests/test_runner_root_prompt.py +++ b/tests/test_runner_root_prompt.py @@ -50,6 +50,7 @@ def _patch_engine_scaffold( timeout=300, prompt_cache=True, extra_headers=None, + max_tokens=12_000, ), runtime=types.SimpleNamespace(max_context_images=3), ) @@ -75,10 +76,15 @@ def _patch_engine_scaffold( monkeypatch.setattr(runner, "build_root_task", lambda _scan_config: "task") monkeypatch.setattr(runner, "build_scope_context", lambda _scan_config: scope_context) - monkeypatch.setattr(runner, "make_model_settings", lambda *_args, **_kwargs: ModelSettings()) captured: dict[str, Any] = {} + def _make_model_settings(*_args: Any, **kwargs: Any) -> ModelSettings: + captured["model_settings_kwargs"] = kwargs + return ModelSettings() + + monkeypatch.setattr(runner, "make_model_settings", _make_model_settings) + def _build_strix_agent(**kwargs: Any) -> object: if kwargs.get("is_root") and "kwargs" not in captured: captured["kwargs"] = kwargs @@ -196,3 +202,20 @@ async def test_unknown_tool_calls_are_returned_to_the_model( ) assert captured["run_config"].tool_not_found_behavior == "return_error_to_model" + + +@pytest.mark.asyncio +async def test_llm_max_tokens_flows_into_model_settings( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, +) -> None: + captured = _patch_engine_scaffold(monkeypatch, tmp_path, {"scope": "built-in"}) + + await runner.run_strix_scan( + scan_config={"targets": [], "scan_mode": "deep"}, + scan_id="scan-token-budget", + image="img", + coordinator=AgentCoordinator(), + ) + + assert captured["model_settings_kwargs"]["max_tokens"] == 12_000