feat: warn when agent transfer runs without a context cache config
Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 956662722
This commit is contained in:
committed by
Copybara-Service
parent
7d164786f4
commit
0cf10a543e
@@ -47,6 +47,7 @@ from .errors.session_not_found_error import SessionNotFoundError
|
||||
from .events.event import Event
|
||||
from .events.event import EventActions
|
||||
from .flows.llm_flows import contents
|
||||
from .flows.llm_flows.agent_transfer import _get_transfer_targets
|
||||
from .flows.llm_flows.functions import find_event_by_function_call_id
|
||||
from .flows.llm_flows.functions import find_matching_function_call
|
||||
from .memory.base_memory_service import BaseMemoryService
|
||||
@@ -71,6 +72,9 @@ logger = logging.getLogger('google_adk.' + __name__)
|
||||
# tracer is imported for backwards compatibility, to avoid breaking change in the API.
|
||||
_ = tracer
|
||||
|
||||
# App names already told that agent transfer runs without a context cache.
|
||||
_UNCACHED_TRANSFER_APPS: set[str] = set()
|
||||
|
||||
|
||||
async def _notify_run_error(
|
||||
plugin_manager: PluginManager,
|
||||
@@ -160,6 +164,22 @@ def _apply_run_config_custom_metadata(
|
||||
}
|
||||
|
||||
|
||||
def _can_transfer_between_agents(root: Any) -> bool:
|
||||
"""Reports whether any agent in the tree can transfer to another agent."""
|
||||
pending = [root]
|
||||
while pending:
|
||||
agent = pending.pop()
|
||||
sub_agents = getattr(agent, 'sub_agents', None)
|
||||
if not isinstance(sub_agents, list):
|
||||
continue
|
||||
if hasattr(agent, 'disallow_transfer_to_parent') and _get_transfer_targets(
|
||||
agent
|
||||
):
|
||||
return True
|
||||
pending.extend(sub_agents)
|
||||
return False
|
||||
|
||||
|
||||
class Runner:
|
||||
"""The Runner class is used to run agents.
|
||||
|
||||
@@ -268,6 +288,7 @@ class Runner:
|
||||
self._agent_origin_dir = None
|
||||
self._app_name_alignment_hint: Optional[str] = None
|
||||
self._enforce_app_name_alignment()
|
||||
self._warn_uncached_agent_transfer()
|
||||
|
||||
@staticmethod
|
||||
def _resolve_app(
|
||||
@@ -428,6 +449,24 @@ class Runner:
|
||||
self._app_name_alignment_hint = f'{mismatch_details} {resolution}'
|
||||
logger.warning('App name mismatch detected. %s', mismatch_details)
|
||||
|
||||
def _warn_uncached_agent_transfer(self) -> None:
|
||||
"""Warns once per app when agent transfer runs with no context cache."""
|
||||
if self.context_cache_config is not None:
|
||||
return
|
||||
if self.app_name in _UNCACHED_TRANSFER_APPS:
|
||||
return
|
||||
if self.agent is None or not _can_transfer_between_agents(self.agent):
|
||||
return
|
||||
_UNCACHED_TRANSFER_APPS.add(self.app_name)
|
||||
logger.warning(
|
||||
'App "%s" can transfer between agents but has no'
|
||||
' context_cache_config. Every transfer swaps the system instruction'
|
||||
' and the tool set, so the request prefix changes and the whole'
|
||||
' prompt is re-sent uncached after each transfer. Set'
|
||||
' context_cache_config on the app to give each agent its own cache.',
|
||||
self.app_name,
|
||||
)
|
||||
|
||||
def _resolve_invocation_id(
|
||||
self,
|
||||
session: Session,
|
||||
|
||||
@@ -15,6 +15,7 @@
|
||||
import asyncio
|
||||
from contextlib import aclosing
|
||||
import importlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import textwrap
|
||||
@@ -22,6 +23,7 @@ from typing import AsyncGenerator
|
||||
from typing import Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from google.adk import runners
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
@@ -1537,6 +1539,76 @@ class TestRunnerCacheConfig:
|
||||
assert str(runner.context_cache_config) == expected_str
|
||||
|
||||
|
||||
class TestRunnerUncachedTransferWarning:
|
||||
"""Tests for the warning about agent transfer without a context cache."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures."""
|
||||
self.session_service = InMemorySessionService()
|
||||
runners._UNCACHED_TRANSFER_APPS.clear()
|
||||
|
||||
def teardown_method(self):
|
||||
runners._UNCACHED_TRANSFER_APPS.clear()
|
||||
|
||||
def _multi_agent(self) -> LlmAgent:
|
||||
return LlmAgent(
|
||||
name="root_agent",
|
||||
model="gemini-1.5-pro",
|
||||
sub_agents=[MockLlmAgent("sub_agent")],
|
||||
)
|
||||
|
||||
def _warnings(self, caplog) -> list[str]:
|
||||
return [
|
||||
record.getMessage()
|
||||
for record in caplog.records
|
||||
if record.levelno == logging.WARNING
|
||||
and "context_cache_config" in record.getMessage()
|
||||
]
|
||||
|
||||
def test_warns_for_multi_agent_app_without_cache_config(self, caplog):
|
||||
"""Transfer is possible and no cache is configured, so warn."""
|
||||
app = App(name="multi_agent_app", root_agent=self._multi_agent())
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
Runner(app=app, session_service=self.session_service)
|
||||
|
||||
messages = self._warnings(caplog)
|
||||
assert len(messages) == 1
|
||||
assert "multi_agent_app" in messages[0]
|
||||
|
||||
def test_no_warning_when_cache_config_present(self, caplog):
|
||||
"""An app that configures a context cache is not warned."""
|
||||
app = App(
|
||||
name="cached_app",
|
||||
root_agent=self._multi_agent(),
|
||||
context_cache_config=ContextCacheConfig(),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
Runner(app=app, session_service=self.session_service)
|
||||
|
||||
assert not self._warnings(caplog)
|
||||
|
||||
def test_no_warning_without_transfer_targets(self, caplog):
|
||||
"""A single-agent app cannot transfer, so nothing is lost."""
|
||||
app = App(name="single_agent_app", root_agent=MockLlmAgent("root_agent"))
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
Runner(app=app, session_service=self.session_service)
|
||||
|
||||
assert not self._warnings(caplog)
|
||||
|
||||
def test_warns_only_once_per_app(self, caplog):
|
||||
"""Rebuilding the runner for the same app does not warn again."""
|
||||
app = App(name="multi_agent_app", root_agent=self._multi_agent())
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
for _ in range(3):
|
||||
Runner(app=app, session_service=self.session_service)
|
||||
|
||||
assert len(self._warnings(caplog)) == 1
|
||||
|
||||
|
||||
class TestRunnerResolveApp:
|
||||
"""Tests for Runner._resolve_app and node support."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user