From 49e359587df8151b406dc07c3c9af237470ef589 Mon Sep 17 00:00:00 2001 From: Hamza Imran Date: Wed, 21 Oct 2026 02:35:44 +0500 Subject: [PATCH] fix(gemini): exclude thought parts from the model answer (fixes #63) --- agentkit/llm/gemini.py | 2 +- docs/PROGRESS.md | 7 +++ tests/unit/test_audit_regressions.py | 78 +++++++++++++++++++++++++++- 3 files changed, 85 insertions(+), 2 deletions(-) diff --git a/agentkit/llm/gemini.py b/agentkit/llm/gemini.py index b527e85..1deb0b1 100644 --- a/agentkit/llm/gemini.py +++ b/agentkit/llm/gemini.py @@ -165,7 +165,7 @@ async def chat( candidate = candidates[0] if candidate.content and candidate.content.parts: for part in candidate.content.parts: - if getattr(part, "text", None): + if getattr(part, "text", None) and not getattr(part, "thought", False): text_parts.append(part.text) if getattr(part, "function_call", None): fc = part.function_call diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2c7b115..984e8d9 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -407,3 +407,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-06: Exclude Thought Parts from Model Answer +- **Status**: Done +- **Branch**: `fix/gemini-thought-part-leak` +- **Issue**: #63 +- **Commits**: 1 +- **Notes**: In `agentkit/llm/gemini.py`, candidate parts with `thought=True` are excluded from `LLMResponse.text`, preventing internal reasoning from leaking into assistant message turns, persisted history, and final API answers. Parts with `thought=True` that carry a `function_call` still properly yield tool calls and preserve their `thought_signature`. Added unit and API regression tests in `tests/unit/test_audit_regressions.py`. diff --git a/tests/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py index e7c0d9d..11c2712 100644 --- a/tests/unit/test_audit_regressions.py +++ b/tests/unit/test_audit_regressions.py @@ -20,7 +20,7 @@ from redis.asyncio import Redis from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from agentkit.api.deps import get_llm_client +from agentkit.api.deps import get_llm_client, get_memory_store, get_trace_sink from agentkit.api.main import create_app from agentkit.api.ratelimit import RateLimiter from agentkit.config import Settings @@ -469,6 +469,82 @@ def candidates(self) -> Any: await client.chat(messages=[]) +async def test_gemini_thought_parts_excluded_from_response_and_api_answer() -> None: + """AUDIT-06: Parts with thought=True must not leak into response.text or API answer.""" + from agentkit.llm.gemini import GeminiClient + + thought_part = SimpleNamespace( + text="INTERNAL CHAIN OF THOUGHT: step 1, step 2", + thought=True, + function_call=None, + thought_signature=None, + ) + answer_part = SimpleNamespace( + text="42 is the answer", + thought=False, + function_call=None, + thought_signature=None, + ) + candidate = SimpleNamespace(content=SimpleNamespace(parts=[thought_part, answer_part])) + raw = SimpleNamespace(candidates=[candidate], usage_metadata=None, text="42 is the answer") + + sdk_client = AsyncMock() + sdk_client.aio.models.generate_content = AsyncMock(return_value=raw) + + gemini_client = GeminiClient(api_key="key", client=sdk_client) + response = await gemini_client.chat(messages=[]) + + assert response.text == "42 is the answer" + assert "INTERNAL CHAIN OF THOUGHT" not in response.text + + # Test via API endpoint POST /runs + settings = Settings(API_KEYS="secret_token") + app = create_app(settings=settings) + app.dependency_overrides[get_memory_store] = lambda: InMemoryMemoryStore() + app.dependency_overrides[get_trace_sink] = lambda: InMemoryTraceSink() + app.dependency_overrides[get_llm_client] = lambda: gemini_client + + test_client = TestClient(app, raise_server_exceptions=False) + headers = {"X-API-Key": "secret_token"} + resp = test_client.post("/runs", json={"goal": "What is the answer?"}, headers=headers) + assert resp.status_code == 200 + data = resp.json() + assert data["status"] == "succeeded" + assert data["answer"] == "42 is the answer" + assert "INTERNAL CHAIN OF THOUGHT" not in (data.get("answer") or "") + + +async def test_gemini_thought_part_with_function_call_yields_tool_call() -> None: + """AUDIT-06: A thought part with a function_call still yields the tool call and signature.""" + from agentkit.llm.gemini import GeminiClient + + fc = SimpleNamespace(name="calculator", args={"expression": "2+2"}) + part = SimpleNamespace( + text="Thinking about calling calculator...", + thought=True, + function_call=fc, + thought_signature=b"opaque_sig", + ) + candidate = SimpleNamespace(content=SimpleNamespace(parts=[part])) + raw = SimpleNamespace(candidates=[candidate], usage_metadata=None, text="") + + sdk_client = AsyncMock() + sdk_client.aio.models.generate_content = AsyncMock(return_value=raw) + + client = GeminiClient(api_key="key", client=sdk_client) + response = await client.chat(messages=[]) + + # Text must exclude the thought part + assert response.text == "" + assert "Thinking about calling" not in response.text + + # Tool call and thought_signature must be preserved + assert len(response.tool_calls) == 1 + assert response.tool_calls[0].name == "calculator" + assert response.tool_calls[0].arguments == {"expression": "2+2"} + assert response.tool_calls[0].thought_signature == b"opaque_sig" + + # --- sql_readonly: server-side row limit ---------------------------------------------------