google-adk 1.7.0__py3-none-any.whl → 1.8.0__py3-none-any.whl

google/adk/cli/fast_api.py

@@ -20,18 +20,21 @@ import json
 import logging
 import os
 from pathlib import Path
+import shutil
 import time
 import traceback
 import typing
 from typing import Any
 from typing import List
 from typing import Literal
+from typing import Mapping
 from typing import Optional
 
 import click
 from fastapi import FastAPI
 from fastapi import HTTPException
 from fastapi import Query
+from fastapi import UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse
 from fastapi.responses import StreamingResponse
@@ -41,7 +44,6 @@ from fastapi.websockets import WebSocketDisconnect
 from google.genai import types
 import graphviz
 from opentelemetry import trace
-from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
 from opentelemetry.sdk.trace import export
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace import TracerProvider
@@ -51,7 +53,6 @@ from starlette.types import Lifespan
 from typing_extensions import override
 from watchdog.events import FileSystemEventHandler
 from watchdog.observers import Observer
-import yaml
 
 from ..agents import RunConfig
 from ..agents.live_request_queue import LiveRequest
@@ -61,6 +62,9 @@ from ..artifacts.gcs_artifact_service import GcsArtifactService
 from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
 from ..auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from ..errors.not_found_error import NotFoundError
+from ..evaluation.base_eval_service import InferenceConfig
+from ..evaluation.base_eval_service import InferenceRequest
+from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
 from ..evaluation.eval_case import EvalCase
 from ..evaluation.eval_case import SessionInput
 from ..evaluation.eval_metrics import EvalMetric
@@ -72,9 +76,7 @@ from ..evaluation.local_eval_sets_manager import LocalEvalSetsManager
 from ..events.event import Event
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..memory.vertex_ai_memory_bank_service import VertexAiMemoryBankService
-from ..memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService
 from ..runners import Runner
-from ..sessions.database_session_service import DatabaseSessionService
 from ..sessions.in_memory_session_service import InMemorySessionService
 from ..sessions.session import Session
 from ..sessions.vertex_ai_session_service import VertexAiSessionService
@@ -198,6 +200,7 @@ class RunEvalResult(common.BaseModel):
   final_eval_status: EvalStatus
   eval_metric_results: list[tuple[EvalMetric, EvalMetricResult]] = Field(
       deprecated=True,
+      default=[],
       description=(
           "This field is deprecated, use overall_eval_metric_results instead."
       ),
@@ -212,18 +215,11 @@ class GetEventGraphResult(common.BaseModel):
   dot_src: str
 
 
-class AgentBuildRequest(common.BaseModel):
-  agent_name: str
-  agent_type: str
-  model: str
-  description: str
-  instruction: str
-
-
 def get_fast_api_app(
     *,
     agents_dir: str,
     session_service_uri: Optional[str] = None,
+    session_db_kwargs: Optional[Mapping[str, Any]] = None,
     artifact_service_uri: Optional[str] = None,
     memory_service_uri: Optional[str] = None,
     eval_storage_uri: Optional[str] = None,
@@ -248,6 +244,8 @@ def get_fast_api_app(
   memory_exporter = InMemoryExporter(session_trace_dict)
   provider.add_span_processor(export.SimpleSpanProcessor(memory_exporter))
   if trace_to_cloud:
+    from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
+
     envs.load_dotenv_for_agent("", agents_dir)
     if project_id := os.environ.get("GOOGLE_CLOUD_PROJECT", None):
       processor = export.BatchSpanProcessor(
@@ -304,9 +302,36 @@ def get_fast_api_app(
     eval_sets_manager = LocalEvalSetsManager(agents_dir=agents_dir)
     eval_set_results_manager = LocalEvalSetResultsManager(agents_dir=agents_dir)
 
+  def _parse_agent_engine_resource_name(agent_engine_id_or_resource_name):
+    if not agent_engine_id_or_resource_name:
+      raise click.ClickException(
+          "Agent engine resource name or resource id can not be empty."
+      )
+
+    # "projects/my-project/locations/us-central1/reasoningEngines/1234567890",
+    if "/" in agent_engine_id_or_resource_name:
+      # Validate resource name.
+      if len(agent_engine_id_or_resource_name.split("/")) != 6:
+        raise click.ClickException(
+            "Agent engine resource name is mal-formatted. It should be of"
+            " format :"
+            " projects/{project_id}/locations/{location}/reasoningEngines/{resource_id}"
+        )
+      project = agent_engine_id_or_resource_name.split("/")[1]
+      location = agent_engine_id_or_resource_name.split("/")[3]
+      agent_engine_id = agent_engine_id_or_resource_name.split("/")[-1]
+    else:
+      envs.load_dotenv_for_agent("", agents_dir)
+      project = os.environ["GOOGLE_CLOUD_PROJECT"]
+      location = os.environ["GOOGLE_CLOUD_LOCATION"]
+      agent_engine_id = agent_engine_id_or_resource_name
+    return project, location, agent_engine_id
+
   # Build the Memory service
   if memory_service_uri:
     if memory_service_uri.startswith("rag://"):
+      from ..memory.vertex_ai_rag_memory_service import VertexAiRagMemoryService
+
       rag_corpus = memory_service_uri.split("://")[1]
       if not rag_corpus:
         raise click.ClickException("Rag corpus can not be empty.")
@@ -315,13 +340,13 @@ def get_fast_api_app(
           rag_corpus=f'projects/{os.environ["GOOGLE_CLOUD_PROJECT"]}/locations/{os.environ["GOOGLE_CLOUD_LOCATION"]}/ragCorpora/{rag_corpus}'
       )
     elif memory_service_uri.startswith("agentengine://"):
-      agent_engine_id = memory_service_uri.split("://")[1]
-      if not agent_engine_id:
-        raise click.ClickException("Agent engine id can not be empty.")
-      envs.load_dotenv_for_agent("", agents_dir)
+      agent_engine_id_or_resource_name = memory_service_uri.split("://")[1]
+      project, location, agent_engine_id = _parse_agent_engine_resource_name(
+          agent_engine_id_or_resource_name
+      )
       memory_service = VertexAiMemoryBankService(
-          project=os.environ["GOOGLE_CLOUD_PROJECT"],
-          location=os.environ["GOOGLE_CLOUD_LOCATION"],
+          project=project,
+          location=location,
           agent_engine_id=agent_engine_id,
       )
     else:
@@ -334,18 +359,24 @@ def get_fast_api_app(
   # Build the Session service
   if session_service_uri:
     if session_service_uri.startswith("agentengine://"):
-      # Create vertex session service
-      agent_engine_id = session_service_uri.split("://")[1]
-      if not agent_engine_id:
-        raise click.ClickException("Agent engine id can not be empty.")
-      envs.load_dotenv_for_agent("", agents_dir)
+      agent_engine_id_or_resource_name = session_service_uri.split("://")[1]
+      project, location, agent_engine_id = _parse_agent_engine_resource_name(
+          agent_engine_id_or_resource_name
+      )
       session_service = VertexAiSessionService(
-          project=os.environ["GOOGLE_CLOUD_PROJECT"],
-          location=os.environ["GOOGLE_CLOUD_LOCATION"],
+          project=project,
+          location=location,
           agent_engine_id=agent_engine_id,
       )
     else:
-      session_service = DatabaseSessionService(db_url=session_service_uri)
+      from ..sessions.database_session_service import DatabaseSessionService
+
+      # Database session additional settings
+      if session_db_kwargs is None:
+        session_db_kwargs = {}
+      session_service = DatabaseSessionService(
+          db_url=session_service_uri, **session_db_kwargs
+      )
   else:
     session_service = InMemorySessionService()
 
@@ -644,63 +675,66 @@ def get_fast_api_app(
       app_name: str, eval_set_id: str, req: RunEvalRequest
   ) -> list[RunEvalResult]:
     """Runs an eval given the details in the eval request."""
-    from .cli_eval import run_evals
-
     # Create a mapping from eval set file to all the evals that needed to be
     # run.
-    eval_set = eval_sets_manager.get_eval_set(app_name, eval_set_id)
+    try:
+      from ..evaluation.local_eval_service import LocalEvalService
+      from .cli_eval import _collect_eval_results
+      from .cli_eval import _collect_inferences
 
-    if not eval_set:
-      raise HTTPException(
-          status_code=400, detail=f"Eval set `{eval_set_id}` not found."
-      )
+      eval_set = eval_sets_manager.get_eval_set(app_name, eval_set_id)
 
-    if req.eval_ids:
-      eval_cases = [e for e in eval_set.eval_cases if e.eval_id in req.eval_ids]
-      eval_set_to_evals = {eval_set_id: eval_cases}
-    else:
-      logger.info("Eval ids to run list is empty. We will run all eval cases.")
-      eval_set_to_evals = {eval_set_id: eval_set.eval_cases}
+      if not eval_set:
+        raise HTTPException(
+            status_code=400, detail=f"Eval set `{eval_set_id}` not found."
+        )
 
-    root_agent = agent_loader.load_agent(app_name)
-    run_eval_results = []
-    eval_case_results = []
-    try:
-      async for eval_case_result in run_evals(
-          eval_set_to_evals,
-          root_agent,
-          getattr(root_agent, "reset_data", None),
-          req.eval_metrics,
+      root_agent = agent_loader.load_agent(app_name)
+
+      eval_case_results = []
+
+      eval_service = LocalEvalService(
+          root_agent=root_agent,
+          eval_sets_manager=eval_sets_manager,
+          eval_set_results_manager=eval_set_results_manager,
           session_service=session_service,
           artifact_service=artifact_service,
-      ):
-        run_eval_results.append(
-            RunEvalResult(
-                app_name=app_name,
-                eval_set_file=eval_case_result.eval_set_file,
-                eval_set_id=eval_set_id,
-                eval_id=eval_case_result.eval_id,
-                final_eval_status=eval_case_result.final_eval_status,
-                eval_metric_results=eval_case_result.eval_metric_results,
-                overall_eval_metric_results=eval_case_result.overall_eval_metric_results,
-                eval_metric_result_per_invocation=eval_case_result.eval_metric_result_per_invocation,
-                user_id=eval_case_result.user_id,
-                session_id=eval_case_result.session_id,
-            )
-        )
-        eval_case_result.session_details = await session_service.get_session(
-            app_name=app_name,
-            user_id=eval_case_result.user_id,
-            session_id=eval_case_result.session_id,
-        )
-        eval_case_results.append(eval_case_result)
+      )
+      inference_request = InferenceRequest(
+          app_name=app_name,
+          eval_set_id=eval_set.eval_set_id,
+          eval_case_ids=req.eval_ids,
+          inference_config=InferenceConfig(),
+      )
+      inference_results = await _collect_inferences(
+          inference_requests=[inference_request], eval_service=eval_service
+      )
+
+      eval_case_results = await _collect_eval_results(
+          inference_results=inference_results,
+          eval_service=eval_service,
+          eval_metrics=req.eval_metrics,
+      )
     except ModuleNotFoundError as e:
       logger.exception("%s", e)
-      raise HTTPException(status_code=400, detail=str(e)) from e
+      raise HTTPException(
+          status_code=400, detail=MISSING_EVAL_DEPENDENCIES_MESSAGE
+      ) from e
 
-    eval_set_results_manager.save_eval_set_result(
-        app_name, eval_set_id, eval_case_results
-    )
+    run_eval_results = []
+    for eval_case_result in eval_case_results:
+      run_eval_results.append(
+          RunEvalResult(
+              eval_set_file=eval_case_result.eval_set_file,
+              eval_set_id=eval_set_id,
+              eval_id=eval_case_result.eval_id,
+              final_eval_status=eval_case_result.final_eval_status,
+              overall_eval_metric_results=eval_case_result.overall_eval_metric_results,
+              eval_metric_result_per_invocation=eval_case_result.eval_metric_result_per_invocation,
+              user_id=eval_case_result.user_id,
+              session_id=eval_case_result.session_id,
+          )
+      )
 
     return run_eval_results
 
@@ -820,26 +854,30 @@ def get_fast_api_app(
 
   @working_in_progress("builder_save is not ready for use.")
   @app.post("/builder/save", response_model_exclude_none=True)
-  async def builder_build(req: AgentBuildRequest):
+  async def builder_build(files: list[UploadFile]) -> bool:
     base_path = Path.cwd() / agents_dir
-    agent = {
-        "agent_class": req.agent_type,
-        "name": req.agent_name,
-        "model": req.model,
-        "description": req.description,
-        "instruction": f"""{req.instruction}""",
-    }
-    try:
-      agent_dir = os.path.join(base_path, req.agent_name)
-      os.makedirs(agent_dir, exist_ok=True)
-      file_path = os.path.join(agent_dir, "root_agent.yaml")
-      with open(file_path, "w") as file:
-        yaml.dump(agent, file, default_flow_style=False)
-      agent_loader.load_agent(agent_name=req.agent_name)
-      return True
-    except Exception as e:
-      logger.exception("Error in builder_build: %s", e)
-      return False
+
+    for file in files:
+      try:
+        # File name format: {app_name}/{agent_name}.yaml
+        if not file.filename:
+          logger.exception("Agent name is missing in the input files")
+          return False
+
+        agent_name, filename = file.filename.split("/")
+
+        agent_dir = os.path.join(base_path, agent_name)
+        os.makedirs(agent_dir, exist_ok=True)
+        file_path = os.path.join(agent_dir, filename)
+
+        with open(file_path, "wb") as buffer:
+          shutil.copyfileobj(file.file, buffer)
+
+      except Exception as e:
+        logger.exception("Error in builder_build: %s", e)
+        return False
+
+    return True
 
   @app.post("/run", response_model_exclude_none=True)
   async def agent_run(req: AgentRunRequest) -> list[Event]:
@@ -857,7 +895,8 @@ def get_fast_api_app(
             new_message=req.new_message,
         )
     ]
-    logger.info("Generated %s events in agent run: %s", len(events), events)
+    logger.info("Generated %s events in agent run", len(events))
+    logger.debug("Events generated: %s", events)
     return events
 
   @app.post("/run_sse")
@@ -883,7 +922,7 @@ def get_fast_api_app(
         ):
           # Format as SSE data
           sse_event = event.model_dump_json(exclude_none=True, by_alias=True)
-          logger.info("Generated event in agent run streaming: %s", sse_event)
+          logger.debug("Generated event in agent run streaming: %s", sse_event)
           yield f"data: {sse_event}\n\n"
       except Exception as e:
         logger.exception("Error in event_generator: %s", e)
@@ -1136,7 +1175,9 @@ def get_fast_api_app(
 
     app.mount(
         "/dev-ui/",
-        StaticFiles(directory=ANGULAR_DIST_PATH, html=True),
+        StaticFiles(
+            directory=ANGULAR_DIST_PATH, html=True, follow_symlink=True
+        ),
         name="static",
     )
   return app

google/adk/evaluation/eval_metrics.py

@@ -36,6 +36,10 @@ class PrebuiltMetrics(Enum):
 
   RESPONSE_MATCH_SCORE = "response_match_score"
 
+  SAFETY_V1 = "safety_v1"
+
+  FINAL_RESPONSE_MATCH_V2 = "final_response_match_v2"
+
 
 MetricName: TypeAlias = Union[str, PrebuiltMetrics]
 

google/adk/evaluation/eval_sets_manager.py

@@ -36,7 +36,11 @@ class EvalSetsManager(ABC):
 
   @abstractmethod
   def list_eval_sets(self, app_name: str) -> list[str]:
-    """Returns a list of EvalSets that belong to the given app_name."""
+    """Returns a list of EvalSets that belong to the given app_name.
+
+    Raises:
+      NotFoundError: If the app_name doesn't exist.
+    """
 
   @abstractmethod
   def get_eval_case(

google/adk/evaluation/final_response_match_v2.py

@@ -21,7 +21,7 @@ from typing import Optional
 from typing_extensions import override
 
 from ..models.llm_response import LlmResponse
-from ..utils.feature_decorator import working_in_progress
+from ..utils.feature_decorator import experimental
 from .eval_case import Invocation
 from .eval_metrics import EvalMetric
 from .evaluator import EvalStatus
@@ -125,7 +125,7 @@ def _parse_critique(response: str) -> Label:
   return label
 
 
-@working_in_progress
+@experimental
 class FinalResponseMatchV2Evaluator(LlmAsJudge):
   """V2 final response match evaluator which uses an LLM to judge responses.
 

google/adk/evaluation/gcs_eval_sets_manager.py

@@ -23,6 +23,7 @@ from google.cloud import exceptions as cloud_exceptions
 from google.cloud import storage
 from typing_extensions import override
 
+from ..errors.not_found_error import NotFoundError
 from ._eval_sets_manager_utils import add_eval_case_to_eval_set
 from ._eval_sets_manager_utils import delete_eval_case_from_eval_set
 from ._eval_sets_manager_utils import get_eval_case_from_eval_set
@@ -130,7 +131,7 @@ class GcsEvalSetsManager(EvalSetsManager):
         eval_sets.append(eval_set_id)
       return sorted(eval_sets)
     except cloud_exceptions.NotFound as e:
-      raise ValueError(
+      raise NotFoundError(
           f"App `{app_name}` not found in GCS bucket `{self.bucket_name}`."
       ) from e
 

google/adk/evaluation/local_eval_service.py

@@ -30,7 +30,7 @@ from ..artifacts.in_memory_artifact_service import InMemoryArtifactService
 from ..errors.not_found_error import NotFoundError
 from ..sessions.base_session_service import BaseSessionService
 from ..sessions.in_memory_session_service import InMemorySessionService
-from ..utils.feature_decorator import working_in_progress
+from ..utils.feature_decorator import experimental
 from .base_eval_service import BaseEvalService
 from .base_eval_service import EvaluateConfig
 from .base_eval_service import EvaluateRequest
@@ -60,7 +60,7 @@ def _get_session_id() -> str:
   return f'{EVAL_SESSION_ID_PREFIX}{str(uuid.uuid4())}'
 
 
-@working_in_progress("Incomplete feature, don't use yet")
+@experimental
 class LocalEvalService(BaseEvalService):
   """An implementation of BaseEvalService, that runs the evals locally."""
 

google/adk/evaluation/local_eval_set_results_manager.py

@@ -60,7 +60,7 @@ class LocalEvalSetResultsManager(EvalSetResultsManager):
         eval_set_result.eval_set_result_name + _EVAL_SET_RESULT_FILE_EXTENSION,
     )
     logger.info("Writing eval result to file: %s", eval_set_result_file_path)
-    with open(eval_set_result_file_path, "w") as f:
+    with open(eval_set_result_file_path, "w", encoding="utf-8") as f:
       f.write(json.dumps(eval_set_result_json, indent=2))
 
   @override
@@ -78,7 +78,7 @@ class LocalEvalSetResultsManager(EvalSetResultsManager):
     )
     if not os.path.exists(maybe_eval_result_file_path):
       raise NotFoundError(f"Eval set result `{eval_set_result_id}` not found.")
-    with open(maybe_eval_result_file_path, "r") as file:
+    with open(maybe_eval_result_file_path, "r", encoding="utf-8") as file:
       eval_result_data = json.load(file)
     return EvalSetResult.model_validate_json(eval_result_data)
 

google/adk/evaluation/local_eval_sets_manager.py

@@ -315,7 +315,7 @@ class LocalEvalSetsManager(EvalSetsManager):
       )
 
   def _write_eval_set_to_path(self, eval_set_path: str, eval_set: EvalSet):
-    with open(eval_set_path, "w") as f:
+    with open(eval_set_path, "w", encoding="utf-8") as f:
       f.write(eval_set.model_dump_json(indent=2))
 
   def _save_eval_set(self, app_name: str, eval_set_id: str, eval_set: EvalSet):

google/adk/evaluation/metric_evaluator_registry.py

@@ -21,7 +21,9 @@ from .eval_metrics import EvalMetric
 from .eval_metrics import MetricName
 from .eval_metrics import PrebuiltMetrics
 from .evaluator import Evaluator
+from .final_response_match_v2 import FinalResponseMatchV2Evaluator
 from .response_evaluator import ResponseEvaluator
+from .safety_evaluator import SafetyEvaluatorV1
 from .trajectory_evaluator import TrajectoryEvaluator
 
 logger = logging.getLogger("google_adk." + __name__)
@@ -71,16 +73,24 @@ def _get_default_metric_evaluator_registry() -> MetricEvaluatorRegistry:
   metric_evaluator_registry = MetricEvaluatorRegistry()
 
   metric_evaluator_registry.register_evaluator(
-      metric_name=PrebuiltMetrics.TOOL_TRAJECTORY_AVG_SCORE,
-      evaluator=type(TrajectoryEvaluator),
+      metric_name=PrebuiltMetrics.TOOL_TRAJECTORY_AVG_SCORE.value,
+      evaluator=TrajectoryEvaluator,
   )
   metric_evaluator_registry.register_evaluator(
-      metric_name=PrebuiltMetrics.RESPONSE_EVALUATION_SCORE,
-      evaluator=type(ResponseEvaluator),
+      metric_name=PrebuiltMetrics.RESPONSE_EVALUATION_SCORE.value,
+      evaluator=ResponseEvaluator,
   )
   metric_evaluator_registry.register_evaluator(
-      metric_name=PrebuiltMetrics.RESPONSE_MATCH_SCORE,
-      evaluator=type(ResponseEvaluator),
+      metric_name=PrebuiltMetrics.RESPONSE_MATCH_SCORE.value,
+      evaluator=ResponseEvaluator,
+  )
+  metric_evaluator_registry.register_evaluator(
+      metric_name=PrebuiltMetrics.SAFETY_V1.value,
+      evaluator=SafetyEvaluatorV1,
+  )
+  metric_evaluator_registry.register_evaluator(
+      metric_name=PrebuiltMetrics.FINAL_RESPONSE_MATCH_V2.value,
+      evaluator=FinalResponseMatchV2Evaluator,
   )
 
   return metric_evaluator_registry

google/adk/evaluation/vertex_ai_eval_facade.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import math
 import os
 from typing import Optional
 
@@ -112,7 +113,12 @@ class _VertexAiEvalFacade(Evaluator):
     return ""
 
   def _get_score(self, eval_result) -> Optional[float]:
-    if eval_result and eval_result.summary_metrics:
+    if (
+        eval_result
+        and eval_result.summary_metrics
+        and isinstance(eval_result.summary_metrics[0].mean_score, float)
+        and not math.isnan(eval_result.summary_metrics[0].mean_score)
+    ):
       return eval_result.summary_metrics[0].mean_score
 
     return None

google/adk/events/event.py

@@ -42,7 +42,6 @@ class Event(LlmResponse):
     branch: The branch of the event.
     id: The unique identifier of the event.
     timestamp: The timestamp of the event.
-    is_final_response: Whether the event is the final response of the agent.
     get_function_calls: Returns the function calls in the event.
   """
 
@@ -92,7 +91,13 @@ class Event(LlmResponse):
       self.id = Event.new_id()
 
   def is_final_response(self) -> bool:
-    """Returns whether the event is the final response of the agent."""
+    """Returns whether the event is the final response of an agent.
+
+    NOTE: This method is ONLY for use by Agent Development Kit.
+
+    Note that when multiple agents participage in one invocation, there could be
+    one event has `is_final_response()` as True for each participating agent.
+    """
     if self.actions.skip_summarization or self.long_running_tool_ids:
       return True
     return (

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:
@@ -194,7 +195,12 @@ class BaseLlmFlow(ABC):
       if live_request.close:
         await llm_connection.close()
         return
-      if live_request.blob:
+
+      if live_request.activity_start:
+        await llm_connection.send_realtime(types.ActivityStart())
+      elif live_request.activity_end:
+        await llm_connection.send_realtime(types.ActivityEnd())
+      elif live_request.blob:
         # Cache audio data here for transcription
         if not invocation_context.transcription_cache:
           invocation_context.transcription_cache = []
@@ -205,6 +211,7 @@ class BaseLlmFlow(ABC):
               TranscriptionEntry(role='user', data=live_request.blob)
           )
         await llm_connection.send_realtime(live_request.blob)
+
       if live_request.content:
         await llm_connection.send_content(live_request.content)
 
@@ -335,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,