feat(llm): add a configurable output-token cap

- some model/gateway combinations emit unbounded reasoning/output and hang or
  trip an idle timeout (e.g. some deepseek-v4-flash deployments)
- add an opt-in max_tokens config knob + TRADINGAGENTS_MAX_TOKENS, forwarded to
  every provider when set (Gemini takes it as max_output_tokens); int-coerced,
  rejects non-positive/boolean values #1204
This commit is contained in:
Yijia-Xiao
2026-08-30 06:26:57 +00:00
parent 539eae8fd6
commit 0ef56e6a33
6 changed files with 155 additions and 2 deletions

View File

@@ -61,6 +61,10 @@ NVIDIA_API_KEY=
# own default (usually 2). Raise it to ride out bursty 429 rate-limit throttling # own default (usually 2). Raise it to ride out bursty 429 rate-limit throttling
# on rate-limited deployments (e.g. Azure OpenAI) instead of aborting the run. # on rate-limited deployments (e.g. Azure OpenAI) instead of aborting the run.
#TRADINGAGENTS_LLM_MAX_RETRIES=6 #TRADINGAGENTS_LLM_MAX_RETRIES=6
# Cap on output tokens forwarded to every provider (Gemini's max_output_tokens
# too). Unset leaves each provider at its default. Set it to bound a model that
# emits unbounded reasoning/output and hangs or trips a gateway idle timeout.
#TRADINGAGENTS_MAX_TOKENS=8192
# Provider-specific reasoning/thinking depth (optional; unset = provider # Provider-specific reasoning/thinking depth (optional; unset = provider
# default). Setting one also skips the matching interactive prompt. # default). Setting one also skips the matching interactive prompt.
#TRADINGAGENTS_OPENAI_REASONING_EFFORT=medium #TRADINGAGENTS_OPENAI_REASONING_EFFORT=medium

View File

@@ -0,0 +1,122 @@
"""Configurable output-token cap (#1204).
Some model/gateway combinations (e.g. deepseek-v4-flash deployments) emit
unbounded reasoning/output and hang or trip an idle timeout. An opt-in
``max_tokens`` config knob is forwarded to every provider so a run can bound it;
Gemini names the parameter ``max_output_tokens``, so it is forwarded under the
right key per provider.
"""
from __future__ import annotations
import importlib
import pytest
import tradingagents.default_config as default_config_module
from tradingagents.graph.trading_graph import TradingAgentsGraph, _coerce_max_tokens
# --- coercion / validation -------------------------------------------------
@pytest.mark.unit
@pytest.mark.parametrize("value,expected", [(1, 1), (8192, 8192), ("4096", 4096)])
def test_coerce_accepts_positive_ints_and_numeric_strings(value, expected):
assert _coerce_max_tokens(value) == expected
@pytest.mark.unit
@pytest.mark.parametrize("bad", [0, -1, "0", "-5"])
def test_coerce_rejects_non_positive(bad):
with pytest.raises(ValueError, match="> 0"):
_coerce_max_tokens(bad)
@pytest.mark.unit
@pytest.mark.parametrize("bad", [True, False])
def test_coerce_rejects_booleans(bad):
with pytest.raises(ValueError, match="boolean"):
_coerce_max_tokens(bad)
@pytest.mark.unit
@pytest.mark.parametrize("bad", ["abc", "1.5", None])
def test_coerce_rejects_non_integers(bad):
with pytest.raises(ValueError, match="integer"):
_coerce_max_tokens(bad)
# --- forwarding into provider kwargs (right key per provider) --------------
def _bare_graph(config):
g = object.__new__(TradingAgentsGraph)
g.config = config
return g
@pytest.mark.unit
def test_not_forwarded_when_unset():
kwargs = _bare_graph({"llm_provider": "openai", "max_tokens": None})._get_provider_kwargs()
assert "max_tokens" not in kwargs
assert "max_output_tokens" not in kwargs
@pytest.mark.unit
@pytest.mark.parametrize("provider", ["openai", "anthropic", "deepseek", "openai_compatible"])
def test_forwarded_as_max_tokens_for_non_google(provider):
kwargs = _bare_graph({"llm_provider": provider, "max_tokens": 8192})._get_provider_kwargs()
assert kwargs["max_tokens"] == 8192
assert "max_output_tokens" not in kwargs
@pytest.mark.unit
def test_forwarded_as_max_output_tokens_for_google():
# Gemini's kwarg name differs; forwarding plain max_tokens would be rejected.
kwargs = _bare_graph({"llm_provider": "google", "max_tokens": 8192})._get_provider_kwargs()
assert kwargs["max_output_tokens"] == 8192
assert "max_tokens" not in kwargs
@pytest.mark.unit
def test_env_string_is_coerced():
kwargs = _bare_graph({"llm_provider": "openai", "max_tokens": "4096"})._get_provider_kwargs()
assert kwargs["max_tokens"] == 4096
@pytest.mark.unit
def test_invalid_value_fails_loudly():
with pytest.raises(ValueError):
_bare_graph({"llm_provider": "openai", "max_tokens": 0})._get_provider_kwargs()
# --- client-side allowlists carry the kwarg --------------------------------
@pytest.mark.unit
def test_openai_and_google_clients_accept_the_kwarg():
from tradingagents.llm_clients import openai_client
from tradingagents.llm_clients.google_client import GoogleClient # noqa: F401
assert "max_tokens" in openai_client._PASSTHROUGH_KWARGS
# Google client forwards max_output_tokens through construction.
llm = GoogleClient("gemini-3.5-flash", api_key="x", max_output_tokens=8192).get_llm()
assert getattr(llm, "max_output_tokens", None) == 8192
# --- env overlay -----------------------------------------------------------
def _reload_with_env(monkeypatch, **overrides):
for key in list(default_config_module._ENV_OVERRIDES):
monkeypatch.delenv(key, raising=False)
for key, val in overrides.items():
monkeypatch.setenv(key, val)
return importlib.reload(default_config_module)
@pytest.mark.unit
def test_default_is_none(monkeypatch):
dc = _reload_with_env(monkeypatch)
assert dc.DEFAULT_CONFIG["max_tokens"] is None
@pytest.mark.unit
def test_env_override_sets_config(monkeypatch):
dc = _reload_with_env(monkeypatch, TRADINGAGENTS_MAX_TOKENS="8192")
assert dc.DEFAULT_CONFIG["max_tokens"] == "8192"
assert _coerce_max_tokens(dc.DEFAULT_CONFIG["max_tokens"]) == 8192

View File

@@ -19,6 +19,7 @@ _ENV_OVERRIDES = {
"TRADINGAGENTS_BENCHMARK_TICKER": "benchmark_ticker", "TRADINGAGENTS_BENCHMARK_TICKER": "benchmark_ticker",
"TRADINGAGENTS_TEMPERATURE": "temperature", "TRADINGAGENTS_TEMPERATURE": "temperature",
"TRADINGAGENTS_LLM_MAX_RETRIES": "llm_max_retries", "TRADINGAGENTS_LLM_MAX_RETRIES": "llm_max_retries",
"TRADINGAGENTS_MAX_TOKENS": "max_tokens",
# Provider-specific reasoning/thinking knobs (None = each provider's own # Provider-specific reasoning/thinking knobs (None = each provider's own
# default). Settable here for non-interactive runs; the CLI also offers an # default). Settable here for non-interactive runs; the CLI also offers an
# interactive choice, which is skipped when the matching var is set. # interactive choice, which is skipped when the matching var is set.
@@ -100,6 +101,11 @@ DEFAULT_CONFIG = _apply_env_overrides({
# provider/SDK at its own default (usually 2). Raise it to ride out bursty # provider/SDK at its own default (usually 2). Raise it to ride out bursty
# 429 throttling on rate-limited deployments instead of aborting a run (#1091). # 429 throttling on rate-limited deployments instead of aborting a run (#1091).
"llm_max_retries": None, "llm_max_retries": None,
# Cap on output tokens forwarded to every provider chat client. None leaves
# each provider at its own default. Set it to bound a model that emits
# unbounded reasoning/output and hangs or trips a gateway idle timeout
# (e.g. some deepseek-v4-flash deployments, #1204).
"max_tokens": None,
# Checkpoint/resume: when True, LangGraph saves state after each node # Checkpoint/resume: when True, LangGraph saves state after each node
# so a crashed run can resume from the last successful step. # so a crashed run can resume from the last successful step.
"checkpoint_enabled": False, "checkpoint_enabled": False,

View File

@@ -62,6 +62,19 @@ def _coerce_max_retries(value):
return n return n
def _coerce_max_tokens(value):
"""Validate a ``max_tokens`` value to a positive int (env vars are strings)."""
if isinstance(value, bool):
raise ValueError(f"max_tokens must be an integer, not a boolean: {value!r}")
try:
n = int(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"max_tokens must be an integer, got {value!r}") from exc
if n <= 0:
raise ValueError(f"max_tokens must be > 0, got {n}")
return n
class TradingAgentsGraph: class TradingAgentsGraph:
"""Main class that orchestrates the trading agents framework.""" """Main class that orchestrates the trading agents framework."""
@@ -183,6 +196,13 @@ class TradingAgentsGraph:
if max_retries is not None and max_retries != "": if max_retries is not None and max_retries != "":
kwargs["max_retries"] = _coerce_max_retries(max_retries) kwargs["max_retries"] = _coerce_max_retries(max_retries)
# Output-token cap is cross-provider, but Gemini names it
# ``max_output_tokens``; forward under the right key when set (#1204).
max_tokens = self.config.get("max_tokens")
if max_tokens is not None and max_tokens != "":
key = "max_output_tokens" if provider == "google" else "max_tokens"
kwargs[key] = _coerce_max_tokens(max_tokens)
return kwargs return kwargs
def _create_tool_nodes(self) -> dict[str, ToolNode]: def _create_tool_nodes(self) -> dict[str, ToolNode]:

View File

@@ -31,7 +31,8 @@ class GoogleClient(BaseLLMClient):
if self.base_url: if self.base_url:
llm_kwargs["base_url"] = self.base_url llm_kwargs["base_url"] = self.base_url
for key in ("timeout", "max_retries", "temperature", "callbacks", "http_client", "http_async_client"): for key in ("timeout", "max_retries", "temperature", "max_output_tokens",
"callbacks", "http_client", "http_async_client"):
if key in self.kwargs: if key in self.kwargs:
llm_kwargs[key] = self.kwargs[key] llm_kwargs[key] = self.kwargs[key]

View File

@@ -164,7 +164,7 @@ class MinimaxChatOpenAI(NormalizedChatOpenAI):
# Kwargs forwarded from user config to ChatOpenAI # Kwargs forwarded from user config to ChatOpenAI
_PASSTHROUGH_KWARGS = ( _PASSTHROUGH_KWARGS = (
"timeout", "max_retries", "reasoning_effort", "temperature", "timeout", "max_retries", "reasoning_effort", "temperature", "max_tokens",
"api_key", "callbacks", "http_client", "http_async_client", "api_key", "callbacks", "http_client", "http_async_client",
) )