refactor: route max_completion_tokens on provider, drop per-model flag (#30)

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.<br><br>- 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.<br><br>- 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.<br><br>- 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? |

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":

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)

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
     )

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