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