4545from .utils import create_tools_config , to_fnc_ctx , to_response_format
4646
4747
48+ def _is_gemini_3_model (model : str ) -> bool :
49+ """Check if model is Gemini 3 series"""
50+ return "gemini-3" in model .lower () or model .lower ().startswith ("gemini-3" )
51+
52+
4853@dataclass
4954class _LLMOptions :
5055 model : ChatModels | str
@@ -157,10 +162,13 @@ def __init__(
157162 # Validate thinking_config
158163 if is_given (thinking_config ):
159164 _thinking_budget = None
165+ _thinking_level = None
160166 if isinstance (thinking_config , dict ):
161167 _thinking_budget = thinking_config .get ("thinking_budget" )
168+ _thinking_level = thinking_config .get ("thinking_level" )
162169 elif isinstance (thinking_config , types .ThinkingConfig ):
163170 _thinking_budget = thinking_config .thinking_budget
171+ _thinking_level = getattr (thinking_config , "thinking_level" , None )
164172
165173 if _thinking_budget is not None :
166174 if not isinstance (_thinking_budget , int ):
@@ -191,6 +199,8 @@ def __init__(
191199 project = gcp_project ,
192200 location = gcp_location ,
193201 )
202+ # Store thought_signatures for Gemini 3 multi-turn function calling
203+ self ._thought_signatures : dict [str , bytes ] = {}
194204
195205 @property
196206 def model (self ) -> str :
@@ -284,9 +294,45 @@ def chat(
284294 if is_given (self ._opts .seed ):
285295 extra ["seed" ] = self ._opts .seed
286296
287- # Add thinking config if thinking_budget is provided
297+ # Handle thinking_config based on model version
288298 if is_given (self ._opts .thinking_config ):
289- extra ["thinking_config" ] = self ._opts .thinking_config
299+ is_gemini_3 = _is_gemini_3_model (self ._opts .model )
300+ thinking_cfg = self ._opts .thinking_config
301+
302+ # Extract both parameters
303+ _budget = None
304+ _level = None
305+ if isinstance (thinking_cfg , dict ):
306+ _budget = thinking_cfg .get ("thinking_budget" )
307+ _level = thinking_cfg .get ("thinking_level" )
308+ elif isinstance (thinking_cfg , types .ThinkingConfig ):
309+ _budget = thinking_cfg .thinking_budget
310+ _level = getattr (thinking_cfg , "thinking_level" , None )
311+
312+ if is_gemini_3 :
313+ # Gemini 3: only support thinking_level
314+ if _budget is not None and _level is None :
315+ logger .warning (
316+ f"Model { self ._opts .model } is Gemini 3 which does not support thinking_budget. "
317+ "Please use thinking_level ('low' or 'high') instead. Ignoring thinking_budget."
318+ )
319+ if _level is not None :
320+ # Use thinking_level only (pass as dict since SDK may not have this field yet)
321+ extra ["thinking_config" ] = {"thinking_level" : _level }
322+ # If neither, let API use default
323+ else :
324+ # Gemini 2.5 and earlier: only support thinking_budget
325+ if _level is not None and _budget is None :
326+ raise ValueError (
327+ f"Model { self ._opts .model } does not support thinking_level. "
328+ "Please use thinking_budget (int) instead for Gemini 2.5 and earlier models."
329+ )
330+ if _budget is not None :
331+ # Use thinking_budget only
332+ extra ["thinking_config" ] = types .ThinkingConfig (thinking_budget = _budget )
333+ else :
334+ # Pass through original config if no specific handling needed
335+ extra ["thinking_config" ] = self ._opts .thinking_config
290336
291337 if is_given (self ._opts .automatic_function_calling_config ):
292338 extra ["automatic_function_calling" ] = self ._opts .automatic_function_calling_config
@@ -333,7 +379,14 @@ async def _run(self) -> None:
333379 request_id = utils .shortuuid ()
334380
335381 try :
336- turns_dict , extra_data = self ._chat_ctx .to_provider_format (format = "google" )
382+ # Pass thought_signatures for Gemini 3 multi-turn function calling
383+ thought_sigs = (
384+ self ._llm ._thought_signatures if _is_gemini_3_model (self ._model ) else None
385+ )
386+ turns_dict , extra_data = self ._chat_ctx .to_provider_format (
387+ format = "google" , thought_signatures = thought_sigs
388+ )
389+
337390 turns = [types .Content .model_validate (turn ) for turn in turns_dict ]
338391 function_declarations = to_fnc_ctx (self ._tools )
339392 tools_config = create_tools_config (
@@ -354,6 +407,7 @@ async def _run(self) -> None:
354407 ),
355408 ** self ._extra_kwargs ,
356409 )
410+
357411 stream = await self ._client .aio .models .generate_content_stream (
358412 model = self ._model ,
359413 contents = cast (types .ContentListUnion , turns ),
@@ -433,17 +487,25 @@ async def _run(self) -> None:
433487
434488 def _parse_part (self , id : str , part : types .Part ) -> llm .ChatChunk | None :
435489 if part .function_call :
490+ tool_call = llm .FunctionToolCall (
491+ arguments = json .dumps (part .function_call .args ),
492+ name = part .function_call .name ,
493+ call_id = part .function_call .id or utils .shortuuid ("function_call_" ),
494+ )
495+
496+ # Store thought_signature for Gemini 3 multi-turn function calling
497+ if (
498+ _is_gemini_3_model (self ._model )
499+ and hasattr (part , "thought_signature" )
500+ and part .thought_signature
501+ ):
502+ self ._llm ._thought_signatures [tool_call .call_id ] = part .thought_signature
503+
436504 chat_chunk = llm .ChatChunk (
437505 id = id ,
438506 delta = llm .ChoiceDelta (
439507 role = "assistant" ,
440- tool_calls = [
441- llm .FunctionToolCall (
442- arguments = json .dumps (part .function_call .args ),
443- name = part .function_call .name ,
444- call_id = part .function_call .id or utils .shortuuid ("function_call_" ),
445- )
446- ],
508+ tool_calls = [tool_call ],
447509 content = part .text ,
448510 ),
449511 )
0 commit comments