diff --git a/.env.example b/.env.example index a81234734..b6f14ed41 100644 --- a/.env.example +++ b/.env.example @@ -61,6 +61,10 @@ NVIDIA_API_KEY= # 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. #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 # default). Setting one also skips the matching interactive prompt. #TRADINGAGENTS_OPENAI_REASONING_EFFORT=medium diff --git a/tests/test_llm_max_tokens.py b/tests/test_llm_max_tokens.py new file mode 100644 index 000000000..87cddba52 --- /dev/null +++ b/tests/test_llm_max_tokens.py @@ -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 diff --git a/tradingagents/default_config.py b/tradingagents/default_config.py index ab75f7508..8487c0d74 100644 --- a/tradingagents/default_config.py +++ b/tradingagents/default_config.py @@ -19,6 +19,7 @@ _ENV_OVERRIDES = { "TRADINGAGENTS_BENCHMARK_TICKER": "benchmark_ticker", "TRADINGAGENTS_TEMPERATURE": "temperature", "TRADINGAGENTS_LLM_MAX_RETRIES": "llm_max_retries", + "TRADINGAGENTS_MAX_TOKENS": "max_tokens", # Provider-specific reasoning/thinking knobs (None = each provider's own # default). Settable here for non-interactive runs; the CLI also offers an # 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 # 429 throttling on rate-limited deployments instead of aborting a run (#1091). "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 # so a crashed run can resume from the last successful step. "checkpoint_enabled": False, diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index a1902177d..9ebdcf51e 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -62,6 +62,19 @@ def _coerce_max_retries(value): 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: """Main class that orchestrates the trading agents framework.""" @@ -183,6 +196,13 @@ class TradingAgentsGraph: if max_retries is not None and 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 def _create_tool_nodes(self) -> dict[str, ToolNode]: diff --git a/tradingagents/llm_clients/google_client.py b/tradingagents/llm_clients/google_client.py index 93bb1d11a..19ae8a213 100644 --- a/tradingagents/llm_clients/google_client.py +++ b/tradingagents/llm_clients/google_client.py @@ -31,7 +31,8 @@ class GoogleClient(BaseLLMClient): if 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: llm_kwargs[key] = self.kwargs[key] diff --git a/tradingagents/llm_clients/openai_client.py b/tradingagents/llm_clients/openai_client.py index bf6e55457..5f78eb6e6 100644 --- a/tradingagents/llm_clients/openai_client.py +++ b/tradingagents/llm_clients/openai_client.py @@ -164,7 +164,7 @@ class MinimaxChatOpenAI(NormalizedChatOpenAI): # Kwargs forwarded from user config to ChatOpenAI _PASSTHROUGH_KWARGS = ( - "timeout", "max_retries", "reasoning_effort", "temperature", + "timeout", "max_retries", "reasoning_effort", "temperature", "max_tokens", "api_key", "callbacks", "http_client", "http_async_client", )