feat: Allow toolset to process llm_request before tools returned by it

PiperOrigin-RevId: 785480813

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -42,6 +42,7 @@ from ...models.llm_response import LlmResponse
 from ...telemetry import trace_call_llm
 from ...telemetry import trace_send_data
 from ...telemetry import tracer
+from ...tools.base_toolset import BaseToolset
 from ...tools.tool_context import ToolContext
 
 if TYPE_CHECKING:
@@ -341,13 +342,25 @@ class BaseLlmFlow(ABC):
         yield event
 
     # Run processors for tools.
-    for tool in await agent.canonical_tools(
-        ReadonlyContext(invocation_context)
-    ):
+    for tool_union in agent.tools:
       tool_context = ToolContext(invocation_context)
-      await tool.process_llm_request(
-          tool_context=tool_context, llm_request=llm_request
+
+      # If it's a toolset, process it first
+      if isinstance(tool_union, BaseToolset):
+        await tool_union.process_llm_request(
+            tool_context=tool_context, llm_request=llm_request
+        )
+
+      from ...agents.llm_agent import _convert_tool_union_to_tools
+
+      # Then process all tools from this tool union
+      tools = await _convert_tool_union_to_tools(
+          tool_union, ReadonlyContext(invocation_context)
       )
+      for tool in tools:
+        await tool.process_llm_request(
+            tool_context=tool_context, llm_request=llm_request
+        )
 
   async def _postprocess_async(
       self,

src/google/adk/tools/base_toolset.py

@@ -20,11 +20,16 @@ from typing import List
 from typing import Optional
 from typing import Protocol
 from typing import runtime_checkable
+from typing import TYPE_CHECKING
 from typing import Union
 
 from ..agents.readonly_context import ReadonlyContext
 from .base_tool import BaseTool
 
+if TYPE_CHECKING:
+  from ..models.llm_request import LlmRequest
+  from .tool_context import ToolContext
+
 
 @runtime_checkable
 class ToolPredicate(Protocol):
@@ -96,3 +101,20 @@ class BaseToolset(ABC):
       return tool.name in self.tool_filter
 
     return False
+
+  async def process_llm_request(
+      self, *, tool_context: ToolContext, llm_request: LlmRequest
+  ) -> None:
+    """Processes the outgoing LLM request for this toolset. This method will be
+    called before each tool processes the llm request.
+
+    Use cases:
+    - Instead of let each tool process the llm request, we can let the toolset
+      process the llm request. e.g. ComputerUseToolset can add computer use
+      tool to the llm request.
+
+    Args:
+      tool_context: The context of the tool.
+      llm_request: The outgoing LLM request, mutable this method.
+    """
+    pass