github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 9 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

perf: avoid quadratic streaming accumulation in the LiteLLM adapter

Browse Source 
Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 957325226
This commit is contained in:
George Weale
2026-07-31 14:07:12 -07:00
committed by Copybara-Service
parent e182146f95
commit 36fd2c8e0c
2 changed files with 111 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/models/lite_llm.py
+20 -14
View File
@@ -2829,12 +2829,15 @@ class LiteLlm(BaseLlm):
completion_args["extra_body"] = http_opts.extra_body
if stream:
text = ""
# Accumulate into lists and join once: `+=` on a closure cell or a dict
# item does not get CPython's in-place unicode concat, so it would copy
# the whole buffer on every streamed chunk.
text_parts: list[str] = []
reasoning_parts: List[types.Part] = []
# Track function calls by index
function_calls: dict[int, dict[str, Any]] = (
{}
) # index -> {name, args, id}
) # index -> {name, args_parts, id}
tool_call_trackers: Dict[int, _BraceDepthTracker] = {}
completion_args["stream"] = True
completion_args["stream_options"] = {"include_usage": True}
@@ -2851,9 +2854,10 @@ class LiteLlm(BaseLlm):
has_incomplete_tool_call_args = False
for index, func_data in function_calls.items():
if func_data["id"]:
args = "".join(func_data["args_parts"])
if finish_reason == "length":
try:
_parse_tool_call_arguments(func_data["args"])
_parse_tool_call_arguments(args)
except json.JSONDecodeError:
has_incomplete_tool_call_args = True
continue
@@ -2863,7 +2867,7 @@ class LiteLlm(BaseLlm):
id=func_data["id"],
function=Function(
name=func_data["name"],
arguments=func_data["args"],
arguments=args,
index=index,
),
)
@@ -2884,7 +2888,7 @@ class LiteLlm(BaseLlm):
llm_response = _message_to_generate_content_response(
ChatCompletionAssistantMessage(
role="assistant",
content=text,
content="".join(text_parts),
tool_calls=tool_calls,
),
model_version=model_version,
@@ -2902,7 +2906,7 @@ class LiteLlm(BaseLlm):
def _finalize_text_response(
*, model_version: str, finish_reason: str
) -> LlmResponse:
message_content = text if text else None
message_content = "".join(text_parts) or None
llm_response = _message_to_generate_content_response(
ChatCompletionAssistantMessage(
role="assistant",
@@ -2921,8 +2925,8 @@ class LiteLlm(BaseLlm):
return llm_response
def _reset_stream_buffers() -> None:
nonlocal text, reasoning_parts
text = ""
nonlocal reasoning_parts
text_parts.clear()
reasoning_parts = []
function_calls.clear()
tool_call_trackers.clear()
@@ -2937,12 +2941,13 @@ class LiteLlm(BaseLlm):
if isinstance(chunk, FunctionChunk):
index = chunk.index or fallback_index
if index not in function_calls:
function_calls[index] = {"name": "", "args": "", "id": None}
function_calls[index] = {"name": "", "args_parts": [], "id": None}
if chunk.name:
function_calls[index]["name"] += chunk.name
if chunk.args:
function_calls[index]["args"] += chunk.args
args_parts = function_calls[index]["args_parts"]
args_parts.append(chunk.args)
# Detect args completion to advance fallback_index (workaround
# for improper chunk indexing) without O(N^2) re-parsing.
@@ -2951,7 +2956,7 @@ class LiteLlm(BaseLlm):
)
if tracker.feed(chunk.args):
try:
json.loads(function_calls[index]["args"])
json.loads("".join(args_parts))
fallback_index += 1
except json.JSONDecodeError:
pass
@@ -2960,7 +2965,8 @@ class LiteLlm(BaseLlm):
chunk.id or function_calls[index]["id"] or str(index)
)
elif isinstance(chunk, TextChunk):
text += chunk.text
if chunk.text:
text_parts.append(chunk.text)
yield _message_to_generate_content_response(
ChatCompletionAssistantMessage(
role="assistant",
@@ -3003,7 +3009,7 @@ class LiteLlm(BaseLlm):
)
)
_reset_stream_buffers()
elif (text or reasoning_parts) and (
elif (text_parts or reasoning_parts) and (
finish_reason == "length"
or (
finish_reason == "stop"
@@ -3024,7 +3030,7 @@ class LiteLlm(BaseLlm):
)
_reset_stream_buffers()
if (text or reasoning_parts) and not aggregated_llm_response:
if (text_parts or reasoning_parts) and not aggregated_llm_response:
aggregated_llm_response = _finalize_text_response(
model_version=part.model,
finish_reason="stop",
tests/unittests/models/test_litellm.py
+91
View File
@@ -6452,6 +6452,97 @@ async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete(
assert args_by_name["other_func"] == json.loads(full_args_b)
def _text_stream_chunks(text_fragments, finish_reason="stop"):
stream = [
ModelResponseStream(
choices=[
StreamingChoices(
finish_reason=None,
delta=Delta(role="assistant", content=fragment),
)
]
)
for fragment in text_fragments
]
stream.append(
ModelResponseStream(
choices=[StreamingChoices(finish_reason=finish_reason, delta=Delta())]
)
)
return stream
@pytest.mark.asyncio
async def test_streaming_text_assembled_from_many_fragments(
mock_completion, lite_llm_instance
):
full_text = "".join(f"token-{i} " for i in range(500))
fragments = _split_into_chunks(full_text, [7] * (len(full_text) // 7))
mock_completion.return_value = iter(_text_stream_chunks(fragments))
responses = [
r
async for r in lite_llm_instance.generate_content_async(
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
)
]
partials = [r for r in responses if r.partial]
aggregated = [r for r in responses if not r.partial]
assert [p.content.parts[0].text for p in partials] == fragments
assert len(aggregated) == 1
assert aggregated[0].content.parts[0].text == full_text
@pytest.mark.asyncio
async def test_streaming_buffers_hold_fragments_instead_of_growing_copies(
mock_completion, lite_llm_instance
):
# `+=` onto a closure cell or a dict item does not get CPython's in-place
# unicode concat, so it re-copies the whole buffer on every chunk and makes
# a stream quadratic in its own length. Both buffers must stay lists of the
# raw fragments, so each chunk costs only its own length.
arg_fragments = ['{"a": ', "1, ", '"b": 2}']
text_fragments = ["alpha ", "beta"]
stream = _stream_chunks_from_function_chunks(
_function_chunks_for_args(arg_fragments)
)[:-1]
stream.extend(_text_stream_chunks(text_fragments)[:-1])
mock_completion.return_value = iter(stream)
responses = lite_llm_instance.generate_content_async(
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
)
try:
# Suspends on the first partial text response, with both buffers filled.
await responses.__anext__()
buffers = responses.ag_frame.f_locals
assert buffers["text_parts"] == text_fragments[:1]
assert buffers["function_calls"][0]["args_parts"] == arg_fragments
finally:
await responses.aclose()
@pytest.mark.asyncio
async def test_streaming_text_buffer_is_reset_between_aggregated_responses(
mock_completion, lite_llm_instance
):
stream = _text_stream_chunks(["first "])
stream.extend(_text_stream_chunks(["second"]))
mock_completion.return_value = iter(stream)
responses = [
r
async for r in lite_llm_instance.generate_content_async(
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
)
]
aggregated = [r for r in responses if not r.partial]
assert len(aggregated) == 1
assert aggregated[0].content.parts[0].text == "second"
def test_model_dump_json_excludes_llm_client():
lite_llm_model = LiteLlm(model="test_model")