You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Add streaming support for Anthropic models
Refactor ToolResultBlockParam content handling to use json.dumps for dict/list results. Implement _generate_content_streaming to handle Anthropic's streaming API Close #3250 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 877613612
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
80c5a24555
commit
5770cd3776
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user