fix(google): route Gemma 4 thinking through thinking_level

livekit-plugins/livekit-plugins-google/livekit/plugins/google/llm.py

@@ -43,7 +43,17 @@
 
 def _is_gemini_3_model(model: str) -> bool:
     """Check if model is Gemini 3 series"""
-    return "gemini-3" in model.lower() or model.lower().startswith("gemini-3")
+    return "gemini-3" in model.lower()
+
+
+def _is_gemma_4_model(model: str) -> bool:
+    """Check if model is Gemma 4 series"""
+    return "gemma-4-" in model.lower()
+
+
+def _supports_thinking_level(model: str) -> bool:
+    """Check if model configures thinking with a level instead of a token budget"""
+    return _is_gemini_3_model(model) or _is_gemma_4_model(model)
 
 
 def _function_calling_config(
@@ -70,7 +80,7 @@ def _function_calling_config(
 def _is_gemini_3_flash_model(model: str) -> bool:
     """Check if model is Gemini 3 Flash"""
     m = model.lower()
-    return m.startswith("gemini-3") and "flash" in m
+    return "gemini-3" in m and "flash" in m
 
 
 def _requires_thought_signatures(model: str) -> bool:
@@ -228,12 +238,26 @@ def __init__(
                 _thinking_level = thinking_config.get("thinking_level")
             elif isinstance(thinking_config, types.ThinkingConfig):
                 _thinking_budget = thinking_config.thinking_budget
-                _thinking_level = getattr(thinking_config, "thinking_level", None)
+                _thinking_level = thinking_config.thinking_level
 
             if _thinking_budget is not None:
                 if not isinstance(_thinking_budget, int):
                     raise ValueError("thinking_budget inside thinking_config must be an integer")
 
+            # gemma 4 toggles thinking on and off with nothing in between, and answers any other
+            # level with 400 INVALID_ARGUMENT
+            # https://ai.google.dev/gemma/docs/core/gemma_on_gemini_api
+            if (
+                _thinking_level is not None
+                and _is_gemma_4_model(model)
+                and _thinking_level.upper()
+                not in (types.ThinkingLevel.MINIMAL, types.ThinkingLevel.HIGH)
+            ):
+                raise ValueError(
+                    f"Model {model} only supports thinking_level 'minimal' (thinking off) "
+                    f"or 'high' (thinking on), got {_thinking_level!r}."
+                )
+
         self._opts = _LLMOptions(
             model=model,
             temperature=temperature,
@@ -364,35 +388,35 @@ def chat(
 
         # Handle thinking_config based on model version
         if is_given(self._opts.thinking_config):
-            is_gemini_3 = _is_gemini_3_model(self._opts.model)
-            is_gemini_3_flash = _is_gemini_3_flash_model(self._opts.model)
             thinking_cfg = self._opts.thinking_config
 
-            # Extract both parameters
+            # Extract the parameters
             _budget = None
             _level: str | types.ThinkingLevel | None = None
+            _include_thoughts: bool | None = None
             if isinstance(thinking_cfg, dict):
                 _budget = thinking_cfg.get("thinking_budget")
                 _level = thinking_cfg.get("thinking_level")
+                _include_thoughts = thinking_cfg.get("include_thoughts")
             elif isinstance(thinking_cfg, types.ThinkingConfig):
                 _budget = thinking_cfg.thinking_budget
-                _level = getattr(thinking_cfg, "thinking_level", None)
+                _level = thinking_cfg.thinking_level
+                _include_thoughts = thinking_cfg.include_thoughts
 
-            if is_gemini_3:
-                # Gemini 3: only support thinking_level
+            if _supports_thinking_level(self._opts.model):
+                # Gemini 3 and Gemma 4 configure thinking with a level, not a token budget
                 if _budget is not None and _level is None:
                     logger.warning(
-                        f"Model {self._opts.model} is Gemini 3 which does not support thinking_budget. "
-                        "Please use thinking_level ('low' or 'high') instead. Ignoring thinking_budget."
+                        f"Model {self._opts.model} does not support thinking_budget. "
+                        "Please use thinking_level instead. Ignoring thinking_budget."
                     )
-                if _level is None:
-                    # If no thinking_level is provided, use the fastest thinking level
-                    if is_gemini_3_flash:
-                        _level = "minimal"
-                    else:
-                        _level = "low"
-                # Use thinking_level only (pass as dict since SDK may not have this field yet)
-                extra["thinking_config"] = {"thinking_level": _level}
+                if _level is None and _is_gemini_3_model(self._opts.model):
+                    # only default the level where the accepted levels are known
+                    _level = "minimal" if _is_gemini_3_flash_model(self._opts.model) else "low"
+                # the level passes through as given, checked against the model in __init__
+                extra["thinking_config"] = types.ThinkingConfig(
+                    thinking_level=_level, include_thoughts=_include_thoughts
+                )
 
             else:
                 # Gemini 2.5 and earlier: only support thinking_budget

tests/test_google_thought_signatures.py

@@ -4,6 +4,7 @@
     _is_gemini_3_flash_model,
     _is_gemini_3_model,
     _requires_thought_signatures,
+    _supports_thinking_level,
 )
 
 pytestmark = pytest.mark.unit
@@ -29,6 +30,7 @@ class TestGeminiModelDetection:
             # Gemini 1.5 models - should return False
             ("gemini-1.5-pro", False),
             # Other models - should return False
+            ("gemma-4-31b-it", False),
             ("gpt-4", False),
             ("claude-3", False),
         ],
@@ -43,6 +45,7 @@ def test_is_gemini_3_model(self, model: str, expected: bool):
             ("gemini-3-flash-preview", True),
             ("gemini-3-flash", True),
             ("GEMINI-3-FLASH", True),  # case insensitive
+            ("models/gemini-3-flash-preview", True),  # qualified name
             # Gemini 3 Pro models - should return False
             ("gemini-3-pro-preview", False),
             ("gemini-3-pro", False),
@@ -54,6 +57,24 @@ def test_is_gemini_3_model(self, model: str, expected: bool):
     def test_is_gemini_3_flash_model(self, model: str, expected: bool):
         assert _is_gemini_3_flash_model(model) == expected
 
+    @pytest.mark.parametrize(
+        "model,expected",
+        [
+            # level models - should return True
+            ("gemini-3-pro-preview", True),
+            ("gemma-4-31b-it", True),
+            ("GEMMA-4-31B-IT", True),  # case insensitive
+            ("models/gemma-4-31b-it", True),  # qualified name
+            ("publishers/google/models/gemma-4-26b-a4b-it", True),
+            # budget models - should return False
+            ("gemini-2.5-flash", False),
+            ("gemma-3-27b-it", False),
+            ("gemma-40-31b-it", False),
+        ],
+    )
+    def test_supports_thinking_level(self, model: str, expected: bool):
+        assert _supports_thinking_level(model) == expected
+
     @pytest.mark.parametrize(
         "model,expected",
         [
@@ -76,6 +97,7 @@ def test_is_gemini_3_flash_model(self, model: str, expected: bool):
             # Gemini 1.5 models - should return False
             ("gemini-1.5-pro", False),
             # Other models - should return False
+            ("gemma-4-31b-it", False),
             ("gpt-4", False),
             ("claude-3", False),
         ],

tests/test_plugin_google_llm.py

@@ -506,3 +506,120 @@ async def test_cached_content_strips_raw_tools_from_extra_kwargs(self) -> None:
         assert config.cached_content == "cachedContents/abc"
         assert config.tools is None
         assert config.tool_config is None
+
+
+class TestThinkingConfigRequestConstruction:
+    """Gemini 3 and Gemma 4 configure thinking with ``thinking_level``, every other model with
+    ``thinking_budget``. These tests drive ``chat()`` against a stubbed
+    ``generate_content_stream`` and assert on the ``ThinkingConfig`` it received."""
+
+    @staticmethod
+    async def _single_response_async_iter():
+        yield types.GenerateContentResponse(
+            candidates=[
+                types.Candidate(
+                    content=types.Content(role="model", parts=[types.Part(text="ok")]),
+                    finish_reason=types.FinishReason.STOP,
+                )
+            ],
+        )
+
+    @classmethod
+    async def _capture_thinking_config(
+        cls, model: str, thinking_config: types.ThinkingConfigOrDict
+    ) -> types.ThinkingConfig | None:
+        captured: dict = {}
+
+        async def fake_stream(**kwargs):
+            captured["config"] = kwargs.get("config")
+            return cls._single_response_async_iter()
+
+        llm_ = LLM(model=model, api_key="test", thinking_config=thinking_config)
+        fake = AsyncMock(side_effect=fake_stream)
+        with patch.object(llm_._client.aio.models, "generate_content_stream", fake):
+            stream = llm_.chat(chat_ctx=ChatContext.empty())
+            try:
+                async for _ in stream:
+                    pass
+            finally:
+                await stream.aclose()
+        return captured["config"].thinking_config
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("model", ["gemma-4-31b-it", "models/gemma-4-31b-it"])
+    @pytest.mark.parametrize("level", [types.ThinkingLevel.MINIMAL, types.ThinkingLevel.HIGH])
+    async def test_gemma_4_level_reaches_the_request(
+        self, model: str, level: types.ThinkingLevel
+    ) -> None:
+        """Gemma 4 takes the level it documents, under a bare or a qualified model name."""
+        thinking_config = await self._capture_thinking_config(
+            model, types.ThinkingConfig(thinking_level=level)
+        )
+
+        assert thinking_config is not None
+        assert thinking_config.thinking_level == level
+
+    @pytest.mark.parametrize("level", [types.ThinkingLevel.LOW, types.ThinkingLevel.MEDIUM])
+    def test_gemma_4_rejects_undocumented_levels(self, level: types.ThinkingLevel) -> None:
+        """Gemma 4 answers 'low' and 'medium' with 400 INVALID_ARGUMENT, so the session fails at
+        construction instead of on the first turn."""
+        with pytest.raises(ValueError, match="only supports thinking_level"):
+            LLM(
+                model="gemma-4-31b-it",
+                api_key="test",
+                thinking_config=types.ThinkingConfig(thinking_level=level),
+            )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("model", ["gemini-3-pro-preview", "gemma-4-31b-it"])
+    async def test_include_thoughts_survives_the_level_branch(self, model: str) -> None:
+        thinking_config = await self._capture_thinking_config(
+            model, {"thinking_level": "high", "include_thoughts": True}
+        )
+
+        assert thinking_config is not None
+        assert thinking_config.thinking_level == types.ThinkingLevel.HIGH
+        assert thinking_config.include_thoughts is True
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "model,expected_level",
+        [
+            ("gemini-3-flash-preview", types.ThinkingLevel.MINIMAL),
+            ("models/gemini-3-flash-preview", types.ThinkingLevel.MINIMAL),
+            ("gemini-3-pro-preview", types.ThinkingLevel.LOW),
+            # no default is invented for a model whose accepted levels are unknown
+            ("gemma-4-31b-it", None),
+        ],
+    )
+    async def test_default_level_only_for_gemini_3(
+        self, model: str, expected_level: types.ThinkingLevel | None
+    ) -> None:
+        thinking_config = await self._capture_thinking_config(
+            model, types.ThinkingConfig(include_thoughts=False)
+        )
+
+        assert thinking_config is not None
+        assert thinking_config.thinking_level == expected_level
+        assert thinking_config.include_thoughts is False
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("model", ["gemini-3-pro-preview", "gemma-4-31b-it"])
+    async def test_budget_alone_is_dropped_on_level_models(self, model: str) -> None:
+        """A budget on a level model is ignored with a warning, on Gemini 3 and Gemma 4 alike."""
+        thinking_config = await self._capture_thinking_config(
+            model, types.ThinkingConfig(thinking_budget=0)
+        )
+
+        assert thinking_config is not None
+        assert thinking_config.thinking_budget is None
+
+    @pytest.mark.asyncio
+    async def test_budget_still_reaches_gemini_2_5(self) -> None:
+        thinking_config = await self._capture_thinking_config(
+            "gemini-2.5-flash", types.ThinkingConfig(thinking_budget=1024)
+        )
+
+        assert thinking_config is not None
+        assert thinking_config.thinking_budget == 1024
+        assert thinking_config.thinking_level is None