fix: cap LLM output at 16k tokens

Author: Henry

backend_api_python/app/services/llm.py

@@ -16,6 +16,8 @@
 
 logger = get_logger(__name__)
 
+DEFAULT_MAX_TOKENS = 16_384
+
 
 class LLMAPIError(ValueError):
     """Provider HTTP error with status and request metadata preserved."""
@@ -198,6 +200,20 @@ def get_code_generation_model(self, provider: LLMProvider = None) -> str:
             return model
         return self.get_default_model(provider)
 
+    def get_max_tokens(self) -> int:
+        """Get the shared maximum output-token budget for every LLM provider."""
+        config = load_addon_config()
+        configured = config.get('llm', {}).get('max_tokens', DEFAULT_MAX_TOKENS)
+        try:
+            return max(1, int(configured))
+        except (TypeError, ValueError):
+            logger.warning(
+                "Invalid LLM_MAX_TOKENS=%r; using %s",
+                configured,
+                DEFAULT_MAX_TOKENS,
+            )
+            return DEFAULT_MAX_TOKENS
+
     def is_configured(self, provider: LLMProvider = None) -> bool:
         """Return whether the provider has enough configuration to make a request."""
         p = provider or self.provider
@@ -301,6 +317,7 @@ def _call_openai_compatible(self, messages: list, model: str, temperature: float
             "model": model,
             "messages": messages,
             "temperature": temperature,
+            "max_tokens": self.get_max_tokens(),
         }
 
         # AtlasCloud documents the OpenAI-compatible ChatCompletion shape, but
@@ -458,6 +475,7 @@ def _call_google_gemini(self, messages: list, model: str, temperature: float,
             "contents": contents,
             "generationConfig": {
                 "temperature": temperature,
+                "maxOutputTokens": self.get_max_tokens(),
                 "responseMimeType": "application/json",
             }
         }
@@ -496,6 +514,7 @@ def _call_litellm(self, messages: list, model: str, temperature: float,
             "model": model,
             "messages": messages,
             "temperature": temperature,
+            "max_tokens": self.get_max_tokens(),
             "timeout": timeout,
             "drop_params": True,
         }
@@ -535,6 +554,7 @@ def _stream_openai_compatible(self, messages: list, model: str, temperature: flo
             "model": model,
             "messages": messages,
             "temperature": temperature,
+            "max_tokens": self.get_max_tokens(),
             "stream": True,
         }
         response = self._llm_post(url, headers=headers, json_payload=data, timeout=timeout, stream=True)

backend_api_python/app/utils/config_loader.py

@@ -92,7 +92,7 @@ def env_get(name: str) -> Optional[str]:
         ('OPENROUTER_API_URL', 'openrouter.api_url', 'string'),
         ('OPENROUTER_MODEL', 'openrouter.model', 'string'),
         ('OPENROUTER_TEMPERATURE', 'openrouter.temperature', 'float'),
-        ('OPENROUTER_MAX_TOKENS', 'openrouter.max_tokens', 'int'),
+        ('LLM_MAX_TOKENS', 'llm.max_tokens', 'int'),
         ('OPENROUTER_TIMEOUT', 'openrouter.timeout', 'int'),
         ('OPENROUTER_CONNECT_TIMEOUT', 'openrouter.connect_timeout', 'int'),
 

backend_api_python/env.example

@@ -487,7 +487,9 @@ SIGNAL_NOTIFY_TIMEOUT_SEC=6
 # LLM advanced tuning
 OPENROUTER_API_URL=https://openrouter.ai/api/v1/chat/completions
 OPENROUTER_TEMPERATURE=0.7
-OPENROUTER_MAX_TOKENS=4000
+# Shared maximum output-token budget for all LLM providers.
+# Providers charge for tokens actually generated, not this full upper bound.
+LLM_MAX_TOKENS=16384
 OPENROUTER_TIMEOUT=300
 OPENROUTER_CONNECT_TIMEOUT=30
 OPENAI_BASE_URL=https://api.openai.com/v1

backend_api_python/tests/test_llm_litellm_provider.py

@@ -26,6 +26,16 @@ def test_litellm_env_mapping(monkeypatch):
     assert cfg["litellm"]["base_url"] == "https://litellm.example/v1"
 
 
+def test_uniform_llm_max_tokens_env_mapping(monkeypatch):
+    monkeypatch.setenv("LLM_MAX_TOKENS", "16384")
+    _reset_config_cache()
+
+    cfg = load_addon_config()
+
+    assert cfg["llm"]["max_tokens"] == 16384
+    assert LLMService().get_max_tokens() == 16384
+
+
 def test_atlascloud_env_mapping(monkeypatch):
     monkeypatch.setenv("ATLASCLOUD_API_KEY", "atlas-key")
     monkeypatch.setenv("ATLASCLOUD_MODEL", "openai/gpt-5.4")
@@ -90,9 +100,82 @@ def fake_post(url, headers, json, timeout):
     assert captured["url"] == "https://api.atlascloud.ai/v1/chat/completions"
     assert captured["headers"]["Authorization"] == "Bearer atlas-key"
     assert captured["json"]["model"] == "deepseek-v3"
+    assert captured["json"]["max_tokens"] == 16384
     assert "response_format" not in captured["json"]
 
 
+def test_google_gemini_uses_uniform_max_tokens(monkeypatch):
+    captured = {}
+
+    class FakeResponse:
+        def raise_for_status(self):
+            return None
+
+        def json(self):
+            return {
+                "candidates": [
+                    {"content": {"parts": [{"text": "{\"ok\": true}"}]}}
+                ]
+            }
+
+    service = LLMService(provider="google")
+    monkeypatch.setattr(service, "get_max_tokens", lambda: 16384)
+
+    def fake_post(url, **kwargs):
+        captured.update({"url": url, **kwargs})
+        return FakeResponse()
+
+    monkeypatch.setattr(service, "_llm_post", fake_post)
+
+    out = service._call_google_gemini(
+        [{"role": "user", "content": "hello"}],
+        "gemini-1.5-flash",
+        0.7,
+        "google-key",
+        "https://generativelanguage.googleapis.com/v1beta",
+        30,
+    )
+
+    assert out == "{\"ok\": true}"
+    assert captured["json_payload"]["generationConfig"]["maxOutputTokens"] == 16384
+
+
+def test_openai_compatible_stream_uses_uniform_max_tokens(monkeypatch):
+    captured = {}
+
+    class FakeResponse:
+        status_code = 200
+
+        def iter_lines(self, decode_unicode=False):
+            return [b"data: [DONE]"]
+
+        def close(self):
+            return None
+
+    service = LLMService(provider="openrouter")
+    monkeypatch.setattr(service, "get_max_tokens", lambda: 16384)
+
+    def fake_post(url, **kwargs):
+        captured.update({"url": url, **kwargs})
+        return FakeResponse()
+
+    monkeypatch.setattr(service, "_llm_post", fake_post)
+
+    chunks = list(
+        service._stream_openai_compatible(
+            [{"role": "user", "content": "hello"}],
+            "openai/gpt-5.4",
+            0.7,
+            "openrouter-key",
+            "https://openrouter.ai/api/v1",
+            30,
+        )
+    )
+
+    assert chunks == []
+    assert captured["json_payload"]["max_tokens"] == 16384
+
+
 @pytest.mark.parametrize(
     ("payload", "expected"),
     [
@@ -400,9 +483,12 @@ def completion(**kwargs):
 
 
 def test_litellm_response_content(monkeypatch):
+    captured = {}
+
     class FakeLiteLLM:
         @staticmethod
         def completion(**kwargs):
+            captured.update(kwargs)
             return SimpleNamespace(
                 choices=[SimpleNamespace(message=SimpleNamespace(content="hello"))]
             )
@@ -421,3 +507,4 @@ def completion(**kwargs):
     )
 
     assert out == "hello"
+    assert captured["max_tokens"] == 16384