feat(gemini): support thought_signature preservation, tool role, and update default model to gemini-2.0-flash

.env.example

@@ -14,7 +14,7 @@ REDIS_URL=redis://localhost:6379/0
 # LLM Provider Configuration
 # Supported providers: "gemini", "openai", "fake"
 LLM_PROVIDER=gemini
-LLM_MODEL=gemini-2.5-flash
+LLM_MODEL=gemini-2.0-flash
 GEMINI_API_KEY=your_gemini_api_key_here
 OPENAI_API_KEY=your_openai_api_key_here
 

agentkit/config.py

@@ -32,7 +32,7 @@ class Settings(BaseSettings):
         description="LLM provider: gemini, openai, or fake.",
     )
     LLM_MODEL: str = Field(
-        default="gemini-2.5-flash",
+        default="gemini-2.0-flash",
         description="LLM model identifier.",
     )
     GEMINI_API_KEY: str | None = Field(

agentkit/llm/factory.py

@@ -37,7 +37,7 @@ def create_llm_client(
     normalized_provider = provider.strip().lower()
 
     if normalized_provider == "gemini":
-        default_model = model or "gemini-2.5-flash"
+        default_model = model or "gemini-2.0-flash"
         return GeminiClient(api_key=api_key, model=default_model, **kwargs)
 
     if normalized_provider == "openai":

agentkit/llm/gemini.py

@@ -18,14 +18,14 @@ class GeminiClient(LLMClient):
     def __init__(
         self,
         api_key: str | None = None,
-        model: str = "gemini-2.5-flash",
+        model: str = "gemini-2.0-flash",
         client: Any | None = None,
     ) -> None:
         """Initialize the Gemini client adapter.
 
         Args:
             api_key: Gemini API key. If omitted, google-genai attempts ADC/env lookup.
-            model: Gemini model identifier (default: gemini-2.5-flash).
+            model: Gemini model identifier (default: gemini-2.0-flash).
             client: Optional pre-configured genai.Client instance (useful for mocking).
         """
         self.model = model
@@ -60,6 +60,13 @@ async def chat(
         system_instruction: str | None = None
         contents: list[types.Content] = []
 
+        # Map tool_call_id to tool_name for matching function responses
+        tool_id_to_name: dict[str, str] = {}
+        for msg in messages:
+            if msg.role == Role.ASSISTANT:
+                for tc in msg.tool_calls:
+                    tool_id_to_name[tc.id] = tc.name
+
         # Convert messages
         for msg in messages:
             if msg.role == Role.SYSTEM:
@@ -76,18 +83,31 @@ async def chat(
                 if msg.content:
                     parts.append(types.Part.from_text(text=msg.content))
                 for tc in msg.tool_calls:
-                    parts.append(
-                        types.Part.from_function_call(
-                            name=tc.name,
-                            args=tc.arguments,
+                    if tc.thought_signature is not None:
+                        parts.append(
+                            types.Part(
+                                function_call=types.FunctionCall(
+                                    name=tc.name,
+                                    args=tc.arguments,
+                                ),
+                                thought_signature=tc.thought_signature,
+                            )
+                        )
+                    else:
+                        parts.append(
+                            types.Part.from_function_call(
+                                name=tc.name,
+                                args=tc.arguments,
+                            )
                         )
-                    )
                 contents.append(types.Content(role="model", parts=parts))
             elif msg.role == Role.TOOL:
-                tool_name = msg.tool_call_id or "tool"
+                tool_name = (
+                    tool_id_to_name.get(msg.tool_call_id or "") or msg.tool_call_id or "tool"
+                )
                 contents.append(
                     types.Content(
-                        role="user",
+                        role="tool",
                         parts=[
                             types.Part.from_function_response(
                                 name=tool_name,
@@ -144,11 +164,14 @@ async def chat(
                         text_parts.append(part.text)
                     if getattr(part, "function_call", None):
                         fc = part.function_call
+                        raw_sig = getattr(part, "thought_signature", None)
+                        sig = raw_sig if isinstance(raw_sig, bytes) else None
                         tool_calls.append(
                             ToolCall(
                                 id=str(uuid.uuid4())[:8],
                                 name=fc.name,
                                 arguments=dict(fc.args or {}),
+                                thought_signature=sig,
                             )
                         )
 
@@ -166,7 +189,9 @@ async def chat(
         metadata: dict[str, Any] = {}
         if hasattr(raw_response, "model_dump"):
             try:
-                metadata = raw_response.model_dump()
+                dumped = raw_response.model_dump()
+                if isinstance(dumped, dict):
+                    metadata = dumped
             except Exception:
                 metadata = {}
 

agentkit/memory/redis_store.py

@@ -72,7 +72,7 @@ async def save_messages(self, session_id: str, messages: list[Message]) -> None:
             messages = self.truncate_messages(messages, self.max_messages)
 
         key = self._key(session_id)
-        serialized = json.dumps([m.model_dump() for m in messages])
+        serialized = json.dumps([m.model_dump(mode="json") for m in messages])
         await self._client.set(key, serialized, ex=self.ttl_s)
         logger.debug(
             "Saved %d messages to Redis for session %s (TTL: %ds)",

agentkit/tools/models.py

@@ -67,6 +67,10 @@ class ToolCall(BaseModel):
         validation_alias=AliasChoices("arguments", "args"),
         description="Parsed tool arguments.",
     )
+    thought_signature: bytes | None = Field(
+        default=None,
+        description="Optional provider thought signature for preserving tool call state across turns.",
+    )
 
 
 class ToolSchema(BaseModel):

docker-compose.yml

@@ -76,7 +76,7 @@ services:
       # Compose keeps ${VAR:-} so `down`, `ps` and `logs` still work without it.
       - API_KEYS=${API_KEYS:-}
       - LLM_PROVIDER=${LLM_PROVIDER:-gemini}
-      - LLM_MODEL=${LLM_MODEL:-gemini-2.5-flash}
+      - LLM_MODEL=${LLM_MODEL:-gemini-2.0-flash}
       - GEMINI_API_KEY=${GEMINI_API_KEY:-}
       - OPENAI_API_KEY=${OPENAI_API_KEY:-}
     volumes:

tests/unit/test_config.py

@@ -14,7 +14,7 @@ def test_default_settings() -> None:
         API_KEYS="key1,key2",
     )
     assert settings.LLM_PROVIDER == "gemini"
-    assert settings.LLM_MODEL == "gemini-2.5-flash"
+    assert settings.LLM_MODEL == "gemini-2.0-flash"
     assert settings.MAX_STEPS == 10
     assert settings.RUN_TIMEOUT_S == 60
     assert settings.TOOL_TIMEOUT_S == 15

tests/unit/test_gemini_adapter.py

@@ -126,3 +126,81 @@ async def test_gemini_multi_turn_history_translation(mock_genai_client: MagicMoc
     response = await adapter.chat(messages)
     assert response.text == "Result is 25."
     assert mock_genai_client.aio.models.generate_content.called
+
+
+@pytest.mark.asyncio
+async def test_gemini_thought_signature_capture_and_forward(mock_genai_client: MagicMock) -> None:
+    """Verify thought_signature from candidate is captured in ToolCall and forwarded back."""
+    mock_fc = MagicMock()
+    mock_fc.name = "calculator"
+    mock_fc.args = {"expression": "2 + 2"}
+
+    mock_part = MagicMock()
+    mock_part.text = None
+    mock_part.function_call = mock_fc
+    mock_part.thought_signature = b"opaque_thought_sig_xyz"
+
+    mock_candidate = MagicMock()
+    mock_candidate.content.parts = [mock_part]
+
+    mock_response = MagicMock()
+    mock_response.candidates = [mock_candidate]
+    mock_response.text = ""
+    mock_response.usage_metadata.prompt_token_count = 10
+    mock_response.usage_metadata.candidates_token_count = 5
+    mock_response.model_dump.return_value = {}
+
+    mock_genai_client.aio.models.generate_content.return_value = mock_response
+
+    adapter = GeminiClient(api_key="fake-key", client=mock_genai_client)
+
+    # 1. Turn 1: model returns tool call with thought_signature
+    messages = [Message.user("What is 2+2?")]
+    turn1_resp = await adapter.chat(messages)
+
+    assert len(turn1_resp.tool_calls) == 1
+    tc = turn1_resp.tool_calls[0]
+    assert tc.name == "calculator"
+    assert tc.thought_signature == b"opaque_thought_sig_xyz"
+
+    # 2. Turn 2: send assistant tool call with thought_signature and tool result back
+    turn2_messages = [
+        Message.user("What is 2+2?"),
+        Message.assistant(content="", tool_calls=[tc]),
+        Message(role=Role.TOOL, content="4", tool_call_id=tc.id),
+    ]
+
+    mock_response2 = MagicMock()
+    mock_response2.candidates = []
+    mock_response2.text = "The answer is 4."
+    mock_genai_client.aio.models.generate_content.return_value = mock_response2
+
+    await adapter.chat(turn2_messages)
+
+    # Verify generate_content was called with correct contents structure
+    call_kwargs = mock_genai_client.aio.models.generate_content.call_args.kwargs
+    contents = call_kwargs["contents"]
+    assert len(contents) == 3
+
+    # User message
+    assert contents[0].role == "user"
+
+    # Model message must preserve thought_signature on the function call part
+    model_content = contents[1]
+    assert model_content.role == "model"
+    model_part = model_content.parts[0]
+    assert model_part.function_call.name == "calculator"
+    assert model_part.thought_signature == b"opaque_thought_sig_xyz"
+
+    # Tool message must have role="tool" and function_response name resolved to "calculator"
+    tool_content = contents[2]
+    assert tool_content.role == "tool"
+    tool_part = tool_content.parts[0]
+    assert tool_part.function_response.name == "calculator"
+    assert tool_part.function_response.response == {"result": "4"}
+
+
+def test_gemini_default_model(mock_genai_client: MagicMock) -> None:
+    """Verify GeminiClient defaults to gemini-2.0-flash."""
+    adapter = GeminiClient(api_key="fake-key", client=mock_genai_client)
+    assert adapter.model == "gemini-2.0-flash"

tests/unit/test_llm_factory.py

@@ -45,12 +45,12 @@ def test_create_unknown_provider_raises() -> None:
 def test_create_from_settings_gemini() -> None:
     settings = Settings(
         LLM_PROVIDER="gemini",
-        LLM_MODEL="gemini-2.5-flash",
+        LLM_MODEL="gemini-2.0-flash",
         GEMINI_API_KEY="test-gemini-key",
     )
     client = create_llm_client_from_settings(settings)
     assert isinstance(client, GeminiClient)
-    assert client.model == "gemini-2.5-flash"
+    assert client.model == "gemini-2.0-flash"
 
 
 def test_create_from_settings_openai() -> None:

tests/unit/test_tool_models.py

@@ -37,8 +37,18 @@ def test_tool_call_model() -> None:
         "id": "call_123",
         "name": "calculator",
         "arguments": {"expression": "2 + 2"},
+        "thought_signature": None,
     }
 
+    call_with_sig = ToolCall(
+        id="call_456",
+        name="search",
+        arguments={"query": "test"},
+        thought_signature=b"test_signature",
+    )
+    assert call_with_sig.thought_signature == b"test_signature"
+    assert call_with_sig.model_dump(mode="json")["thought_signature"] == "test_signature"
+
 
 def test_tool_schema_model() -> None:
     """Verify ToolSchema definition."""