fix(dspy): support Gemini 3.5 thinking levels, regional prefixes, and streaming usage

dspy/_vendor/lm15/providers/gemini.py

@@ -5,6 +5,7 @@
 import base64
 import json
 import os
+import re
 import struct
 import urllib.parse
 from dataclasses import dataclass, field, replace
@@ -120,15 +121,23 @@
     "codeExecutionResult",
 }
 
+def _normalize_gemini_model(model: str) -> str:
+    return re.sub(
+        r"^[a-z]{2,}\.(?=(?:gemini|gemma)-)",
+        "",
+        model.lower().rsplit("/", 1)[-1],
+    )
+
+
 def gemini_level_class(model: str) -> bool:
-    """True for the Gemini 3.x class: `thinkingLevel`, no full off (MAP-7 rule 10).
+    """True for the Gemini 3.x+ class: `thinkingLevel`, no full off (MAP-7 rule 10).
 
     Model-name table, receipted 2026-09-02: 2.5 models reject thinkingLevel,
-    3.x models take it; 3.7 Flash accepts thinkingBudget: 0 and still spends
+    3.x+ models take it; 3.7 Flash accepts thinkingBudget: 0 and still spends
     thinking tokens.  A table that rots; `extensions` overrides.
     """
-    lowered = model.lower()
-    return lowered.startswith("models/gemini-3") or lowered.startswith("gemini-3")
+    normalized = _normalize_gemini_model(model)
+    return bool(re.match(r"^gemini-([3-9]|\d{2,})(?:\.\d+|-|$)", normalized))
 
 
 
@@ -772,19 +781,25 @@ def _payload(self, request: Request) -> dict[str, Any]:
             generation_config.update(_response_format_to_gemini_config(fmt))
         if request.config.reasoning is not None:
             reasoning = request.config.reasoning
+            normalized_model = _normalize_gemini_model(request.model)
+            is_pro = "-pro" in normalized_model
+            was_off = reasoning.is_off
             level_class = gemini_level_class(request.model)
             if reasoning.is_off and level_class:
                 # MAP-13 (decision 2026-09-14 §4.2): the Gemini 3 class has no
                 # honoured off switch (thinkingBudget 0 accepted, 58 tokens
                 # still spent on 3.7 Flash, live 2026-09-02); the closest to
                 # "none" is the lowest level, and the spend shows in
-                # usage.reasoning_tokens.
+                # usage.reasoning_tokens. Pro models do not support "minimal".
+                min_level = "low" if is_pro else "minimal"
                 adapt("config.reasoning.effort", "substituted",
                       f"{request.model} cannot disable thinking (the Gemini 3 class honours no off switch); "
                       "the lowest level was sent and the thinking spend is visible in usage",
-                      asked="off", applied="minimal", provider=self.provider)
-                reasoning = replace(reasoning, effort="minimal")
-            if reasoning.is_off:
+                      asked="off", applied=min_level, provider=self.provider)
+                reasoning = replace(reasoning, effort=min_level)
+            if reasoning.is_off and is_pro:
+                generation_config["thinkingConfig"] = {"thinkingBudget": 128}
+            elif reasoning.is_off:
                 generation_config["thinkingConfig"] = {"thinkingBudget": 0}
             else:
                 if reasoning.summary in ("concise", "detailed"):
@@ -793,13 +808,23 @@ def _payload(self, request: Request) -> dict[str, Any]:
                           asked=reasoning.summary, applied="auto", provider=self.provider)
                     reasoning = replace(reasoning, summary="auto")
                 thinking: dict[str, Any] = {}
-                if reasoning.summary is not None:
-                    thinking["includeThoughts"] = True  # MAP-7 rule 7: only when asked
+                if not was_off and reasoning.summary != "none":
+                    thinking["includeThoughts"] = True
                 if reasoning.thinking_budget is not None:
                     thinking["thinkingBudget"] = reasoning.thinking_budget  # 3.x: accepted, docs warn
                 elif level_class:
                     effort = reasoning.effort
-                    if effort in ("xhigh", "max"):
+                    if effort == "minimal" and is_pro:
+                        adapt("config.reasoning.effort", "clamped",
+                              f"{request.model} does not support thinkingLevel 'minimal'; clamped to 'low'",
+                              asked=effort, applied="low", provider=self.provider)
+                        effort = "low"
+                    elif effort == "medium" and normalized_model in ("gemini-3-pro", "gemini-3-pro-preview"):
+                        adapt("config.reasoning.effort", "clamped",
+                              f"{request.model} does not support thinkingLevel 'medium'; clamped to 'high'",
+                              asked=effort, applied="high", provider=self.provider)
+                        effort = "high"
+                    elif effort in ("xhigh", "max"):
                         adapt("config.reasoning.effort", "clamped",
                               "the Gemini 3 class has thinkingLevel minimal|low|medium|high; 'high' is the ceiling",
                               asked=effort, applied="high", provider=self.provider)
@@ -1078,10 +1103,11 @@ def parse_stream_events(self, request: Request, raw_event: SSEEvent) -> Iterator
                         yield StreamDeltaEvent(delta=ImageDelta(data=data, part_index=idx, media_type=mime))
             finish = candidate.get("finishReason")
 
+        usage = self._usage_from_payload(payload) if "usageMetadata" in payload else None
         if finish:
-            yield StreamEndEvent(finish_reason=_finish_reason(finish, has_tool_call=saw_tool), usage=self._usage_from_payload(payload), provider_data=payload)
-        elif not yielded_delta and "usageMetadata" in payload:
-            yield StreamEndEvent(finish_reason="stop", usage=self._usage_from_payload(payload), provider_data=payload)
+            yield StreamEndEvent(finish_reason=_finish_reason(finish, has_tool_call=saw_tool), usage=usage, provider_data=payload)
+        elif not yielded_delta and usage is not None:
+            yield StreamEndEvent(finish_reason=None, usage=usage, provider_data=payload)
 
     def _usage_from_payload(self, payload: dict[str, Any]) -> Usage:
         return _gemini_usage(payload.get("usageMetadata"), output_keys=("candidatesTokenCount", "responseTokenCount"))
@@ -1142,9 +1168,11 @@ def _stream_via_live_completion(self, request: Request) -> Iterator[StreamEvent]
                 raw = ws.recv()
                 events, turn_complete, usage = self._decode_live_completion_stream_events(raw)
                 acc_usage = Usage(
-                    input_tokens=max(acc_usage.input_tokens, usage.input_tokens),
-                    output_tokens=max(acc_usage.output_tokens, usage.output_tokens),
+                    input_tokens=max(acc_usage.input_tokens or 0, usage.input_tokens or 0),
+                    output_tokens=max(acc_usage.output_tokens or 0, usage.output_tokens or 0),
                     total_tokens=max(acc_usage.total_tokens or 0, usage.total_tokens or 0),
+                    cache_read_tokens=max(acc_usage.cache_read_tokens or 0, usage.cache_read_tokens or 0) or None,
+                    reasoning_tokens=max(acc_usage.reasoning_tokens or 0, usage.reasoning_tokens or 0) or None,
                 )
                 for event in events:
                     if event.type == "delta" and isinstance(event.delta, ToolCallDelta):

tests/adapters/test_reasoning.py

@@ -127,3 +127,28 @@ class Sig(dspy.Signature):
 
     assert value["answer"] == "42"
     assert value["reasoning"] == dspy.Reasoning(content="(omitted)")
+
+
+def test_gemini_native_reasoning_sets_include_thoughts_and_thinking_level():
+    from dataclasses import replace
+
+    from dspy._vendor.lm15.providers.gemini import GeminiLM
+    from dspy.clients.execution import _canonical, prepare
+
+    class Sig(dspy.Signature):
+        question: str = dspy.InputField()
+        reasoning: dspy.Reasoning = dspy.OutputField()
+        answer: str = dspy.OutputField()
+
+    lm = dspy.LM("gemini/au.gemini-3.5-flash", api_key="fake", cache=False)
+    lm_kwargs: dict = {}
+    adapted = dspy.Reasoning.adapt_to_native_lm_feature(Sig, "reasoning", lm, lm_kwargs)
+    assert "reasoning" not in adapted.output_fields
+    assert lm_kwargs["reasoning_effort"] == "low"
+
+    call = prepare(lm, "What is 2+2?", None, lm_kwargs)
+    req = _canonical(call)
+    provider_lm = GeminiLM(api_key="fake")
+    payload = provider_lm._payload(replace(req, model="au.gemini-3.5-flash"))
+    thinking_cfg = payload["generationConfig"]["thinkingConfig"]
+    assert thinking_cfg == {"includeThoughts": True, "thinkingLevel": "low"}

tests/clients/test_lm15_adaptations_and_timeouts.py

@@ -259,3 +259,184 @@ def test_stop_word_spanning_two_text_parts_is_honoured(monkeypatch):
     ))
     lm = dspy.LM("openai/gpt-5", model_type="responses", cache=False)
     assert lm("hello", stop=["STOP"])[0]["text"] == "alpha "
+
+
+def test_gemini_level_class_and_regional_prefixes():
+    from dspy._vendor.lm15.providers.gemini import gemini_level_class
+
+    assert gemini_level_class("gemini-3-flash") is True
+    assert gemini_level_class("gemini-3.5-pro") is True
+    assert gemini_level_class("au.gemini-3.5-flash") is True
+    assert gemini_level_class("eu.gemini-3.5-pro") is True
+    assert gemini_level_class("publishers/google/models/gemini-3.5-flash") is True
+    assert gemini_level_class("publishers/google/models/au.gemini-3.5-pro") is True
+    assert gemini_level_class("gemini-4-pro") is True
+    assert gemini_level_class("gemini-2.5-pro") is False
+    assert gemini_level_class("au.gemini-2.5-flash") is False
+    assert gemini_level_class("au.gemma-3-27b") is False
+
+
+def test_gemini_reasoning_off_and_clamping_for_pro_and_flash():
+    from dspy._vendor.lm15.adaptation import collecting
+    from dspy._vendor.lm15.providers.gemini import GeminiLM
+    from dspy._vendor.lm15.types import Config, Message, Reasoning, Request
+
+    lm = GeminiLM(api_key="fake")
+
+    def _build(model: str, reasoning: Reasoning):
+        req = Request(
+            model=model,
+            messages=(Message.user("hi"),),
+            config=Config(reasoning=reasoning),
+        )
+        with collecting("note") as scope:
+            payload = lm._payload(req)
+        return payload["generationConfig"]["thinkingConfig"], scope.records
+
+    # 1. reasoning="off" on Gemini 3.5 Pro -> thinkingLevel="low", no includeThoughts
+    cfg, adapts = _build("au.gemini-3.5-pro", Reasoning(effort="off"))
+    assert cfg == {"thinkingLevel": "low"}
+    assert len(adapts) == 1 and adapts[0].applied == "low"
+
+    # 2. reasoning="off" on Gemini 3.5 Flash -> thinkingLevel="minimal", no includeThoughts
+    cfg, adapts = _build("au.gemini-3.5-flash", Reasoning(effort="off"))
+    assert cfg == {"thinkingLevel": "minimal"}
+    assert len(adapts) == 1 and adapts[0].applied == "minimal"
+
+    # 3. reasoning="off" on Gemini 2.5 Pro -> thinkingBudget=128
+    cfg, _ = _build("au.gemini-2.5-pro", Reasoning(effort="off"))
+    assert cfg == {"thinkingBudget": 128}
+
+    # 4. reasoning="off" on Gemini 2.5 Flash -> thinkingBudget=0
+    cfg, _ = _build("gemini-2.5-flash", Reasoning(effort="off"))
+    assert cfg == {"thinkingBudget": 0}
+
+    # 5. effort="minimal" on Gemini 3.5 Pro clamped to "low"
+    cfg, adapts = _build("eu.gemini-3.5-pro", Reasoning(effort="minimal"))
+    assert cfg == {"includeThoughts": True, "thinkingLevel": "low"}
+    assert any(a.action == "clamped" and a.applied == "low" for a in adapts)
+
+    # 6. effort="medium" clamped to "high" on gemini-3-pro, preserved on gemini-3.1-pro and gemini-3.5-pro
+    cfg_30, adapts_30 = _build("gemini-3-pro", Reasoning(effort="medium"))
+    assert cfg_30 == {"includeThoughts": True, "thinkingLevel": "high"}
+    assert any(a.action == "clamped" and a.applied == "high" for a in adapts_30)
+
+    cfg_31, adapts_31 = _build("gemini-3.1-pro", Reasoning(effort="medium"))
+    assert cfg_31 == {"includeThoughts": True, "thinkingLevel": "medium"}
+    assert not adapts_31
+
+    cfg_35, adapts_35 = _build("au.gemini-3.5-pro", Reasoning(effort="medium"))
+    assert cfg_35 == {"includeThoughts": True, "thinkingLevel": "medium"}
+    assert not adapts_35
+
+
+def test_gemini_streaming_trailing_usage_chunk_preserves_finish_reason():
+    from dspy._vendor.lm15.providers.gemini import GeminiLM
+    from dspy._vendor.lm15.result import coalesce_stream
+    from dspy._vendor.lm15.sse import SSEEvent
+    from dspy._vendor.lm15.types import Message, Request
+
+    lm = GeminiLM(api_key="fake")
+    req = Request(model="gemini-3.5-flash", messages=(Message.user("hi"),))
+
+    chunk1 = SSEEvent(
+        event="message",
+        data=json.dumps({
+            "candidates": [
+                {
+                    "content": {
+                        "role": "model",
+                        "parts": [{"functionCall": {"name": "lookup", "args": {"q": "x"}}}],
+                    },
+                    "finishReason": "STOP",
+                }
+            ]
+        })
+    )
+    chunk2_usage_only = SSEEvent(
+        event="message",
+        data=json.dumps({
+            "candidates": [],
+            "usageMetadata": {
+                "promptTokenCount": 25,
+                "candidatesTokenCount": 12,
+                "thoughtsTokenCount": 40,
+                "cachedContentTokenCount": 10,
+                "totalTokenCount": 77,
+            },
+        })
+    )
+
+    raw_events = [
+        ev
+        for sse in (chunk1, chunk2_usage_only)
+        for ev in lm.parse_stream_events(req, sse)
+    ]
+    coalesced = list(coalesce_stream(iter(raw_events), model="gemini-3.5-flash"))
+    end_events = [e for e in coalesced if e.type == "end"]
+    assert len(end_events) == 1
+    end = end_events[0]
+    assert end.finish_reason == "tool_call"
+    assert end.usage.input_tokens == 25
+    assert end.usage.output_tokens == 12
+    assert end.usage.reasoning_tokens == 40
+    assert end.usage.cache_read_tokens == 10
+    assert end.usage.total_tokens == 77
+
+
+def test_gemini_live_completion_accumulates_reasoning_and_cache_tokens(monkeypatch):
+    from dspy._vendor.lm15.providers.gemini import GeminiLM
+    from dspy._vendor.lm15.types import Message, Request
+
+    class FakeWS:
+        def __init__(self):
+            self.sent = []
+            self.frames = iter([
+                json.dumps({"setupComplete": {}}),
+                json.dumps({
+                    "serverContent": {
+                        "modelTurn": {"parts": [{"text": "Hello"}]},
+                        "turnComplete": False,
+                    },
+                    "usageMetadata": {
+                        "promptTokenCount": 20,
+                        "responseTokenCount": 5,
+                        "thoughtsTokenCount": 14,
+                        "cachedContentTokenCount": 8,
+                        "totalTokenCount": 39,
+                    },
+                }),
+                json.dumps({
+                    "serverContent": {
+                        "modelTurn": {"parts": [{"text": " world"}]},
+                        "turnComplete": True,
+                    },
+                    "usageMetadata": {
+                        "promptTokenCount": 20,
+                        "responseTokenCount": 10,
+                        "thoughtsTokenCount": 18,
+                        "cachedContentTokenCount": 8,
+                        "totalTokenCount": 48,
+                    },
+                }),
+            ])
+
+        def send(self, msg):
+            self.sent.append(msg)
+
+        def recv(self):
+            return next(self.frames)
+
+        def close(self):
+            pass
+
+    lm = GeminiLM(api_key="fake")
+    monkeypatch.setattr(lm, "_live_connect", lambda url: FakeWS())
+    req = Request(model="gemini-2.5-flash-live", messages=(Message.user("hi"),))
+    events = list(lm._stream_via_live_completion(req))
+    end = [e for e in events if e.type == "end"][-1]
+    assert end.usage.input_tokens == 20
+    assert end.usage.output_tokens == 10
+    assert end.usage.reasoning_tokens == 18
+    assert end.usage.cache_read_tokens == 8
+    assert end.usage.total_tokens == 48