diff --git a/docs/reference/provider-capabilities.md b/docs/reference/provider-capabilities.md index 752efe0..a53cd90 100644 --- a/docs/reference/provider-capabilities.md +++ b/docs/reference/provider-capabilities.md @@ -91,6 +91,11 @@ jobs, see [Building With Deferred Delivery](../building-with-deferred-delivery.m when the selected model and reasoning mode support it. Pollux forwards both `reasoning_effort` and `reasoning_budget_tokens`; model-specific acceptance is enforced by the Gemini API. +- Gemini function-call thought signatures and call IDs are preserved in opaque + `continuation` state and replayed with tool results. This applies to both + `interact()` and `stream()` and is required for Gemini 3 tool loops. Keep the + continuation untouched until the pending tool exchange is complete; portable + application-authored `history` does not contain provider signatures. ### OpenAI diff --git a/src/pollux/providers/gemini.py b/src/pollux/providers/gemini.py index a9e12c1..bfa0c91 100644 --- a/src/pollux/providers/gemini.py +++ b/src/pollux/providers/gemini.py @@ -3,7 +3,9 @@ from __future__ import annotations import asyncio +import base64 import contextlib +from dataclasses import replace from datetime import datetime import json import logging @@ -48,6 +50,7 @@ logger = logging.getLogger(__name__) _GEMINI_BATCH_INLINE_LIMIT_BYTES = 20_000_000 +_GEMINI_FUNCTION_CALLS_KEY = "gemini_function_calls" def _provider_hint_payload( @@ -306,15 +309,23 @@ def _build_contents( self, parts: list[Any], history: list[Message] | None, + provider_state: dict[str, object] | None = None, ) -> Any: """Build Gemini contents from history + current parts.""" from google.genai import types input_contents: list[Any] = [] call_id_to_name: dict[str, str] = {} + history_states = ( + provider_state.get("history", []) + if isinstance(provider_state, dict) + else [] + ) + if not isinstance(history_states, list): + history_states = [] if history: - for item in history: + for item_index, item in enumerate(history): if item.role == "tool": call_id = item.tool_call_id name = "unknown_tool" @@ -334,15 +345,25 @@ def _build_contents( types.Content( role="user", parts=[ - types.Part.from_function_response( - name=name, - response=parsed_content, + types.Part( + function_response=types.FunctionResponse( + id=call_id, + name=name, + response=parsed_content, + ) ) ], ) ) elif item.role == "assistant": ast_parts: list[Any] = [] + item_state = ( + history_states[item_index] + if item_index < len(history_states) + and isinstance(history_states[item_index], dict) + else None + ) + signatures = _function_call_signatures(item_state) if item.content: ast_parts.append(types.Part.from_text(text=item.content)) if item.tool_calls: @@ -353,8 +374,13 @@ def _build_contents( except Exception: args_dict = {} ast_parts.append( - types.Part.from_function_call( - name=tc.name, args=args_dict + types.Part( + function_call=types.FunctionCall( + id=tc.id, + name=tc.name, + args=args_dict, + ), + thought_signature=signatures.get(tc.id), ) ) if ast_parts: @@ -430,9 +456,11 @@ async def submit_deferred( config, upload_cache=upload_cache, ) - history, _prev, _ps = _compile.prior_turns(inp) + history, _prev, provider_state = _compile.prior_turns(inp) inlined_request = types.InlinedRequest( - contents=self._build_contents(parts, history or None), + contents=self._build_contents( + parts, history or None, provider_state + ), config=types.GenerateContentConfig( **self._build_config_kwargs(parts, snapshot, requirements) ), @@ -621,9 +649,9 @@ async def generate( from google.genai import types parts = _compile.request_parts(snapshot, input) - history, _previous_response_id, _provider_state = _compile.prior_turns(input) + history, _previous_response_id, provider_state = _compile.prior_turns(input) config_kwargs = self._build_config_kwargs(parts, snapshot, requirements) - contents = self._build_contents(parts, history or None) + contents = self._build_contents(parts, history or None, provider_state) try: response = await client.aio.models.generate_content( @@ -656,20 +684,22 @@ async def stream_generate( ) -> AsyncIterator[ProviderStreamChunk]: """Stream normalized deltas from Gemini's streaming content endpoint. - Gemini replays conversation via history (no provider_state), so thought - text is surfaced for display only. Function calls arrive whole rather - than fragmented, so each one is emitted as a single-fragment tool-call - delta with its own slot index. + Gemini replays conversation via history. Thought text is surfaced for + display, while signed function-call state is retained opaquely for the + next continuation turn. Function calls arrive whole rather than + fragmented, so each one is emitted as a single-fragment tool-call delta + with its own slot index. """ client = self._get_client() from google.genai import types parts = _compile.request_parts(snapshot, input) - history, _previous_response_id, _provider_state = _compile.prior_turns(input) + history, _previous_response_id, provider_state = _compile.prior_turns(input) config_kwargs = self._build_config_kwargs(parts, snapshot, requirements) - contents = self._build_contents(parts, history or None) + contents = self._build_contents(parts, history or None, provider_state) tool_call_index = 0 + function_call_state: list[dict[str, str]] = [] stream: Any = None try: stream = await client.aio.models.generate_content_stream( @@ -681,7 +711,23 @@ async def stream_generate( for chunk in self._stream_response_to_chunks(response, tool_call_index): if chunk.tool_calls: tool_call_index += len(chunk.tool_calls) + if chunk.provider_state: + entries = chunk.provider_state.get(_GEMINI_FUNCTION_CALLS_KEY) + if isinstance(entries, list): + function_call_state.extend( + entry for entry in entries if isinstance(entry, dict) + ) + remaining_state = dict(chunk.provider_state) + remaining_state.pop(_GEMINI_FUNCTION_CALLS_KEY, None) + chunk = replace( + chunk, + provider_state=remaining_state or None, + ) yield chunk + if function_call_state: + yield ProviderStreamChunk( + provider_state={_GEMINI_FUNCTION_CALLS_KEY: function_call_state} + ) except asyncio.CancelledError: raise except Exception as e: @@ -717,6 +763,7 @@ def _stream_response_to_chunks( _field(function_call, "id") or f"call_{uuid.uuid4().hex[:8]}" ) arguments = json.dumps(_field(function_call, "args") or {}) + state_entry = _function_call_state_entry(str(call_id), part) chunks.append( ProviderStreamChunk( tool_calls=( @@ -726,7 +773,12 @@ def _stream_response_to_chunks( name=str(_field(function_call, "name", "")), arguments=arguments, ), - ) + ), + provider_state=( + {_GEMINI_FUNCTION_CALLS_KEY: [state_entry]} + if state_entry is not None + else None + ), ) ) local_tool_count += 1 @@ -1238,6 +1290,26 @@ def _parse_response(self, response: Any) -> ProviderResponse: ) ) + function_call_parts: list[Any] = [] + candidates = _field(response, "candidates", []) or [] + if isinstance(candidates, list) and candidates: + content = _field(candidates[0], "content") + candidate_parts = ( + _field(content, "parts", []) if content is not None else [] + ) + if isinstance(candidate_parts, list): + function_call_parts = [ + part + for part in candidate_parts + if _field(part, "function_call") is not None + ] + + function_call_state = [ + state + for call, part in zip(tool_calls, function_call_parts, strict=False) + if (state := _function_call_state_entry(call.id, part)) is not None + ] + finish_reason: str | None = None reasoning_parts: list[str] = [] artifacts: dict[str, Any] = {} @@ -1273,10 +1345,52 @@ def _parse_response(self, response: Any) -> ProviderResponse: tool_calls=tool_calls if tool_calls else None, response_id=None, finish_reason=finish_reason, + provider_state=( + {_GEMINI_FUNCTION_CALLS_KEY: function_call_state} + if function_call_state + else None + ), artifacts=artifacts or None, ) +def _function_call_state_entry(call_id: str, part: Any) -> dict[str, str] | None: + """Encode one Gemini function-call thought signature for continuation replay.""" + signature = _field(part, "thought_signature") + if isinstance(signature, bytes) and signature: + encoded = base64.b64encode(signature).decode("ascii") + elif isinstance(signature, str) and signature: + encoded = signature + else: + return None + return {"id": call_id, "thought_signature": encoded} + + +def _function_call_signatures( + provider_state: object, +) -> dict[str, bytes]: + """Decode preserved Gemini signatures, keyed by normalized tool-call ID.""" + if not isinstance(provider_state, dict): + return {} + entries = provider_state.get(_GEMINI_FUNCTION_CALLS_KEY) + if not isinstance(entries, list): + return {} + + signatures: dict[str, bytes] = {} + for entry in entries: + if not isinstance(entry, dict): + continue + call_id = entry.get("id") + encoded = entry.get("thought_signature") + if not isinstance(call_id, str) or not isinstance(encoded, str): + continue + try: + signatures[call_id] = base64.b64decode(encoded, validate=True) + except ValueError: + continue + return signatures + + def _field(obj: Any, key: str, default: Any = None) -> Any: """Read a field from either an SDK object or a raw JSON dict.""" if isinstance(obj, dict): diff --git a/tests/providers/test_gemini_contract.py b/tests/providers/test_gemini_contract.py index 06b3ba4..0aad8b6 100644 --- a/tests/providers/test_gemini_contract.py +++ b/tests/providers/test_gemini_contract.py @@ -838,6 +838,39 @@ def test_gemini_parse_response_extracts_tool_calls() -> None: assert result.tool_calls[0].arguments == '{"city": "NYC"}' +def test_gemini_parse_response_preserves_function_call_thought_signature() -> None: + """Gemini function-call signatures survive in opaque replay state.""" + provider = GeminiProvider("test-key") + + fc = _obj(id="call_abc", name="get_weather", args={"city": "NYC"}) + function_call_part = _obj( + function_call=fc, + thought_signature=b"signed-thought", + text=None, + thought=False, + ) + fake_response = _obj( + text="", + parsed=None, + usage_metadata=None, + function_calls=[fc], + candidates=[ + _obj(content=_obj(parts=[function_call_part]), finish_reason="STOP") + ], + ) + + result = provider._parse_response(fake_response) + + assert result.provider_state == { + "gemini_function_calls": [ + { + "id": "call_abc", + "thought_signature": "c2lnbmVkLXRob3VnaHQ=", + } + ] + } + + def test_gemini_parse_response_generates_id_when_missing() -> None: """Gemini may omit fc.id; a stable synthetic ID should be generated.""" provider = GeminiProvider("test-key") diff --git a/tests/providers/test_gemini_stream.py b/tests/providers/test_gemini_stream.py index 84b1c22..3a26d4b 100644 --- a/tests/providers/test_gemini_stream.py +++ b/tests/providers/test_gemini_stream.py @@ -29,6 +29,7 @@ def _part(**kwargs: Any) -> SimpleNamespace: kwargs.setdefault("text", None) kwargs.setdefault("thought", False) kwargs.setdefault("function_call", None) + kwargs.setdefault("thought_signature", None) return SimpleNamespace(**kwargs) @@ -54,7 +55,8 @@ def _thinking_tool_chunks() -> list[SimpleNamespace]: _part( function_call=SimpleNamespace( id=None, name="get_weather", args={"city": "NYC"} - ) + ), + thought_signature=b"stream-signature", ) ], finish="STOP", @@ -154,6 +156,11 @@ async def test_gemini_stream_through_interaction_assembles_output() -> None: assert done.usage.total_tokens == 15 assert done.usage.reasoning_tokens == 2 assert done.metrics.finish_reason == "stop" + assert done.continuation is not None + state = done.continuation.to_jsonable()["provider_state"] + assert state["gemini_function_calls"][0]["thought_signature"] == ( + "c3RyZWFtLXNpZ25hdHVyZQ==" + ) @pytest.mark.asyncio diff --git a/tests/providers/test_tool_history.py b/tests/providers/test_tool_history.py index baa3f93..61bbb5f 100644 --- a/tests/providers/test_tool_history.py +++ b/tests/providers/test_tool_history.py @@ -246,6 +246,68 @@ async def fake_generate_content( assert contents[2].parts[1].text == "Continue the conversation" +@pytest.mark.asyncio +async def test_gemini_replays_function_call_signature_and_ids() -> None: + """Opaque Gemini replay state restores signatures and call/result IDs.""" + captured: dict[str, Any] = {} + + async def fake_generate_content( + *, + model: str, # noqa: ARG001 + contents: Any, + config: Any, # noqa: ARG001 + ) -> Any: + captured["contents"] = contents + return MagicMock(text="ok", parsed=None, usage_metadata=None) + + provider = GeminiProvider("test-key") + fake_models = MagicMock() + fake_models.generate_content = fake_generate_content + fake_aio = MagicMock() + fake_aio.models = fake_models + provider._client = MagicMock() + provider._client.aio = fake_aio + + await provider.generate( + *_gemini( + continuation=make_continuation( + provider="gemini", + messages=( + Message( + role="assistant", + tool_calls=( + ToolCall.from_text( + id="call_abc", + name="get_weather", + arguments_text='{"location": "NYC"}', + ), + ), + ), + ), + provider_state={ + "history": [ + { + "gemini_function_calls": [ + { + "id": "call_abc", + "thought_signature": "c2lnbmVkLXRob3VnaHQ=", + } + ] + } + ] + }, + ), + tool_results=[ToolResult(call_id="call_abc", content='{"temp": 72}')], + ) + ) + + contents = captured["contents"] + function_call = contents[0].parts[0] + assert function_call.function_call.id == "call_abc" + assert function_call.thought_signature == b"signed-thought" + assert contents[1].parts[0].function_response.id == "call_abc" + + @pytest.mark.asyncio async def test_gemini_merges_prompt_into_tool_response_content() -> None: """When history ends with a tool response, the prompt is merged in. diff --git a/tests/test_api.py b/tests/test_api.py index 01fc54e..460610f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -393,6 +393,54 @@ async def test_gemini_reasoning_roundtrip_on_gemini3(gemini_api_key: str) -> Non assert result.reasoning.strip() +@pytest.mark.asyncio +async def test_gemini3_tool_call_signature_roundtrip(gemini_api_key: str) -> None: + """E2E: Gemini 3 accepts a continued tool call with its thought signature.""" + config = Config( + provider="gemini", + model=_GEMINI_REASONING_MODEL, + api_key=gemini_api_key, + ) + environment = Environment( + tools=[ + ToolDeclaration( + name="get_secret", + description="Return a secret code.", + parameters={ + "type": "object", + "properties": {"topic": {"type": "string"}}, + "required": ["topic"], + }, + ) + ] + ) + + first = await pollux.interact( + environment, + Input("Call get_secret once with topic='orbit'."), + config=config, + tool_choice="required", + reasoning_effort="low", + ) + + assert first.tool_calls + second = await pollux.interact( + Environment(), + Input( + continuation=first.continuation, + tool_results=[ + ToolResult.from_value( + call_id=first.tool_calls[0].id, + value={"code": "K9-ORBIT"}, + ) + ], + ), + config=config, + ) + + assert "k9-orbit" in second.text.lower() + + @pytest.mark.asyncio async def test_gemini_url_context_roundtrip( gemini_api_key: str,