diff --git a/README.md b/README.md index 7ea2a34..680cc6c 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,7 @@ If you find AttackGen useful, please consider starring the repository on GitHub. | What's new? | Why is it useful? | | ----------- | ----------------- | | Assistant Refines the Detection & Response Output | - Purple-Team Editing: The AttackGen Assistant can now refine the **Detection & Response** narrative, not just the scenario. When a scenario is generated with the 🟣 Purple-team narrative toggle (Threat Group / Custom pages), an **Editing:** selector appears on the Assistant page so the chat can target either the scenario or the defender's walkthrough — previously the Assistant only had the scenario text.

- Combined Editing Mode: A third **Scenario + Detection & Response** target refines both together, so a cross-cutting change — a different industry, threat actor, technique or timeline — is applied consistently across the attacker's scenario and the defender's narrative rather than being made twice and drifting apart.

- Deterministic Facts Untouched: Only the LLM narrative is a refinement target; the local STIX detection join is never LLM-edited. Each editing mode keeps its own chat history, and the combined pass is tagged `purple_team_narrative` in LangSmith. | +| Simpler Model Registry | - One Less Thing Per Model: The `max_completion_tokens` kwarg and the reasoning-model temperature rule are now routed from the **provider** rather than a per-model `uses_max_completion_tokens` flag — every current OpenAI chat model uses them, so the flag duplicated a fact the provider key already captured. Adding a new OpenAI model stays a one-entry edit to `MODELS`, with no extra flag to remember.

- No Behaviour Change: An internal simplification only — the wrapper builds identical LiteLLM kwargs for every current model. | ### v0.13.1 | What's new? | Why is it useful? | diff --git a/core/llm.py b/core/llm.py index 9a8dceb..8201a81 100644 --- a/core/llm.py +++ b/core/llm.py @@ -28,7 +28,6 @@ from core.models import ( get_litellm_prefix, get_provider, - model_uses_completion_tokens, ) from core.schemas import LLMConfig @@ -78,10 +77,11 @@ def _build_litellm_kwargs(config: LLMConfig) -> dict: prefix = get_litellm_prefix(config.provider) model = prefix + config.model_name - # OpenAI reasoning models (gpt-5.x) use max_completion_tokens instead of - # max_tokens, and reject any temperature other than the default (1) — sending - # temperature=0.7 to them is a 400 BadRequest. Detect once; drive both below. - is_openai_reasoning = model_uses_completion_tokens(config.provider, config.model_name) + # OpenAI's chat models use max_completion_tokens instead of the deprecated + # max_tokens, and the gpt-5.x reasoning family rejects any temperature other + # than the default (1) — sending temperature=0.7 to them is a 400 BadRequest. + # Both are properties of the OpenAI provider, so route on it directly. + is_openai = config.provider == "OpenAI API" kwargs: dict = { "model": model, @@ -91,7 +91,7 @@ def _build_litellm_kwargs(config: LLMConfig) -> dict: } # Only send temperature to models that accept a custom value. - if not is_openai_reasoning: + if not is_openai: kwargs["temperature"] = config.temperature if config.api_key: @@ -100,7 +100,7 @@ def _build_litellm_kwargs(config: LLMConfig) -> dict: if config.api_base: kwargs["api_base"] = config.api_base - if is_openai_reasoning: + if is_openai: if config.max_tokens: kwargs["max_completion_tokens"] = config.max_tokens elif config.provider == "Anthropic API": diff --git a/core/models.py b/core/models.py index e3420c7..b572790 100644 --- a/core/models.py +++ b/core/models.py @@ -32,7 +32,6 @@ class ProviderInfo: class ModelInfo: model_id: str provider_key: str - uses_max_completion_tokens: bool = False # OpenAI reasoning models (gpt-5.x) supports_thinking: bool = False # Claude/Gemini extended thinking help_text: str = "" @@ -94,30 +93,27 @@ class ModelInfo: # --------------------------------------------------------------------------- MODELS: list[ModelInfo] = [ - # --- OpenAI (all gpt-5.x are reasoning models: max_completion_tokens, and - # they reject any temperature other than the default — see core/llm.py) --- + # --- OpenAI (all current chat models use max_completion_tokens, and the + # gpt-5.x reasoning models reject any temperature other than the + # default — both handled by provider in core/llm.py) --- ModelInfo( model_id="gpt-5.6-sol", provider_key="OpenAI API", - uses_max_completion_tokens=True, help_text="GPT-5.6 Sol is OpenAI's flagship frontier reasoning model for complex work.", ), ModelInfo( model_id="gpt-5.6-terra", provider_key="OpenAI API", - uses_max_completion_tokens=True, help_text="GPT-5.6 Terra balances intelligence and cost for everyday work.", ), ModelInfo( model_id="gpt-5.6-luna", provider_key="OpenAI API", - uses_max_completion_tokens=True, help_text="GPT-5.6 Luna is the fast, cost-efficient option for high-volume workloads.", ), ModelInfo( model_id="gpt-5.5", provider_key="OpenAI API", - uses_max_completion_tokens=True, help_text="GPT-5.5 is OpenAI's previous-generation flagship model.", ), # --- Anthropic --- @@ -230,8 +226,3 @@ def get_provider(provider_key: str) -> ProviderInfo | None: def get_litellm_prefix(provider_key: str) -> str: provider = PROVIDERS.get(provider_key) return provider.litellm_prefix if provider else "" - - -def model_uses_completion_tokens(provider_key: str, model_id: str) -> bool: - model = get_model(provider_key, model_id) - return bool(model and model.uses_max_completion_tokens) diff --git a/tests/test_llm.py b/tests/test_llm.py index 4ec0b46..0e65893 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -18,7 +18,9 @@ def test_kwargs_always_set_num_retries_to_three() -> None: assert kwargs["num_retries"] == 3 -def test_openai_reasoning_model_uses_max_completion_tokens() -> None: +def test_openai_provider_uses_max_completion_tokens() -> None: + """Routing is keyed off the OpenAI provider, not a per-model flag: every + current OpenAI chat model takes max_completion_tokens, not max_tokens.""" config = LLMConfig( provider="OpenAI API", model_name="gpt-5.5", @@ -32,16 +34,17 @@ def test_openai_reasoning_model_uses_max_completion_tokens() -> None: assert kwargs["model"] == "gpt-5.5" -def test_openai_reasoning_model_without_max_tokens_omits_completion_tokens() -> None: +def test_openai_provider_without_max_tokens_omits_completion_tokens() -> None: config = LLMConfig(provider="OpenAI API", model_name="gpt-5.5", api_key="k") kwargs = _build_litellm_kwargs(config) assert "max_completion_tokens" not in kwargs assert "max_tokens" not in kwargs -def test_openai_reasoning_model_omits_temperature() -> None: +def test_openai_provider_omits_temperature() -> None: """gpt-5.x reject any temperature but the default (1), so we must not send - one — otherwise litellm raises BadRequestError (regression: v0.13).""" + one — otherwise litellm raises BadRequestError (regression: v0.13). The + whole OpenAI family is the reasoning family, so we gate on the provider.""" config = LLMConfig( provider="OpenAI API", model_name="gpt-5.5", api_key="k", temperature=0.7 ) diff --git a/tests/test_models.py b/tests/test_models.py index 0d56bd0..8ac2930 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,7 +1,7 @@ """Tests for the provider/model registry in `core.models`. Scope: only behaviour that's specific to this project's registry decisions -(fallback shapes, the OpenAI gpt-5 invariant). Tautological lookups that just +(fallback shapes, provider references). Tautological lookups that just re-derive expected values from MODELS/PROVIDERS are intentionally omitted. """ @@ -14,7 +14,6 @@ get_model, get_models_for_provider, get_provider, - model_uses_completion_tokens, ) @@ -52,17 +51,3 @@ def test_get_models_for_provider_unknown_returns_empty() -> None: def test_get_model_unknown_returns_none() -> None: assert get_model("OpenAI API", "no-such-model") is None assert get_model("Nope", "gpt-5.5") is None - - -def test_model_uses_completion_tokens_true_only_for_openai_gpt5() -> None: - """Today only OpenAI gpt-5.x entries should set the flag. Adding a new - model that flips this requires updating the assertion, which is the point — - it's an intentional choice, not an accident.""" - for m in MODELS: - if m.uses_max_completion_tokens: - assert m.provider_key == "OpenAI API" - assert m.model_id.startswith("gpt-5.") - - -def test_model_uses_completion_tokens_unknown_model_is_false() -> None: - assert model_uses_completion_tokens("OpenAI API", "no-such-model") is False