From 8c36e1483d2e318f10bb9ee88e7ffe50f0310740 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Thu, 1 Oct 2026 12:01:24 +0530 Subject: [PATCH 1/2] test(gemini): add tests for thought_signature propagation and tool name resolution --- tests/unit/test_gemini_adapter.py | 78 +++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/tests/unit/test_gemini_adapter.py b/tests/unit/test_gemini_adapter.py index 8d94885..1dc9d9d 100644 --- a/tests/unit/test_gemini_adapter.py +++ b/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" From f2364f85d489aebb54a4f166af0a60aae2288df3 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Thu, 1 Oct 2026 12:04:13 +0530 Subject: [PATCH 2/2] feat(gemini): support thought_signature preservation, tool role, and update default model to gemini-2.0-flash --- .env.example | 2 +- agentkit/config.py | 2 +- agentkit/llm/factory.py | 2 +- agentkit/llm/gemini.py | 45 ++++++++++++++++++++++++++-------- agentkit/memory/redis_store.py | 2 +- agentkit/tools/models.py | 4 +++ docker-compose.yml | 2 +- tests/unit/test_config.py | 2 +- tests/unit/test_llm_factory.py | 4 +-- tests/unit/test_tool_models.py | 10 ++++++++ 10 files changed, 57 insertions(+), 18 deletions(-) diff --git a/.env.example b/.env.example index 23bd7e5..4ead17f 100644 --- a/.env.example +++ b/.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 diff --git a/agentkit/config.py b/agentkit/config.py index 4fea7d4..3fc702f 100644 --- a/agentkit/config.py +++ b/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( diff --git a/agentkit/llm/factory.py b/agentkit/llm/factory.py index 5a1de8c..e840cf0 100644 --- a/agentkit/llm/factory.py +++ b/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": diff --git a/agentkit/llm/gemini.py b/agentkit/llm/gemini.py index b18ee85..d9815a6 100644 --- a/agentkit/llm/gemini.py +++ b/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 = {} diff --git a/agentkit/memory/redis_store.py b/agentkit/memory/redis_store.py index 8daf841..39bb896 100644 --- a/agentkit/memory/redis_store.py +++ b/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)", diff --git a/agentkit/tools/models.py b/agentkit/tools/models.py index 41af815..391b8f3 100644 --- a/agentkit/tools/models.py +++ b/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): diff --git a/docker-compose.yml b/docker-compose.yml index f9b6fca..b153448 100644 --- a/docker-compose.yml +++ b/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: diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f192812..f06845d 100644 --- a/tests/unit/test_config.py +++ b/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 diff --git a/tests/unit/test_llm_factory.py b/tests/unit/test_llm_factory.py index 6f79293..f0d6cc0 100644 --- a/tests/unit/test_llm_factory.py +++ b/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: diff --git a/tests/unit/test_tool_models.py b/tests/unit/test_tool_models.py index 5382320..f93e326 100644 --- a/tests/unit/test_tool_models.py +++ b/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."""