You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
feat: Add a app level config for resumable applications
PiperOrigin-RevId: 811272046
This commit is contained in:
committed by
Copybara-Service
parent
c6b6b6f3c6
commit
cbb6e4945a
@@ -15,7 +15,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING
|
||||
import uuid
|
||||
|
||||
from google.genai import types
|
||||
@@ -24,6 +23,7 @@ from pydantic import ConfigDict
|
||||
from pydantic import Field
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from ..apps.app import ResumabilityConfig
|
||||
from ..artifacts.base_artifact_service import BaseArtifactService
|
||||
from ..auth.credential_service.base_credential_service import BaseCredentialService
|
||||
from ..events.event import Event
|
||||
@@ -31,7 +31,6 @@ from ..memory.base_memory_service import BaseMemoryService
|
||||
from ..plugins.plugin_manager import PluginManager
|
||||
from ..sessions.base_session_service import BaseSessionService
|
||||
from ..sessions.session import Session
|
||||
from ..utils.feature_decorator import working_in_progress
|
||||
from .active_streaming_tool import ActiveStreamingTool
|
||||
from .base_agent import BaseAgent
|
||||
from .context_cache_config import ContextCacheConfig
|
||||
@@ -189,6 +188,9 @@ class InvocationContext(BaseModel):
|
||||
run_config: Optional[RunConfig] = None
|
||||
"""Configurations for live agents under this invocation."""
|
||||
|
||||
resumability_config: Optional[ResumabilityConfig] = None
|
||||
"""The resumability config that applies to all agents under this invocation."""
|
||||
|
||||
plugin_manager: PluginManager = Field(default_factory=PluginManager)
|
||||
"""The manager for keeping track of plugins in this invocation."""
|
||||
|
||||
@@ -220,7 +222,6 @@ class InvocationContext(BaseModel):
|
||||
def user_id(self) -> str:
|
||||
return self.session.user_id
|
||||
|
||||
@working_in_progress("incomplete feature, don't use yet")
|
||||
def get_events(
|
||||
self,
|
||||
current_invocation: bool = False,
|
||||
|
||||
@@ -13,7 +13,9 @@
|
||||
# limitations under the License.
|
||||
|
||||
from .app import App
|
||||
from .app import ResumabilityConfig
|
||||
|
||||
__all__ = [
|
||||
'App',
|
||||
'ResumabilityConfig',
|
||||
]
|
||||
|
||||
@@ -26,6 +26,28 @@ from ..plugins.base_plugin import BasePlugin
|
||||
from ..utils.feature_decorator import experimental
|
||||
|
||||
|
||||
@experimental
|
||||
class ResumabilityConfig(BaseModel):
|
||||
"""The config of the resumability for an application.
|
||||
|
||||
The "resumability" in ADK refers to the ability to:
|
||||
1. pause an invocation upon a long running function call.
|
||||
2. resume an invocation from the last event, if it's paused or failed midway
|
||||
through.
|
||||
|
||||
Note: ADK resumes the invocation in a best-effort manner. Edge cases include:
|
||||
1. The invocation crashed before receiving the response of a tool call, once
|
||||
resumed, the same tool call will be invoked again.
|
||||
2. If agent transfer forms a loop (root->sub1->sub2->root[pause]), once
|
||||
resumed, the agent transfer loop will be ignored.
|
||||
"""
|
||||
|
||||
is_resumable: bool = False
|
||||
"""Whether the app supports agent resumption.
|
||||
If enabled, the feature will be enabled for all agents in the app.
|
||||
"""
|
||||
|
||||
|
||||
@experimental
|
||||
class App(BaseModel):
|
||||
"""Represents an LLM-backed agentic application.
|
||||
@@ -57,3 +79,9 @@ class App(BaseModel):
|
||||
|
||||
context_cache_config: Optional[ContextCacheConfig] = None
|
||||
"""Context cache configuration that applies to all LLM agents in the app."""
|
||||
|
||||
resumability_config: Optional[ResumabilityConfig] = None
|
||||
"""
|
||||
The config of the resumability for the application.
|
||||
If configured, will be applied to all agents in the app.
|
||||
"""
|
||||
|
||||
@@ -36,6 +36,7 @@ from .agents.live_request_queue import LiveRequestQueue
|
||||
from .agents.llm_agent import LlmAgent
|
||||
from .agents.run_config import RunConfig
|
||||
from .apps.app import App
|
||||
from .apps.app import ResumabilityConfig
|
||||
from .artifacts.base_artifact_service import BaseArtifactService
|
||||
from .artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from .auth.credential_service.base_credential_service import BaseCredentialService
|
||||
@@ -74,6 +75,8 @@ class Runner:
|
||||
session_service: The session service for the runner.
|
||||
memory_service: The memory service for the runner.
|
||||
credential_service: The credential service for the runner.
|
||||
context_cache_config: The context cache config for the runner.
|
||||
resumability_config: The resumability config for the application.
|
||||
"""
|
||||
|
||||
app_name: str
|
||||
@@ -90,6 +93,10 @@ class Runner:
|
||||
"""The memory service for the runner."""
|
||||
credential_service: Optional[BaseCredentialService] = None
|
||||
"""The credential service for the runner."""
|
||||
context_cache_config: Optional[ContextCacheConfig] = None
|
||||
"""The context cache config for the runner."""
|
||||
resumability_config: Optional[ResumabilityConfig] = None
|
||||
"""The resumability config for the application."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -110,11 +117,11 @@ class Runner:
|
||||
`ValueError`. Providing `app` is the recommended way to create a runner.
|
||||
|
||||
Args:
|
||||
app: An optional `App` instance. If provided, `app_name` and `agent`
|
||||
should not be specified.
|
||||
app_name: The application name of the runner. Required if `app` is not
|
||||
provided.
|
||||
agent: The root agent to run. Required if `app` is not provided.
|
||||
app: An optional `App` instance. If provided, `app_name` and `agent`
|
||||
should not be specified.
|
||||
plugins: Deprecated. A list of plugins for the runner. Please use the
|
||||
`app` argument to provide plugins instead.
|
||||
artifact_service: The artifact service for the runner.
|
||||
@@ -126,9 +133,13 @@ class Runner:
|
||||
ValueError: If `app` is provided along with `app_name` or `plugins`, or
|
||||
if `app` is not provided but either `app_name` or `agent` is missing.
|
||||
"""
|
||||
self.app_name, self.agent, self.context_cache_config, plugins = (
|
||||
self._validate_runner_params(app, app_name, agent, plugins)
|
||||
)
|
||||
(
|
||||
self.app_name,
|
||||
self.agent,
|
||||
self.context_cache_config,
|
||||
self.resumability_config,
|
||||
plugins,
|
||||
) = self._validate_runner_params(app, app_name, agent, plugins)
|
||||
self.artifact_service = artifact_service
|
||||
self.session_service = session_service
|
||||
self.memory_service = memory_service
|
||||
@@ -142,7 +153,11 @@ class Runner:
|
||||
agent: Optional[BaseAgent],
|
||||
plugins: Optional[List[BasePlugin]],
|
||||
) -> tuple[
|
||||
str, BaseAgent, Optional[ContextCacheConfig], Optional[List[BasePlugin]]
|
||||
str,
|
||||
BaseAgent,
|
||||
Optional[ContextCacheConfig],
|
||||
Optional[ResumabilityConfig],
|
||||
Optional[List[BasePlugin]],
|
||||
]:
|
||||
"""Validates and extracts runner parameters.
|
||||
|
||||
@@ -153,7 +168,8 @@ class Runner:
|
||||
plugins: A list of plugins for the runner.
|
||||
|
||||
Returns:
|
||||
A tuple containing (app_name, agent, context_cache_config, plugins).
|
||||
A tuple containing (app_name, agent, context_cache_config,
|
||||
resumability_config, plugins).
|
||||
|
||||
Raises:
|
||||
ValueError: If parameters are invalid.
|
||||
@@ -174,12 +190,14 @@ class Runner:
|
||||
agent = app.root_agent
|
||||
plugins = app.plugins
|
||||
context_cache_config = app.context_cache_config
|
||||
resumability_config = app.resumability_config
|
||||
elif not app_name or not agent:
|
||||
raise ValueError(
|
||||
'Either app or both app_name and agent must be provided.'
|
||||
)
|
||||
else:
|
||||
context_cache_config = None
|
||||
resumability_config = None
|
||||
|
||||
if plugins:
|
||||
warnings.warn(
|
||||
@@ -187,7 +205,7 @@ class Runner:
|
||||
' to provide plugins instead.',
|
||||
DeprecationWarning,
|
||||
)
|
||||
return app_name, agent, context_cache_config, plugins
|
||||
return app_name, agent, context_cache_config, resumability_config, plugins
|
||||
|
||||
def run(
|
||||
self,
|
||||
@@ -264,6 +282,7 @@ class Runner:
|
||||
user_id: The user ID of the session.
|
||||
session_id: The session ID of the session.
|
||||
new_message: A new message to append to the session.
|
||||
state_delta: Optional state changes to apply to the session.
|
||||
run_config: The run config for the agent.
|
||||
|
||||
Yields:
|
||||
@@ -687,6 +706,7 @@ class Runner:
|
||||
user_content=new_message,
|
||||
live_request_queue=live_request_queue,
|
||||
run_config=run_config,
|
||||
resumability_config=self.resumability_config,
|
||||
)
|
||||
|
||||
def _new_invocation_context_for_live(
|
||||
|
||||
@@ -17,6 +17,7 @@ from unittest.mock import Mock
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.apps.app import ResumabilityConfig
|
||||
from google.adk.plugins.base_plugin import BasePlugin
|
||||
|
||||
|
||||
@@ -68,6 +69,23 @@ class TestApp:
|
||||
assert app.context_cache_config.ttl_seconds == 3600
|
||||
assert app.context_cache_config.min_tokens == 1024
|
||||
|
||||
def test_app_initialization_with_resumability_config(self):
|
||||
"""Test that the app is initialized correctly with app config."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
resumability_config = ResumabilityConfig(
|
||||
is_resumable=True,
|
||||
)
|
||||
app = App(
|
||||
name="test_app",
|
||||
root_agent=mock_agent,
|
||||
resumability_config=resumability_config,
|
||||
)
|
||||
|
||||
assert app.name == "test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.resumability_config == resumability_config
|
||||
assert app.resumability_config.is_resumable
|
||||
|
||||
def test_app_with_all_components(self):
|
||||
"""Test app with all components: agent, plugins, and cache config."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
@@ -75,18 +93,24 @@ class TestApp:
|
||||
cache_config = ContextCacheConfig(
|
||||
cache_intervals=20, ttl_seconds=7200, min_tokens=2048
|
||||
)
|
||||
resumability_config = ResumabilityConfig(
|
||||
is_resumable=True,
|
||||
)
|
||||
|
||||
app = App(
|
||||
name="full_test_app",
|
||||
root_agent=mock_agent,
|
||||
plugins=[mock_plugin],
|
||||
context_cache_config=cache_config,
|
||||
resumability_config=resumability_config,
|
||||
)
|
||||
|
||||
assert app.name == "full_test_app"
|
||||
assert app.root_agent == mock_agent
|
||||
assert app.plugins == [mock_plugin]
|
||||
assert app.context_cache_config == cache_config
|
||||
assert app.resumability_config == resumability_config
|
||||
assert app.resumability_config.is_resumable
|
||||
|
||||
def test_app_cache_config_defaults(self):
|
||||
"""Test that cache config has proper defaults when created."""
|
||||
@@ -118,3 +142,29 @@ class TestApp:
|
||||
context_cache_config=None,
|
||||
)
|
||||
assert app.context_cache_config is None
|
||||
|
||||
def test_app_resumability_config_defaults(self):
|
||||
"""Test that app config has proper defaults when created."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
|
||||
app = App(
|
||||
name="default_resumability_config_app",
|
||||
root_agent=mock_agent,
|
||||
resumability_config=ResumabilityConfig(),
|
||||
)
|
||||
assert app.resumability_config is not None
|
||||
assert not app.resumability_config.is_resumable # Default
|
||||
|
||||
def test_app_resumability_config_is_optional(self):
|
||||
"""Test that resumability_config is truly optional."""
|
||||
mock_agent = Mock(spec=BaseAgent)
|
||||
|
||||
app = App(name="no_resumability_config_app", root_agent=mock_agent)
|
||||
assert app.resumability_config is None
|
||||
|
||||
app = App(
|
||||
name="explicit_none_resumability_config_app",
|
||||
root_agent=mock_agent,
|
||||
resumability_config=None,
|
||||
)
|
||||
assert app.resumability_config is None
|
||||
|
||||
@@ -19,6 +19,7 @@ from google.adk.agents.context_cache_config import ContextCacheConfig
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.apps.app import ResumabilityConfig
|
||||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.plugins.base_plugin import BasePlugin
|
||||
@@ -565,6 +566,7 @@ class TestRunnerCacheConfig:
|
||||
name="order_test_app",
|
||||
root_agent=self.root_agent,
|
||||
context_cache_config=cache_config,
|
||||
resumability_config=ResumabilityConfig(is_resumable=True),
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
@@ -574,7 +576,7 @@ class TestRunnerCacheConfig:
|
||||
)
|
||||
|
||||
# Test the validation method directly
|
||||
app_name, agent, context_cache_config, plugins = (
|
||||
app_name, agent, context_cache_config, resumability_config, plugins = (
|
||||
runner._validate_runner_params(app, None, None, None)
|
||||
)
|
||||
|
||||
@@ -582,6 +584,7 @@ class TestRunnerCacheConfig:
|
||||
assert agent == self.root_agent
|
||||
assert context_cache_config == cache_config
|
||||
assert context_cache_config.cache_intervals == 25
|
||||
assert resumability_config == app.resumability_config
|
||||
assert plugins == []
|
||||
|
||||
def test_runner_validate_params_without_app(self):
|
||||
@@ -593,13 +596,14 @@ class TestRunnerCacheConfig:
|
||||
artifact_service=self.artifact_service,
|
||||
)
|
||||
|
||||
app_name, agent, context_cache_config, plugins = (
|
||||
app_name, agent, context_cache_config, resumability_config, plugins = (
|
||||
runner._validate_runner_params(None, "test_app", self.root_agent, None)
|
||||
)
|
||||
|
||||
assert app_name == "test_app"
|
||||
assert agent == self.root_agent
|
||||
assert context_cache_config is None
|
||||
assert resumability_config is None
|
||||
assert plugins is None
|
||||
|
||||
def test_runner_app_name_and_agent_extracted_correctly(self):
|
||||
|
||||
Reference in New Issue
Block a user