mirror of
https://github.com/usestrix/strix.git
synced 2026-08-28 05:25:00 +00:00
feat(config): add per-request token budget
This commit is contained in:
parent
8551339130
commit
18135d9702
8 changed files with 50 additions and 2 deletions
|
|
@ -31,6 +31,11 @@ Configure Strix using environment variables or a config file.
|
|||
Request timeout in seconds for LLM calls.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="STRIX_LLM_MAX_TOKENS" type="integer">
|
||||
Optional maximum output tokens for each agent LLM request. Leave unset to use
|
||||
the provider/model default.
|
||||
</ParamField>
|
||||
|
||||
<ParamField path="STRIX_LLM_MAX_RETRIES" default="5" type="integer">
|
||||
Maximum number of retries for LLM API calls on transient failures.
|
||||
</ParamField>
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue