fix: cache canonical tools to avoid multiple calls when streaming

[hcadioli]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

• Closes: Getting increase latency with MCPToolset after version 1.16.0 #3237
• Related: Getting increase latency with MCPToolset after version 1.16.0 #3237

Testing Plan

Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes.

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

Please include a summary of passed pytest results.

Manual End-to-End (E2E) Tests:

Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules.

Additional context

Add any other context or screenshots about the feature request here.

[gemini-code-assist]

Summary of Changes

Hello @hcadioli, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a performance issue by implementing a caching strategy for canonical tools. By storing the result of agent.canonical_tools in the InvocationContext, the system avoids repeated, potentially expensive, calls to retrieve these tools, especially in streaming contexts. This change aims to improve the overall efficiency and responsiveness of the application by reducing redundant processing.

Highlights

• Performance Optimization: Introduced a caching mechanism for canonical tools within the InvocationContext to prevent redundant computations, especially during streaming operations.
• Reduced Redundant Calls: Modified the _maybe_add_grounding_metadata function to utilize the new canonical_tools_cache, ensuring agent.canonical_tools is not called multiple times.
• Data Structure Enhancement: Added a new canonical_tools_cache field of type Optional[list[BaseTool]] to the InvocationContext class for storing cached tools.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[adk-bot]

Response from ADK Triaging Agent

Hello @hcadioli, thank you for creating this PR!

This PR is missing a link to a Github issue. Could you please associate an issue with this PR? If there is no existing issue, could you please create one?

In addition, could you please provide a testing plan?

This information will help reviewers to review your PR more efficiently. Thanks!

[gemini-code-assist]

Code Review

This pull request aims to improve performance by caching canonical tools to avoid multiple calls during streaming. A new field canonical_tools_cache is added to InvocationContext for this purpose.

My review identifies a critical issue with the caching implementation: the cache is read from but never populated, rendering the optimization ineffective. I've provided a code suggestion to correct this logic. Additionally, I've suggested a minor style improvement for import consistency. Addressing the caching logic is essential for this fix to work as intended.

[gemini-code-assist]

Code Review

This pull request introduces a caching mechanism for canonical tools to optimize performance during streaming by avoiding redundant computations. The implementation correctly adds a cache field to InvocationContext and utilizes it within BaseLlmFlow. My review includes a suggestion to make the caching logic more concise and Pythonic, which will improve code readability and maintainability.

[gemini-code-assist]

Code Review

This pull request introduces a caching mechanism for canonical tools within the invocation context. This is a good optimization to avoid redundant computations, especially in streaming scenarios. The implementation is sound, and I have one suggestion to make the caching logic slightly more concise and readable.

[gemini-code-assist]

Code Review

This pull request introduces a caching mechanism for canonical tools within the InvocationContext to prevent redundant computations during streaming operations. The implementation is sound and correctly utilizes the cache in _maybe_add_grounding_metadata. My primary feedback is the absence of unit tests to validate this new caching logic, which is important for ensuring the fix is effective and preventing future regressions.

[gemini-code-assist]

Code Review

This pull request introduces a caching mechanism for canonical tools within the InvocationContext to avoid redundant calls to agent.canonical_tools when streaming. This optimization is implemented in src/google/adk/flows/llm_flows/base_llm_flow.py and is accompanied by a unit test in tests/unittests/flows/llm_flows/test_base_llm_flow.py to ensure the caching mechanism works as expected. The addition of the canonical_tools_cache field in src/google/adk/agents/invocation_context.py provides a place to store the cached tools.