From 50b081b2ee2bb8d3f48d1c3d990f2dd97744aea0 Mon Sep 17 00:00:00 2001 From: Hamza Imran Date: Wed, 21 Oct 2026 02:44:26 +0500 Subject: [PATCH] fix(gemini): group parallel function responses into one turn (fixes #62) --- agentkit/llm/gemini.py | 41 +++++++++----- docs/PROGRESS.md | 8 +++ tests/unit/test_gemini_adapter.py | 89 +++++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 14 deletions(-) diff --git a/agentkit/llm/gemini.py b/agentkit/llm/gemini.py index b527e85..4070234 100644 --- a/agentkit/llm/gemini.py +++ b/agentkit/llm/gemini.py @@ -105,17 +105,27 @@ async def chat( 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", - parts=[ - types.Part.from_function_response( - name=tool_name, - response={"result": msg.content}, - ) - ], - ) + tool_part = types.Part.from_function_response( + name=tool_name, + response={"result": msg.content}, ) + if msg.tool_call_id and tool_part.function_response is not None: + tool_part.function_response.id = msg.tool_call_id + + prev_parts = ( + contents[-1].parts if contents and contents[-1].role == "user" else None + ) + if prev_parts is not None and any( + getattr(p, "function_response", None) for p in prev_parts + ): + prev_parts.append(tool_part) + else: + contents.append( + types.Content( + role="user", + parts=[tool_part], + ) + ) # Convert tools to Gemini FunctionDeclarations config_tools: list[Any] | None = None @@ -171,12 +181,15 @@ async def chat( fc = part.function_call raw_sig = getattr(part, "thought_signature", None) sig = raw_sig if isinstance(raw_sig, bytes) else None + fc_id = getattr(fc, "id", None) + call_id = ( + fc_id if isinstance(fc_id, str) and fc_id else str(uuid.uuid4()) + ) tool_calls.append( ToolCall( - # Full UUID: tool_call_id is the join key between an - # assistant tool call and its tool result, so a truncated - # 32-bit id risks collisions within a long session. - id=str(uuid.uuid4()), + # Preserve the provider's FunctionCall.id when present; + # fall back to a full UUID so tool_call_id remains unique. + id=call_id, name=fc.name, arguments=dict(fc.args or {}), thought_signature=sig, diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2c7b115..939d8a3 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -61,6 +61,7 @@ This document tracks execution progress across all backlog tasks, including pull | AUD-09 | AUD | Live End-to-End Infrastructure Verification & Documentation Proof | `docs/e2e-verification-report` | #53 | 3 | Done | Standardized model, verified live stack, and documented in E2E_REPORT.md | | AUD-10 | AUD | Rate Limiter Outage Fallback & Provider-Valid Session History | `fix/ratelimit-and-session-repair` | #56 | 3 | Done | In-memory fallback on Redis connection failures; synthetic tool results repair aborted sessions | | M13-T06 | M13 | Whole-Project Audit Remediation | `fix/audit-hardening` | #47 | 3 | Done | Fail-closed config and LLM resolution, SSRF post-connect check, atomic fallback | +| AUDIT-05 | AUD | Gemini Parallel Function Calls Grouping & ID Preservation | `fix/gemini-parallel-function-calls` | #79 | 1 | Done | Group parallel function responses into single user turn; preserve provider call IDs | --- @@ -407,3 +408,10 @@ This document tracks execution progress across all backlog tasks, including pull - **PR**: #47 - **Commits**: 3 - **Notes**: Remediated 17 findings from a full-codebase audit, rebased onto the AUD-01/AUD-04/AUD-05 follow-ups. Critical: removed the `ak_test_key_12345` default from `Settings.API_KEYS` and made the field required; replaced the `FakeLLMClient` fallback in `get_llm_client` with a `ServiceUnavailableError` (503) boundary around provider construction; kept the AUD-04 factory-level `RetryingLLMClient` wrapping and made it configurable through new `LLM_MAX_RETRIES`/`LLM_RETRY_BASE_DELAY_S`/`LLM_RETRY_MAX_DELAY_S` settings; made the per-request `Agent` assembly's override explicit as `app.state.agent_override`; forwarded `get_run_trace`/`list_runs` through `TieredMemoryStore`; moved `http_fetch` DNS off the event loop with `asyncio.to_thread`, added a post-connect `server_addr` check against the validated addresses to close the DNS-rebinding TOCTOU window, and replaced full-body buffering with a byte-budgeted stream; replaced the rate limiter's non-scripting fallback with `WATCH`/`MULTI`/`EXEC`, leaving the AUD-10 in-memory fallback for Redis outages intact. Moderate: corrected the `RedisMemoryStore.truncate_messages` off-by-one that returned `max_messages + 1` entries when a system prompt was present; switched generated `tool_call_id`s to full UUIDs; wrapped response parsing in both LLM adapters inside the error boundary so malformed responses and malformed tool-argument JSON raise `LLMProviderError`; guarded `RetryingLLMClient.__getattr__` against recursion on a partially constructed instance; sanitized `/health` details so the unauthenticated endpoint no longer echoes driver error strings; made `sql_readonly` mask string literals and quoted identifiers before its keyword scan, allow `#>`/`#>>`/`#-`, apply the row limit in SQL and push `statement_timeout` to PostgreSQL. Minor: stopped logging raw run goals at INFO (length only, content at DEBUG); replaced `split(";")` in `scripts/seed.py` with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Added `tests/unit/test_audit_regressions.py` and a real-Redis concurrency regression test pinning the fallback fix. Recorded in ADR-009, ADR-010, ADR-011. Suite: 298 passed, 25 skipped; coverage 92.86% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean. + +### AUDIT-05: Gemini Parallel Function Calls Grouping & ID Preservation +- **Status**: Done +- **PR**: #79 +- **Commits**: 1 +- **Notes**: In `agentkit/llm/gemini.py`, resolved upstream Google Gemini API requirement for parallel function call responses: grouped consecutive `Role.TOOL` messages following an assistant tool call turn into a single `types.Content(role="user")` container with one `function_response` part per tool call. Preserved `FunctionCall.id` when provided by the API and set `function_response.id` matching `tool_call_id`. Added comprehensive regression tests in `tests/unit/test_gemini_adapter.py`. + diff --git a/tests/unit/test_gemini_adapter.py b/tests/unit/test_gemini_adapter.py index d383f2d..29a5b60 100644 --- a/tests/unit/test_gemini_adapter.py +++ b/tests/unit/test_gemini_adapter.py @@ -207,3 +207,92 @@ def test_gemini_default_model(mock_genai_client: MagicMock) -> None: """Verify GeminiClient defaults to gemini-3.1-flash-lite.""" adapter = GeminiClient(api_key="fake-key", client=mock_genai_client) assert adapter.model == "gemini-3.1-flash-lite" + + +@pytest.mark.asyncio +async def test_gemini_parallel_function_calls_grouped_and_preserve_ids( + mock_genai_client: MagicMock, +) -> None: + """AUDIT-05: Parallel tool calls must preserve FunctionCall.id and group responses into one Content.""" + # 1. Test extraction: FunctionCall.id is preserved + fc1 = MagicMock() + fc1.id = "call_calc_1" + fc1.name = "calculator" + fc1.args = {"expression": "2 + 2"} + + fc2 = MagicMock() + fc2.id = "call_calc_2" + fc2.name = "calculator" + fc2.args = {"expression": "10 * 10"} + + mock_part1 = MagicMock() + mock_part1.text = None + mock_part1.thought = False + mock_part1.function_call = fc1 + mock_part1.thought_signature = None + + mock_part2 = MagicMock() + mock_part2.text = None + mock_part2.thought = False + mock_part2.function_call = fc2 + mock_part2.thought_signature = None + + mock_candidate = MagicMock() + mock_candidate.content.parts = [mock_part1, mock_part2] + + 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 = 20 + + mock_genai_client.aio.models.generate_content.return_value = mock_response + + adapter = GeminiClient(api_key="fake-key", client=mock_genai_client) + res = await adapter.chat([Message.user("Do two math problems")]) + + assert len(res.tool_calls) == 2 + assert res.tool_calls[0].id == "call_calc_1" + assert res.tool_calls[0].name == "calculator" + assert res.tool_calls[1].id == "call_calc_2" + assert res.tool_calls[1].name == "calculator" + + # 2. Test outbound conversion: tool responses must be grouped into ONE user Content + messages = [ + Message.user("Do two math problems"), + Message.assistant(content="", tool_calls=res.tool_calls), + Message(role=Role.TOOL, content="4", tool_call_id="call_calc_1"), + Message(role=Role.TOOL, content="100", tool_call_id="call_calc_2"), + ] + + mock_response2 = MagicMock() + mock_response2.candidates = [] + mock_response2.text = "Results are 4 and 100" + mock_genai_client.aio.models.generate_content.return_value = mock_response2 + + await adapter.chat(messages) + + call_kwargs = mock_genai_client.aio.models.generate_content.call_args.kwargs + contents = call_kwargs["contents"] + + # Should have 3 contents: [user, model, user(with 2 tool response parts)] + assert len(contents) == 3 + assert contents[0].role == "user" + assert contents[1].role == "model" + assert len(contents[1].parts) == 2 # 2 function calls + + tool_content = contents[2] + assert tool_content.role == "user" + # Exactly as many function_response parts as the assistant turn had function_call parts + assert len(tool_content.parts) == 2 + + # Verify responses are emitted in call order with matching ids + part_resp1 = tool_content.parts[0] + assert part_resp1.function_response.name == "calculator" + assert part_resp1.function_response.id == "call_calc_1" + assert part_resp1.function_response.response == {"result": "4"} + + part_resp2 = tool_content.parts[1] + assert part_resp2.function_response.name == "calculator" + assert part_resp2.function_response.id == "call_calc_2" + assert part_resp2.function_response.response == {"result": "100"}