Skip to content

fix(llm): prevent base llm flow truncation (#522) #659

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Conversation

AlankritVerma01
Copy link
Contributor

Fixes: #522

Previously, BaseLlmFlow.run_async would loop indefinitely whenever the LLM response was truncated due to hitting its max_tokens limit. The loop’s break condition only checked for “no events” or last_event.is_final_response(), neither of which fires for a partial/truncated chunk—so the generator re-invoked _run_one_step_async endlessly.

With this change, we:

  • Treat any partial (truncated) event as terminal by adding or last_event.partial to the break condition in run_async.
  • Ensure that when an LLM emits a truncated chunk (partial=True), run_async yields it exactly once and then terminates, avoiding infinite loops.

Test Plan

  • New unit test test_base_llm_flow_truncation.py verifies that:
    • Given a flow that only ever yields one partial event, run_async emits exactly one event and then stops.
  • Run it with:
    pytest tests/unittests/flows/llm_flows/test_base_llm_flow_truncation.py

@AlankritVerma01
Copy link
Contributor Author

Hey @Jacksunwei @hangfei
Please review the PR
Thanks team!

@AlankritVerma01 AlankritVerma01 changed the title Fix/522 base llm flow truncation fix(llm): prevent base llm flow truncation (#522) May 10, 2025
@Jacksunwei Jacksunwei added the ready to pull [Status] This PR is ready to be importing back to Google label May 27, 2025
@Jacksunwei Jacksunwei added this to the FixIt Week milestone May 27, 2025
copybara-service bot pushed a commit that referenced this pull request May 28, 2025
…ted state and can happen when the model output limit is maxed out.

Copybara import of the project:

--
a6e3a22 by Alankrit Verma <[email protected]>:

test(base_llm_flow): add test for infinite loop on truncated responses

--
b5f2245 by Alankrit Verma <[email protected]>:

feat(base_llm_flow): break run_async loop on partial/truncated events

--
cbbae4c by Wei Sun (Jack) <[email protected]>:

Raise ValueError if the last event is partial.

This is invalid state for llm flow.
--
6eebae0 by Wei Sun (Jack) <[email protected]>:

Delete tests/unittests/flows/llm_flows/test_base_llm_flow_truncation.py
--
e08b0ab by Wei Sun (Jack) <[email protected]>:

format base_llm_flow.py

COPYBARA_INTEGRATE_REVIEW=#659 from AlankritVerma01:fix/522-base-llm-flow-truncation e08b0ab
PiperOrigin-RevId: 764380150
raflFaisal pushed a commit to raflFaisal/adk-python that referenced this pull request May 28, 2025
…ted state and can happen when the model output limit is maxed out.

Copybara import of the project:

--
a6e3a22 by Alankrit Verma <[email protected]>:

test(base_llm_flow): add test for infinite loop on truncated responses

--
b5f2245 by Alankrit Verma <[email protected]>:

feat(base_llm_flow): break run_async loop on partial/truncated events

--
cbbae4c by Wei Sun (Jack) <[email protected]>:

Raise ValueError if the last event is partial.

This is invalid state for llm flow.
--
6eebae0 by Wei Sun (Jack) <[email protected]>:

Delete tests/unittests/flows/llm_flows/test_base_llm_flow_truncation.py
--
e08b0ab by Wei Sun (Jack) <[email protected]>:

format base_llm_flow.py

COPYBARA_INTEGRATE_REVIEW=google#659 from AlankritVerma01:fix/522-base-llm-flow-truncation e08b0ab
PiperOrigin-RevId: 764380150
@Jacksunwei
Copy link
Collaborator

Merged in 2a07eef

@Jacksunwei Jacksunwei closed this May 28, 2025
whhlct pushed a commit to whhlct/adk-python that referenced this pull request Jun 6, 2025
…ted state and can happen when the model output limit is maxed out.

Copybara import of the project:

--
a6e3a22 by Alankrit Verma <[email protected]>:

test(base_llm_flow): add test for infinite loop on truncated responses

--
b5f2245 by Alankrit Verma <[email protected]>:

feat(base_llm_flow): break run_async loop on partial/truncated events

--
cbbae4c by Wei Sun (Jack) <[email protected]>:

Raise ValueError if the last event is partial.

This is invalid state for llm flow.
--
6eebae0 by Wei Sun (Jack) <[email protected]>:

Delete tests/unittests/flows/llm_flows/test_base_llm_flow_truncation.py
--
e08b0ab by Wei Sun (Jack) <[email protected]>:

format base_llm_flow.py

COPYBARA_INTEGRATE_REVIEW=google#659 from AlankritVerma01:fix/522-base-llm-flow-truncation e08b0ab
PiperOrigin-RevId: 764380150
coder-aditi pushed a commit to coder-aditi/adk-python that referenced this pull request Jul 2, 2025
…ted state and can happen when the model output limit is maxed out.

Copybara import of the project:

--
a6e3a22 by Alankrit Verma <[email protected]>:

test(base_llm_flow): add test for infinite loop on truncated responses

--
b5f2245 by Alankrit Verma <[email protected]>:

feat(base_llm_flow): break run_async loop on partial/truncated events

--
cbbae4c by Wei Sun (Jack) <[email protected]>:

Raise ValueError if the last event is partial.

This is invalid state for llm flow.
--
6eebae0 by Wei Sun (Jack) <[email protected]>:

Delete tests/unittests/flows/llm_flows/test_base_llm_flow_truncation.py
--
e08b0ab by Wei Sun (Jack) <[email protected]>:

format base_llm_flow.py

COPYBARA_INTEGRATE_REVIEW=google#659 from AlankritVerma01:fix/522-base-llm-flow-truncation e08b0ab
PiperOrigin-RevId: 764380150
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ready to pull [Status] This PR is ready to be importing back to Google
Projects
None yet
Development

Successfully merging this pull request may close these issues.

base_llm_flow.py has running problem
2 participants