From e99b0283d28ded90c9d3823559707610bd980add Mon Sep 17 00:00:00 2001 From: Sebastien Coutu Date: Sun, 19 Apr 2026 13:04:37 -0400 Subject: [PATCH 1/3] feat: add thinking config for anthropic llm --- src/google/adk/models/anthropic_llm.py | 103 ++- tests/unittests/models/test_anthropic_llm.py | 619 +++++++++++++++++++ 2 files changed, 720 insertions(+), 2 deletions(-) diff --git a/src/google/adk/models/anthropic_llm.py b/src/google/adk/models/anthropic_llm.py index a14c767f23..882094a82b 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -62,6 +62,13 @@ class _ToolUseAccumulator: args_json: str +class _ThinkingAccumulator(BaseModel): + """Accumulates streamed thinking content block data.""" + + thinking: str = "" + signature: str = "" + + class ClaudeRequest(BaseModel): system_instruction: str messages: Iterable[anthropic_types.MessageParam] @@ -108,7 +115,24 @@ def part_to_message_block( anthropic_types.DocumentBlockParam, anthropic_types.ToolUseBlockParam, anthropic_types.ToolResultBlockParam, + anthropic_types.ThinkingBlockParam, + anthropic_types.RedactedThinkingBlockParam, ]: + if part.thought: + signature_str = ( + part.thought_signature.decode("utf-8") if part.thought_signature else "" + ) + if part.text: + return anthropic_types.ThinkingBlockParam( + type="thinking", + thinking=part.text, + signature=signature_str, + ) + else: + return anthropic_types.RedactedThinkingBlockParam( + type="redacted_thinking", + data=signature_str, + ) if part.text: return anthropic_types.TextBlockParam(text=part.text, type="text") elif part.function_call: @@ -229,6 +253,18 @@ def content_block_to_part( ) part.function_call.id = content_block.id return part + if isinstance(content_block, anthropic_types.ThinkingBlock): + return types.Part( + text=content_block.thinking, + thought=True, + thought_signature=content_block.signature.encode("utf-8"), + ) + if isinstance(content_block, anthropic_types.RedactedThinkingBlock): + return types.Part( + text="", + thought=True, + thought_signature=content_block.data.encode("utf-8"), + ) raise NotImplementedError("Not supported yet.") @@ -349,6 +385,26 @@ def function_declaration_to_tool_param( ) +def _build_thinking_param( + thinking_config: Optional[types.ThinkingConfig], + max_tokens: int, +) -> Union[anthropic_types.ThinkingConfigEnabledParam, NotGiven]: + """Converts ADK ThinkingConfig to Anthropic ThinkingConfigEnabledParam. + + Returns NOT_GIVEN if thinking is not configured or budget is 0. + Clamps budget_tokens to max_tokens - 1 to satisfy the API constraint. + """ + if thinking_config is None: + return NOT_GIVEN + budget = thinking_config.thinking_budget + if not budget: + return NOT_GIVEN + return anthropic_types.ThinkingConfigEnabledParam( + type="enabled", + budget_tokens=min(budget, max_tokens - 1), + ) + + class AnthropicLlm(BaseLlm): """Integration with Claude models via the Anthropic API. @@ -401,6 +457,10 @@ async def generate_content_async( if llm_request.tools_dict else NOT_GIVEN ) + thinking = _build_thinking_param( + llm_request.config.thinking_config if llm_request.config else None, + self.max_tokens, + ) if not stream: message = await self._anthropic_client.messages.create( @@ -410,11 +470,12 @@ async def generate_content_async( tools=tools, tool_choice=tool_choice, max_tokens=self.max_tokens, + thinking=thinking, ) yield message_to_generate_content_response(message) else: async for response in self._generate_content_streaming( - llm_request, messages, tools, tool_choice + llm_request, messages, tools, tool_choice, thinking ): yield response @@ -424,6 +485,9 @@ async def _generate_content_streaming( messages: list[anthropic_types.MessageParam], tools: Union[Iterable[anthropic_types.ToolUnionParam], NotGiven], tool_choice: Union[anthropic_types.ToolChoiceParam, NotGiven], + thinking: Union[ + anthropic_types.ThinkingConfigEnabledParam, NotGiven + ] = NOT_GIVEN, ) -> AsyncGenerator[LlmResponse, None]: """Handles streaming responses from Anthropic models. @@ -439,12 +503,15 @@ async def _generate_content_streaming( tool_choice=tool_choice, max_tokens=self.max_tokens, stream=True, + thinking=thinking, ) # Track content blocks being built during streaming. # Each entry maps a block index to its accumulated state. text_blocks: dict[int, str] = {} tool_use_blocks: dict[int, _ToolUseAccumulator] = {} + thinking_blocks: dict[int, _ThinkingAccumulator] = {} + redacted_thinking_blocks: dict[int, str] = {} input_tokens = 0 output_tokens = 0 @@ -463,6 +530,10 @@ async def _generate_content_streaming( name=block.name, args_json="", ) + elif isinstance(block, anthropic_types.ThinkingBlock): + thinking_blocks[event.index] = _ThinkingAccumulator() + elif isinstance(block, anthropic_types.RedactedThinkingBlock): + redacted_thinking_blocks[event.index] = block.data elif event.type == "content_block_delta": delta = event.delta @@ -479,6 +550,12 @@ async def _generate_content_streaming( elif isinstance(delta, anthropic_types.InputJSONDelta): if event.index in tool_use_blocks: tool_use_blocks[event.index].args_json += delta.partial_json + elif isinstance(delta, anthropic_types.ThinkingDelta): + if event.index in thinking_blocks: + thinking_blocks[event.index].thinking += delta.thinking + elif isinstance(delta, anthropic_types.SignatureDelta): + if event.index in thinking_blocks: + thinking_blocks[event.index].signature = delta.signature elif event.type == "message_delta": output_tokens = event.usage.output_tokens @@ -486,9 +563,31 @@ async def _generate_content_streaming( # Build the final aggregated response with all content. all_parts: list[types.Part] = [] all_indices = sorted( - set(list(text_blocks.keys()) + list(tool_use_blocks.keys())) + set( + list(text_blocks.keys()) + + list(tool_use_blocks.keys()) + + list(thinking_blocks.keys()) + + list(redacted_thinking_blocks.keys()) + ) ) for idx in all_indices: + if idx in thinking_blocks: + acc = thinking_blocks[idx] + all_parts.append( + types.Part( + text=acc.thinking, + thought=True, + thought_signature=acc.signature.encode("utf-8"), + ) + ) + if idx in redacted_thinking_blocks: + all_parts.append( + types.Part( + text="", + thought=True, + thought_signature=redacted_thinking_blocks[idx].encode("utf-8"), + ) + ) if idx in text_blocks: all_parts.append(types.Part.from_text(text=text_blocks[idx])) if idx in tool_use_blocks: diff --git a/tests/unittests/models/test_anthropic_llm.py b/tests/unittests/models/test_anthropic_llm.py index fb44d5c8e7..6d0b8a599f 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -1350,3 +1350,622 @@ async def test_non_streaming_does_not_pass_stream_param(): mock_client.messages.create.assert_called_once() _, kwargs = mock_client.messages.create.call_args assert "stream" not in kwargs + + +# --------------------------------------------------------------------------- +# Extended thinking support (issue #3079) +# --------------------------------------------------------------------------- + +# --- Group E: _build_thinking_param unit tests --- + + +def test_build_thinking_param_none_config(): + """None thinking_config returns NOT_GIVEN.""" + from anthropic import NOT_GIVEN + from google.adk.models.anthropic_llm import _build_thinking_param + + result = _build_thinking_param(None, 8192) + assert result is NOT_GIVEN + + +def test_build_thinking_param_zero_budget(): + """thinking_budget=0 returns NOT_GIVEN.""" + from anthropic import NOT_GIVEN + from google.adk.models.anthropic_llm import _build_thinking_param + + result = _build_thinking_param(types.ThinkingConfig(thinking_budget=0), 8192) + assert result is NOT_GIVEN + + +def test_build_thinking_param_valid_budget(): + """Valid budget returns ThinkingConfigEnabledParam with correct budget_tokens.""" + from google.adk.models.anthropic_llm import _build_thinking_param + + result = _build_thinking_param( + types.ThinkingConfig(thinking_budget=2048), 8192 + ) + assert result["type"] == "enabled" + assert result["budget_tokens"] == 2048 + + +def test_build_thinking_param_clamps_to_max_tokens(): + """budget exceeding max_tokens is clamped to max_tokens - 1.""" + from google.adk.models.anthropic_llm import _build_thinking_param + + result = _build_thinking_param( + types.ThinkingConfig(thinking_budget=10000), 8192 + ) + assert result["type"] == "enabled" + assert result["budget_tokens"] == 8191 + + +# --- Group A: part_to_message_block with thought parts --- + + +def test_part_to_message_block_thought_regular(): + """A thought part with text and signature becomes a ThinkingBlockParam.""" + part = Part( + text="I reasoned about this", thought=True, thought_signature=b"sig_abc" + ) + + result = part_to_message_block(part) + + assert isinstance(result, dict) + assert result["type"] == "thinking" + assert result["thinking"] == "I reasoned about this" + assert result["signature"] == "sig_abc" + + +def test_part_to_message_block_thought_redacted(): + """A thought part with no text becomes a RedactedThinkingBlockParam.""" + part = Part(text="", thought=True, thought_signature=b"redacted_data_xyz") + + result = part_to_message_block(part) + + assert isinstance(result, dict) + assert result["type"] == "redacted_thinking" + assert result["data"] == "redacted_data_xyz" + + +def test_part_to_message_block_thought_checked_before_text(): + """The thought branch fires before the text branch, not producing a TextBlockParam.""" + part = Part( + text="should not become TextBlock", thought=True, thought_signature=b"s" + ) + + result = part_to_message_block(part) + + assert result["type"] == "thinking" + + +# --- Group B: content_block_to_part with thinking blocks --- + + +def test_content_block_to_part_thinking_block(): + """ThinkingBlock converts to a thought Part with text and thought_signature.""" + from google.adk.models.anthropic_llm import content_block_to_part + + block = anthropic_types.ThinkingBlock( + thinking="my reasoning", signature="base64sig==", type="thinking" + ) + + part = content_block_to_part(block) + + assert part.thought is True + assert part.text == "my reasoning" + assert part.thought_signature == b"base64sig==" + + +def test_content_block_to_part_redacted_thinking_block(): + """RedactedThinkingBlock converts to a thought Part with empty text and thought_signature.""" + from google.adk.models.anthropic_llm import content_block_to_part + + block = anthropic_types.RedactedThinkingBlock( + data="opaque_data", type="redacted_thinking" + ) + + part = content_block_to_part(block) + + assert part.thought is True + assert part.text == "" + assert part.thought_signature == b"opaque_data" + + +# --- Group F: round-trip multi-turn --- + + +def test_thought_part_round_trips_in_content_to_message_param(): + """Thought parts in model content survive the round-trip through content_to_message_param.""" + content = Content( + role="model", + parts=[Part(text="reasoning", thought=True, thought_signature=b"sig")], + ) + + result = content_to_message_param(content) + + assert result["role"] == "assistant" + assert len(result["content"]) == 1 + block = result["content"][0] + assert block["type"] == "thinking" + assert block["thinking"] == "reasoning" + assert block["signature"] == "sig" + + +# --- Group C: generate_content_async thinking param plumbing --- + + +@pytest.mark.asyncio +async def test_generate_content_async_passes_thinking_param(): + """thinking_config is passed as ThinkingConfigEnabledParam to messages.create.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + mock_message = anthropic_types.Message( + id="msg_think", + content=[ + anthropic_types.TextBlock(text="Answer", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(return_value=mock_message) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[ + Content(role="user", parts=[Part.from_text(text="Think hard")]) + ], + config=types.GenerateContentConfig( + system_instruction="You are a reasoner", + thinking_config=types.ThinkingConfig(thinking_budget=2048), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + _ = [r async for r in llm.generate_content_async(llm_req, stream=False)] + + mock_client.messages.create.assert_called_once() + _, kwargs = mock_client.messages.create.call_args + assert kwargs["thinking"]["type"] == "enabled" + assert kwargs["thinking"]["budget_tokens"] == 2048 + + +@pytest.mark.asyncio +async def test_generate_content_async_no_thinking_config_passes_not_given(): + """Without thinking_config, messages.create receives NOT_GIVEN for thinking.""" + from anthropic import NOT_GIVEN + + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + mock_message = anthropic_types.Message( + id="msg_no_think", + content=[ + anthropic_types.TextBlock(text="Hi", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=5, + output_tokens=2, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(return_value=mock_message) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig(system_instruction="Test"), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + _ = [r async for r in llm.generate_content_async(llm_req, stream=False)] + + _, kwargs = mock_client.messages.create.call_args + assert kwargs["thinking"] is NOT_GIVEN + + +@pytest.mark.asyncio +async def test_generate_content_async_thinking_budget_zero_passes_not_given(): + """thinking_budget=0 results in NOT_GIVEN being passed to messages.create.""" + from anthropic import NOT_GIVEN + + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + mock_message = anthropic_types.Message( + id="msg_zero", + content=[ + anthropic_types.TextBlock(text="Hi", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=5, + output_tokens=2, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(return_value=mock_message) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig( + system_instruction="Test", + thinking_config=types.ThinkingConfig(thinking_budget=0), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + _ = [r async for r in llm.generate_content_async(llm_req, stream=False)] + + _, kwargs = mock_client.messages.create.call_args + assert kwargs["thinking"] is NOT_GIVEN + + +@pytest.mark.asyncio +async def test_generate_content_async_thinking_budget_clamped_to_max_tokens(): + """thinking_budget exceeding max_tokens is clamped to max_tokens - 1.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514", max_tokens=4096) + + mock_message = anthropic_types.Message( + id="msg_clamp", + content=[ + anthropic_types.TextBlock(text="Hi", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=5, + output_tokens=2, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock(return_value=mock_message) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig( + system_instruction="Test", + thinking_config=types.ThinkingConfig(thinking_budget=5000), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + _ = [r async for r in llm.generate_content_async(llm_req, stream=False)] + + _, kwargs = mock_client.messages.create.call_args + assert kwargs["thinking"]["budget_tokens"] == 4095 + + +# --- Group D: streaming with thinking blocks --- + + +@pytest.mark.asyncio +async def test_streaming_thinking_block_in_final_response(): + """Thinking block is accumulated from deltas and emitted in the final response.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock(usage=MagicMock(input_tokens=10, output_tokens=0)), + ), + MagicMock( + type="content_block_start", + index=0, + content_block=anthropic_types.ThinkingBlock( + thinking="", signature="", type="thinking" + ), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.ThinkingDelta( + thinking="I think ", type="thinking_delta" + ), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.ThinkingDelta( + thinking="carefully", type="thinking_delta" + ), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.SignatureDelta( + signature="sig123", type="signature_delta" + ), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="content_block_start", + index=1, + content_block=anthropic_types.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=1, + delta=anthropic_types.TextDelta(text="Answer", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=1), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="end_turn"), + usage=MagicMock(output_tokens=20), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Think")])], + config=types.GenerateContentConfig( + system_instruction="You are a reasoner", + thinking_config=types.ThinkingConfig(thinking_budget=1024), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_req, stream=True) + ] + + # 1 partial for the text delta + 1 final + assert len(responses) == 2 + + partial = responses[0] + assert partial.partial is True + assert partial.content.parts[0].text == "Answer" + + final = responses[1] + assert final.partial is False + assert len(final.content.parts) == 2 + + thought_part = final.content.parts[0] + assert thought_part.thought is True + assert thought_part.text == "I think carefully" + assert thought_part.thought_signature == b"sig123" + + text_part = final.content.parts[1] + assert text_part.text == "Answer" + + +@pytest.mark.asyncio +async def test_streaming_redacted_thinking_block_in_final_response(): + """Redacted thinking block (data from content_block_start) appears in final response.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock(usage=MagicMock(input_tokens=10, output_tokens=0)), + ), + MagicMock( + type="content_block_start", + index=0, + content_block=anthropic_types.RedactedThinkingBlock( + data="opaque123", type="redacted_thinking" + ), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="content_block_start", + index=1, + content_block=anthropic_types.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=1, + delta=anthropic_types.TextDelta(text="Done", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=1), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="end_turn"), + usage=MagicMock(output_tokens=5), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig(system_instruction="Test"), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_req, stream=True) + ] + + final = responses[-1] + assert final.partial is False + assert len(final.content.parts) == 2 + + redacted_part = final.content.parts[0] + assert redacted_part.thought is True + assert redacted_part.text == "" + assert redacted_part.thought_signature == b"opaque123" + + assert final.content.parts[1].text == "Done" + + +@pytest.mark.asyncio +async def test_streaming_thinking_does_not_yield_partial(): + """Thinking deltas are not yielded as partial responses.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock(usage=MagicMock(input_tokens=10, output_tokens=0)), + ), + MagicMock( + type="content_block_start", + index=0, + content_block=anthropic_types.ThinkingBlock( + thinking="", signature="", type="thinking" + ), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.ThinkingDelta( + thinking="private thoughts", type="thinking_delta" + ), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.SignatureDelta( + signature="s", type="signature_delta" + ), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="content_block_start", + index=1, + content_block=anthropic_types.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=1, + delta=anthropic_types.TextDelta(text="Reply", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=1), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="end_turn"), + usage=MagicMock(output_tokens=5), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig( + system_instruction="Test", + thinking_config=types.ThinkingConfig(thinking_budget=1024), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_req, stream=True) + ] + + partial_responses = [r for r in responses if r.partial] + assert len(partial_responses) == 1 + assert all( + not getattr(p, "thought", False) + for p in partial_responses[0].content.parts + ) + assert partial_responses[0].content.parts[0].text == "Reply" + + +@pytest.mark.asyncio +async def test_streaming_passes_thinking_param(): + """When thinking_config is set, streaming passes thinking param to messages.create.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock(usage=MagicMock(input_tokens=5, output_tokens=0)), + ), + MagicMock( + type="content_block_start", + index=0, + content_block=anthropic_types.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.TextDelta(text="Ok", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="end_turn"), + usage=MagicMock(output_tokens=1), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_req = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig( + system_instruction="Test", + thinking_config=types.ThinkingConfig(thinking_budget=1024), + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + _ = [r async for r in llm.generate_content_async(llm_req, stream=True)] + + mock_client.messages.create.assert_called_once() + _, kwargs = mock_client.messages.create.call_args + assert kwargs["thinking"]["type"] == "enabled" + assert kwargs["thinking"]["budget_tokens"] == 1024 From a70b50affde922dcf3d7506b2739c0ce6c4853c0 Mon Sep 17 00:00:00 2001 From: Sebastien Coutu Date: Sun, 19 Apr 2026 13:54:20 -0400 Subject: [PATCH 2/3] fix: add e2e tests for anthropic thinking feature --- .../__init__.py | 13 ++++ .../hello_world_anthropic_thinking/agent.py | 61 +++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 contributing/samples/hello_world_anthropic_thinking/__init__.py create mode 100644 contributing/samples/hello_world_anthropic_thinking/agent.py diff --git a/contributing/samples/hello_world_anthropic_thinking/__init__.py b/contributing/samples/hello_world_anthropic_thinking/__init__.py new file mode 100644 index 0000000000..58d482ea38 --- /dev/null +++ b/contributing/samples/hello_world_anthropic_thinking/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/contributing/samples/hello_world_anthropic_thinking/agent.py b/contributing/samples/hello_world_anthropic_thinking/agent.py new file mode 100644 index 0000000000..8b34315a62 --- /dev/null +++ b/contributing/samples/hello_world_anthropic_thinking/agent.py @@ -0,0 +1,61 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from google.adk import Agent +from google.adk.models.anthropic_llm import AnthropicLlm +from google.adk.planners.built_in_planner import BuiltInPlanner +from google.genai import types + + +def check_prime(nums: list[int]) -> str: + """Check if a given list of numbers are prime. + + Args: + nums: The list of numbers to check. + + Returns: + A str indicating which numbers are prime. + """ + primes = set() + for number in nums: + number = int(number) + if number <= 1: + continue + is_prime = True + for i in range(2, int(number**0.5) + 1): + if number % i == 0: + is_prime = False + break + if is_prime: + primes.add(number) + return ( + "No prime numbers found." + if not primes + else f"{', '.join(str(num) for num in sorted(primes))} are prime numbers." + ) + + +root_agent = Agent( + model=AnthropicLlm(model="claude-sonnet-4-6"), + name="anthropic_thinking_agent", + description="An agent that uses Claude extended thinking via Vertex AI.", + instruction=""" + You are a helpful assistant. Use your reasoning carefully before answering. + When asked to check prime numbers, use the check_prime tool. + """, + tools=[check_prime], + planner=BuiltInPlanner( + thinking_config=types.ThinkingConfig(thinking_budget=5000), + ), +) From 566f7de2264d78913b7818322480519dbf50ae03 Mon Sep 17 00:00:00 2001 From: Sebastien Coutu Date: Sun, 19 Apr 2026 13:54:59 -0400 Subject: [PATCH 3/3] chore: run autoformat had more files --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - 2 files changed, 2 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index f3751206a8..2710c3894c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index d857da9635..e31db15788 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string(