diff --git a/src/google/adk/models/anthropic_llm.py b/src/google/adk/models/anthropic_llm.py index 97992096..1f7f37b0 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -17,7 +17,9 @@ from __future__ import annotations import base64 +import dataclasses from functools import cached_property +import json import logging import os from typing import Any @@ -31,6 +33,7 @@ from typing import Union from anthropic import AsyncAnthropic from anthropic import AsyncAnthropicVertex from anthropic import NOT_GIVEN +from anthropic import NotGiven from anthropic import types as anthropic_types from google.genai import types from pydantic import BaseModel @@ -48,6 +51,15 @@ __all__ = ["AnthropicLlm", "Claude"] logger = logging.getLogger("google_adk." + __name__) +@dataclasses.dataclass +class _ToolUseAccumulator: + """Accumulates streamed tool_use content block data.""" + + id: str + name: str + args_json: str + + class ClaudeRequest(BaseModel): system_instruction: str messages: Iterable[anthropic_types.MessageParam] @@ -115,12 +127,15 @@ def part_to_message_block( else: content_items.append(str(item)) content = "\n".join(content_items) if content_items else "" - # Handle traditional result format - elif "result" in response_data and response_data["result"]: - # Transformation is required because the content is a list of dict. - # ToolResultBlockParam content doesn't support list of dict. Converting - # to str to prevent anthropic.BadRequestError from being thrown. - content = str(response_data["result"]) + # We serialize to str here + # SDK ref: anthropic.types.tool_result_block_param + # https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/types/tool_result_block_param.py + elif "result" in response_data and response_data["result"] is not None: + result = response_data["result"] + if isinstance(result, (dict, list)): + content = json.dumps(result) + else: + content = str(result) return anthropic_types.ToolResultBlockParam( tool_use_id=part.function_response.id or "", @@ -305,16 +320,111 @@ class AnthropicLlm(BaseLlm): if llm_request.tools_dict else NOT_GIVEN ) - # TODO(b/421255973): Enable streaming for anthropic models. - message = await self._anthropic_client.messages.create( + + if not stream: + message = await self._anthropic_client.messages.create( + model=llm_request.model, + system=llm_request.config.system_instruction, + messages=messages, + tools=tools, + tool_choice=tool_choice, + max_tokens=self.max_tokens, + ) + yield message_to_generate_content_response(message) + else: + async for response in self._generate_content_streaming( + llm_request, messages, tools, tool_choice + ): + yield response + + async def _generate_content_streaming( + self, + llm_request: LlmRequest, + messages: list[anthropic_types.MessageParam], + tools: Union[Iterable[anthropic_types.ToolUnionParam], NotGiven], + tool_choice: Union[anthropic_types.ToolChoiceParam, NotGiven], + ) -> AsyncGenerator[LlmResponse, None]: + """Handles streaming responses from Anthropic models. + + Yields partial LlmResponse objects as content arrives, followed by + a final aggregated LlmResponse with all content. + """ + raw_stream = await self._anthropic_client.messages.create( model=llm_request.model, system=llm_request.config.system_instruction, messages=messages, tools=tools, tool_choice=tool_choice, max_tokens=self.max_tokens, + stream=True, + ) + + # 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] = {} + input_tokens = 0 + output_tokens = 0 + + async for event in raw_stream: + if event.type == "message_start": + input_tokens = event.message.usage.input_tokens + output_tokens = event.message.usage.output_tokens + + elif event.type == "content_block_start": + block = event.content_block + if isinstance(block, anthropic_types.TextBlock): + text_blocks[event.index] = block.text + elif isinstance(block, anthropic_types.ToolUseBlock): + tool_use_blocks[event.index] = _ToolUseAccumulator( + id=block.id, + name=block.name, + args_json="", + ) + + elif event.type == "content_block_delta": + delta = event.delta + if isinstance(delta, anthropic_types.TextDelta): + text_blocks.setdefault(event.index, "") + text_blocks[event.index] += delta.text + yield LlmResponse( + content=types.Content( + role="model", + parts=[types.Part.from_text(text=delta.text)], + ), + partial=True, + ) + elif isinstance(delta, anthropic_types.InputJSONDelta): + if event.index in tool_use_blocks: + tool_use_blocks[event.index].args_json += delta.partial_json + + elif event.type == "message_delta": + output_tokens = event.usage.output_tokens + + # 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())) + ) + for idx in all_indices: + if idx in text_blocks: + all_parts.append(types.Part.from_text(text=text_blocks[idx])) + if idx in tool_use_blocks: + acc = tool_use_blocks[idx] + args = json.loads(acc.args_json) if acc.args_json else {} + part = types.Part.from_function_call(name=acc.name, args=args) + part.function_call.id = acc.id + all_parts.append(part) + + yield LlmResponse( + content=types.Content(role="model", parts=all_parts), + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=input_tokens, + candidates_token_count=output_tokens, + total_token_count=input_tokens + output_tokens, + ), + partial=False, ) - yield message_to_generate_content_response(message) @cached_property def _anthropic_client(self) -> AsyncAnthropic: diff --git a/tests/unittests/models/test_anthropic_llm.py b/tests/unittests/models/test_anthropic_llm.py index fac5f462..50759659 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import sys from unittest import mock +from unittest.mock import AsyncMock +from unittest.mock import MagicMock from anthropic import types as anthropic_types from google.adk import version as adk_version @@ -23,6 +26,7 @@ from google.adk.models.anthropic_llm import AnthropicLlm from google.adk.models.anthropic_llm import Claude from google.adk.models.anthropic_llm import content_to_message_param from google.adk.models.anthropic_llm import function_declaration_to_tool_param +from google.adk.models.anthropic_llm import part_to_message_block from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.genai import types @@ -598,3 +602,354 @@ def test_content_to_message_param_with_images( ) else: mock_logger.warning.assert_not_called() + + +# --- Tests for Bug #2: json.dumps for dict/list function results --- + + +def test_part_to_message_block_dict_result_serialized_as_json(): + """Dict results should be serialized with json.dumps, not str().""" + response_part = types.Part.from_function_response( + name="get_topic", + response={"result": {"topic": "travel", "active": True, "count": None}}, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + content = result["content"] + + # Must be valid JSON (json.dumps produces "true"/"null", not "True"/"None") + parsed = json.loads(content) + assert parsed["topic"] == "travel" + assert parsed["active"] is True + assert parsed["count"] is None + + +def test_part_to_message_block_list_result_serialized_as_json(): + """List results should be serialized with json.dumps.""" + response_part = types.Part.from_function_response( + name="get_items", + response={"result": ["item1", "item2", "item3"]}, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + content = result["content"] + + parsed = json.loads(content) + assert parsed == ["item1", "item2", "item3"] + + +def test_part_to_message_block_empty_dict_result_not_dropped(): + """Empty dict results should produce '{}', not empty string.""" + response_part = types.Part.from_function_response( + name="some_tool", + response={"result": {}}, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + assert result["content"] == "{}" + + +def test_part_to_message_block_empty_list_result_not_dropped(): + """Empty list results should produce '[]', not empty string.""" + response_part = types.Part.from_function_response( + name="some_tool", + response={"result": []}, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + assert result["content"] == "[]" + + +def test_part_to_message_block_string_result_unchanged(): + """String results should still work as before (backward compat).""" + response_part = types.Part.from_function_response( + name="simple_tool", + response={"result": "plain text result"}, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + assert result["content"] == "plain text result" + + +def test_part_to_message_block_nested_dict_result(): + """Nested dict with arrays should produce valid JSON.""" + response_part = types.Part.from_function_response( + name="search", + response={ + "result": { + "results": [ + {"id": 1, "tags": ["a", "b"]}, + {"id": 2, "meta": {"key": "val"}}, + ], + "has_more": False, + } + }, + ) + response_part.function_response.id = "test_id" + + result = part_to_message_block(response_part) + parsed = json.loads(result["content"]) + assert parsed["has_more"] is False + assert parsed["results"][0]["tags"] == ["a", "b"] + + +# --- Tests for Bug #1: Streaming support --- + + +def _make_mock_stream_events(events): + """Helper to create an async iterable from a list of events.""" + + async def _stream(): + for event in events: + yield event + + return _stream() + + +@pytest.mark.asyncio +async def test_streaming_text_yields_partial_and_final(): + """Streaming text should yield partial chunks then a 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.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.TextDelta(text="Hello ", type="text_delta"), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.TextDelta(text="world!", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=0), + 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_request = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + config=types.GenerateContentConfig( + system_instruction="You are helpful", + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_request, stream=True) + ] + + # 2 partial text chunks + 1 final aggregated + assert len(responses) == 3 + assert responses[0].partial is True + assert responses[0].content.parts[0].text == "Hello " + assert responses[1].partial is True + assert responses[1].content.parts[0].text == "world!" + assert responses[2].partial is False + assert responses[2].content.parts[0].text == "Hello world!" + assert responses[2].usage_metadata.prompt_token_count == 10 + assert responses[2].usage_metadata.candidates_token_count == 5 + + +@pytest.mark.asyncio +async def test_streaming_tool_use_yields_function_call(): + """Streaming tool_use should accumulate args and yield in final.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock(usage=MagicMock(input_tokens=20, 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="Checking.", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="content_block_start", + index=1, + content_block=anthropic_types.ToolUseBlock( + id="toolu_abc", + name="get_weather", + input={}, + type="tool_use", + ), + ), + MagicMock( + type="content_block_delta", + index=1, + delta=anthropic_types.InputJSONDelta( + partial_json='{"city": "Paris"}', + type="input_json_delta", + ), + ), + MagicMock(type="content_block_stop", index=1), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="tool_use"), + usage=MagicMock(output_tokens=12), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_request = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[ + Content( + role="user", + parts=[Part.from_text(text="Weather?")], + ) + ], + config=types.GenerateContentConfig( + system_instruction="You are helpful", + ), + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_request, stream=True) + ] + + # 1 text partial + 1 final + assert len(responses) == 2 + + final = responses[-1] + assert final.partial is False + assert len(final.content.parts) == 2 + assert final.content.parts[0].text == "Checking." + assert final.content.parts[1].function_call.name == "get_weather" + assert final.content.parts[1].function_call.args == {"city": "Paris"} + assert final.content.parts[1].function_call.id == "toolu_abc" + + +@pytest.mark.asyncio +async def test_streaming_passes_stream_true_to_create(): + """When stream=True, messages.create should be called with stream=True.""" + 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="Hi", 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_request = 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_request, stream=True)] + + mock_client.messages.create.assert_called_once() + _, kwargs = mock_client.messages.create.call_args + assert kwargs["stream"] is True + + +@pytest.mark.asyncio +async def test_non_streaming_does_not_pass_stream_param(): + """When stream=False, messages.create should NOT get stream param.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + mock_message = anthropic_types.Message( + id="msg_test", + content=[ + anthropic_types.TextBlock(text="Hello!", 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_request = 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_request, stream=False) + ] + + assert len(responses) == 1 + mock_client.messages.create.assert_called_once() + _, kwargs = mock_client.messages.create.call_args + assert "stream" not in kwargs