You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
chore: Introduce a unified Context
Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 868890552
This commit is contained in:
Xuan Yang
committed by
Copybara-Service
Copybara-Service
parent
8cd22fb746
commit
b7f9110b52
@@ -0,0 +1,373 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from .readonly_context import ReadonlyContext
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from google.genai import types
|
||||
|
||||
from ..artifacts.base_artifact_service import ArtifactVersion
|
||||
from ..auth.auth_credential import AuthCredential
|
||||
from ..auth.auth_tool import AuthConfig
|
||||
from ..events.event import Event
|
||||
from ..events.event_actions import EventActions
|
||||
from ..memory.base_memory_service import SearchMemoryResponse
|
||||
from ..sessions.state import State
|
||||
from ..tools.tool_confirmation import ToolConfirmation
|
||||
from .invocation_context import InvocationContext
|
||||
|
||||
|
||||
class Context(ReadonlyContext):
|
||||
"""The context within an agent run."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
invocation_context: InvocationContext,
|
||||
*,
|
||||
event_actions: EventActions | None = None,
|
||||
function_call_id: str | None = None,
|
||||
tool_confirmation: ToolConfirmation | None = None,
|
||||
) -> None:
|
||||
"""Initializes the Context.
|
||||
|
||||
Args:
|
||||
invocation_context: The invocation context.
|
||||
event_actions: The event actions for state and artifact deltas.
|
||||
function_call_id: The function call id of the current tool call. Required
|
||||
for tool-specific methods like request_credential and
|
||||
request_confirmation.
|
||||
tool_confirmation: The tool confirmation of the current tool call.
|
||||
"""
|
||||
super().__init__(invocation_context)
|
||||
|
||||
from ..events.event_actions import EventActions
|
||||
from ..sessions.state import State
|
||||
|
||||
self._event_actions = event_actions or EventActions()
|
||||
self._state = State(
|
||||
value=invocation_context.session.state,
|
||||
delta=self._event_actions.state_delta,
|
||||
)
|
||||
self._function_call_id = function_call_id
|
||||
self._tool_confirmation = tool_confirmation
|
||||
|
||||
@property
|
||||
def function_call_id(self) -> str | None:
|
||||
"""The function call id of the current tool call."""
|
||||
return self._function_call_id
|
||||
|
||||
@property
|
||||
def tool_confirmation(self) -> ToolConfirmation | None:
|
||||
"""The tool confirmation of the current tool call."""
|
||||
return self._tool_confirmation
|
||||
|
||||
@property
|
||||
@override
|
||||
def state(self) -> State:
|
||||
"""The delta-aware state of the current session.
|
||||
|
||||
For any state change, you can mutate this object directly,
|
||||
e.g. `ctx.state['foo'] = 'bar'`
|
||||
"""
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def actions(self) -> EventActions:
|
||||
"""The event actions for the current context."""
|
||||
return self._event_actions
|
||||
|
||||
# ============================================================================
|
||||
# Artifact methods
|
||||
# ============================================================================
|
||||
|
||||
async def load_artifact(
|
||||
self, filename: str, version: int | None = None
|
||||
) -> types.Part | None:
|
||||
"""Loads an artifact attached to the current session.
|
||||
|
||||
Args:
|
||||
filename: The filename of the artifact.
|
||||
version: The version of the artifact. If None, the latest version will be
|
||||
returned.
|
||||
|
||||
Returns:
|
||||
The artifact.
|
||||
"""
|
||||
if self._invocation_context.artifact_service is None:
|
||||
raise ValueError("Artifact service is not initialized.")
|
||||
return await self._invocation_context.artifact_service.load_artifact(
|
||||
app_name=self._invocation_context.app_name,
|
||||
user_id=self._invocation_context.user_id,
|
||||
session_id=self._invocation_context.session.id,
|
||||
filename=filename,
|
||||
version=version,
|
||||
)
|
||||
|
||||
async def save_artifact(
|
||||
self,
|
||||
filename: str,
|
||||
artifact: types.Part,
|
||||
custom_metadata: dict[str, Any] | None = None,
|
||||
) -> int:
|
||||
"""Saves an artifact and records it as delta for the current session.
|
||||
|
||||
Args:
|
||||
filename: The filename of the artifact.
|
||||
artifact: The artifact to save.
|
||||
custom_metadata: Custom metadata to associate with the artifact.
|
||||
|
||||
Returns:
|
||||
The version of the artifact.
|
||||
"""
|
||||
if self._invocation_context.artifact_service is None:
|
||||
raise ValueError("Artifact service is not initialized.")
|
||||
version = await self._invocation_context.artifact_service.save_artifact(
|
||||
app_name=self._invocation_context.app_name,
|
||||
user_id=self._invocation_context.user_id,
|
||||
session_id=self._invocation_context.session.id,
|
||||
filename=filename,
|
||||
artifact=artifact,
|
||||
custom_metadata=custom_metadata,
|
||||
)
|
||||
self._event_actions.artifact_delta[filename] = version
|
||||
return version
|
||||
|
||||
async def get_artifact_version(
|
||||
self, filename: str, version: int | None = None
|
||||
) -> ArtifactVersion | None:
|
||||
"""Gets artifact version info.
|
||||
|
||||
Args:
|
||||
filename: The filename of the artifact.
|
||||
version: The version of the artifact. If None, the latest version will be
|
||||
returned.
|
||||
|
||||
Returns:
|
||||
The artifact version info.
|
||||
"""
|
||||
if self._invocation_context.artifact_service is None:
|
||||
raise ValueError("Artifact service is not initialized.")
|
||||
return await self._invocation_context.artifact_service.get_artifact_version(
|
||||
app_name=self._invocation_context.app_name,
|
||||
user_id=self._invocation_context.user_id,
|
||||
session_id=self._invocation_context.session.id,
|
||||
filename=filename,
|
||||
version=version,
|
||||
)
|
||||
|
||||
async def list_artifacts(self) -> list[str]:
|
||||
"""Lists the filenames of the artifacts attached to the current session."""
|
||||
if self._invocation_context.artifact_service is None:
|
||||
raise ValueError("Artifact service is not initialized.")
|
||||
return await self._invocation_context.artifact_service.list_artifact_keys(
|
||||
app_name=self._invocation_context.app_name,
|
||||
user_id=self._invocation_context.user_id,
|
||||
session_id=self._invocation_context.session.id,
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# Credential methods
|
||||
# ============================================================================
|
||||
|
||||
async def save_credential(self, auth_config: AuthConfig) -> None:
|
||||
"""Saves a credential to the credential service.
|
||||
|
||||
Args:
|
||||
auth_config: The authentication configuration containing the credential.
|
||||
"""
|
||||
if self._invocation_context.credential_service is None:
|
||||
raise ValueError("Credential service is not initialized.")
|
||||
await self._invocation_context.credential_service.save_credential(
|
||||
auth_config, self
|
||||
)
|
||||
|
||||
async def load_credential(
|
||||
self, auth_config: AuthConfig
|
||||
) -> AuthCredential | None:
|
||||
"""Loads a credential from the credential service.
|
||||
|
||||
Args:
|
||||
auth_config: The authentication configuration for the credential.
|
||||
|
||||
Returns:
|
||||
The loaded credential, or None if not found.
|
||||
"""
|
||||
if self._invocation_context.credential_service is None:
|
||||
raise ValueError("Credential service is not initialized.")
|
||||
return await self._invocation_context.credential_service.load_credential(
|
||||
auth_config, self
|
||||
)
|
||||
|
||||
def get_auth_response(self, auth_config: AuthConfig) -> AuthCredential | None:
|
||||
"""Gets the auth response credential from session state.
|
||||
|
||||
This method retrieves an authentication credential that was previously
|
||||
stored in session state after a user completed an OAuth flow or other
|
||||
authentication process.
|
||||
|
||||
Args:
|
||||
auth_config: The authentication configuration for the credential.
|
||||
|
||||
Returns:
|
||||
The auth credential from the auth response, or None if not found.
|
||||
"""
|
||||
from ..auth.auth_handler import AuthHandler
|
||||
|
||||
return AuthHandler(auth_config).get_auth_response(self.state)
|
||||
|
||||
def request_credential(self, auth_config: AuthConfig) -> None:
|
||||
"""Requests a credential for the current tool call.
|
||||
|
||||
This method can only be called in a tool context where function_call_id
|
||||
is set. For callback contexts, use save_credential/load_credential instead.
|
||||
|
||||
Args:
|
||||
auth_config: The authentication configuration for the credential.
|
||||
|
||||
Raises:
|
||||
ValueError: If function_call_id is not set.
|
||||
"""
|
||||
from ..auth.auth_handler import AuthHandler
|
||||
|
||||
if not self.function_call_id:
|
||||
raise ValueError(
|
||||
"request_credential requires function_call_id. "
|
||||
"This method can only be used in a tool context, not a callback "
|
||||
"context. Consider using save_credential/load_credential instead."
|
||||
)
|
||||
self._event_actions.requested_auth_configs[self.function_call_id] = (
|
||||
AuthHandler(auth_config).generate_auth_request()
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# Tool methods
|
||||
# ============================================================================
|
||||
|
||||
def request_confirmation(
|
||||
self,
|
||||
*,
|
||||
hint: str | None = None,
|
||||
payload: Any | None = None,
|
||||
) -> None:
|
||||
"""Requests confirmation for the current tool call.
|
||||
|
||||
This method can only be called in a tool context where function_call_id
|
||||
is set.
|
||||
|
||||
Args:
|
||||
hint: A hint to the user on how to confirm the tool call.
|
||||
payload: The payload used to confirm the tool call.
|
||||
|
||||
Raises:
|
||||
ValueError: If function_call_id is not set.
|
||||
"""
|
||||
from ..tools.tool_confirmation import ToolConfirmation
|
||||
|
||||
if not self.function_call_id:
|
||||
raise ValueError(
|
||||
"request_confirmation requires function_call_id. "
|
||||
"This method can only be used in a tool context."
|
||||
)
|
||||
self._event_actions.requested_tool_confirmations[self.function_call_id] = (
|
||||
ToolConfirmation(
|
||||
hint=hint,
|
||||
payload=payload,
|
||||
)
|
||||
)
|
||||
|
||||
# ============================================================================
|
||||
# Memory methods
|
||||
# ============================================================================
|
||||
|
||||
async def add_session_to_memory(self) -> None:
|
||||
"""Triggers memory generation for the current session.
|
||||
|
||||
This method saves the current session's events to the memory service,
|
||||
enabling the agent to recall information from past interactions.
|
||||
|
||||
Raises:
|
||||
ValueError: If memory service is not available.
|
||||
|
||||
Example:
|
||||
```python
|
||||
async def my_after_agent_callback(ctx: Context):
|
||||
# Save conversation to memory at the end of each interaction
|
||||
await ctx.add_session_to_memory()
|
||||
```
|
||||
"""
|
||||
if self._invocation_context.memory_service is None:
|
||||
raise ValueError(
|
||||
"Cannot add session to memory: memory service is not available."
|
||||
)
|
||||
await self._invocation_context.memory_service.add_session_to_memory(
|
||||
self._invocation_context.session
|
||||
)
|
||||
|
||||
async def add_events_to_memory(
|
||||
self,
|
||||
*,
|
||||
events: Sequence[Event],
|
||||
custom_metadata: Mapping[str, object] | None = None,
|
||||
) -> None:
|
||||
"""Adds an explicit list of events to the memory service.
|
||||
|
||||
Uses this callback's current session identifiers as memory scope.
|
||||
|
||||
Args:
|
||||
events: Explicit events to add to memory.
|
||||
custom_metadata: Optional standard metadata for memory generation.
|
||||
|
||||
Raises:
|
||||
ValueError: If memory service is not available.
|
||||
"""
|
||||
if self._invocation_context.memory_service is None:
|
||||
raise ValueError(
|
||||
"Cannot add events to memory: memory service is not available."
|
||||
)
|
||||
await self._invocation_context.memory_service.add_events_to_memory(
|
||||
app_name=self._invocation_context.session.app_name,
|
||||
user_id=self._invocation_context.session.user_id,
|
||||
session_id=self._invocation_context.session.id,
|
||||
events=events,
|
||||
custom_metadata=custom_metadata,
|
||||
)
|
||||
|
||||
async def search_memory(self, query: str) -> SearchMemoryResponse:
|
||||
"""Searches the memory of the current user.
|
||||
|
||||
Args:
|
||||
query: The search query.
|
||||
|
||||
Returns:
|
||||
The search results from the memory service.
|
||||
|
||||
Raises:
|
||||
ValueError: If memory service is not available.
|
||||
"""
|
||||
if self._invocation_context.memory_service is None:
|
||||
raise ValueError("Memory service is not available.")
|
||||
return await self._invocation_context.memory_service.search_memory(
|
||||
app_name=self._invocation_context.app_name,
|
||||
user_id=self._invocation_context.user_id,
|
||||
query=query,
|
||||
)
|
||||
@@ -0,0 +1,488 @@
|
||||
# Copyright 2026 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
|
||||
from google.adk.agents.context import Context
|
||||
from google.adk.auth import auth_handler
|
||||
from google.adk.auth.auth_credential import AuthCredential
|
||||
from google.adk.auth.auth_credential import AuthCredentialTypes
|
||||
from google.adk.auth.auth_tool import AuthConfig
|
||||
from google.adk.memory.base_memory_service import SearchMemoryResponse
|
||||
from google.adk.tools.tool_confirmation import ToolConfirmation
|
||||
from google.genai.types import Part
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_invocation_context():
|
||||
"""Create a mock invocation context for testing."""
|
||||
mock_context = MagicMock()
|
||||
mock_context.invocation_id = "test-invocation-id"
|
||||
mock_context.agent.name = "test-agent-name"
|
||||
mock_context.session.state = {"key1": "value1", "key2": "value2"}
|
||||
mock_context.session.id = "test-session-id"
|
||||
mock_context.app_name = "test-app"
|
||||
mock_context.user_id = "test-user"
|
||||
mock_context.artifact_service = None
|
||||
mock_context.credential_service = None
|
||||
mock_context.memory_service = None
|
||||
return mock_context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_artifact_service():
|
||||
"""Create a mock artifact service for testing."""
|
||||
mock_service = AsyncMock()
|
||||
mock_service.list_artifact_keys.return_value = [
|
||||
"file1.txt",
|
||||
"file2.txt",
|
||||
"file3.txt",
|
||||
]
|
||||
return mock_service
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_config(mocker):
|
||||
"""Create a mock auth config for testing."""
|
||||
return mocker.create_autospec(AuthConfig, instance=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_credential(mocker):
|
||||
"""Create a mock auth credential for testing."""
|
||||
mock_credential = mocker.create_autospec(AuthCredential, instance=True)
|
||||
mock_credential.auth_type = AuthCredentialTypes.OAUTH2
|
||||
return mock_credential
|
||||
|
||||
|
||||
class TestContextInitialization:
|
||||
"""Test Context initialization."""
|
||||
|
||||
def test_initialization_without_function_call_id(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test Context initialization without function_call_id."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
assert context._invocation_context == mock_invocation_context
|
||||
assert context._event_actions is not None
|
||||
assert context._state is not None
|
||||
assert context.function_call_id is None
|
||||
assert context.tool_confirmation is None
|
||||
|
||||
def test_initialization_with_function_call_id(self, mock_invocation_context):
|
||||
"""Test Context initialization with function_call_id."""
|
||||
context = Context(
|
||||
mock_invocation_context,
|
||||
function_call_id="test-function-call-id",
|
||||
)
|
||||
|
||||
assert context.function_call_id == "test-function-call-id"
|
||||
assert context.tool_confirmation is None
|
||||
|
||||
def test_initialization_with_tool_confirmation(self, mock_invocation_context):
|
||||
"""Test Context initialization with tool_confirmation."""
|
||||
tool_confirmation = ToolConfirmation(
|
||||
hint="test hint", payload={"key": "value"}
|
||||
)
|
||||
context = Context(
|
||||
mock_invocation_context,
|
||||
function_call_id="test-function-call-id",
|
||||
tool_confirmation=tool_confirmation,
|
||||
)
|
||||
|
||||
assert context.function_call_id == "test-function-call-id"
|
||||
assert context.tool_confirmation == tool_confirmation
|
||||
assert context.tool_confirmation.hint == "test hint"
|
||||
assert context.tool_confirmation.payload == {"key": "value"}
|
||||
|
||||
def test_state_property(self, mock_invocation_context):
|
||||
"""Test that state property returns mutable state."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
assert context.state["key1"] == "value1"
|
||||
assert context.state["key2"] == "value2"
|
||||
|
||||
def test_actions_property(self, mock_invocation_context):
|
||||
"""Test that actions property returns event_actions."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
assert context.actions is context._event_actions
|
||||
|
||||
|
||||
class TestContextListArtifacts:
|
||||
"""Test the list_artifacts method in Context."""
|
||||
|
||||
async def test_list_artifacts_returns_artifact_keys(
|
||||
self, mock_invocation_context, mock_artifact_service
|
||||
):
|
||||
"""Test that list_artifacts returns the artifact keys from the service."""
|
||||
mock_invocation_context.artifact_service = mock_artifact_service
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
result = await context.list_artifacts()
|
||||
|
||||
assert result == ["file1.txt", "file2.txt", "file3.txt"]
|
||||
mock_artifact_service.list_artifact_keys.assert_called_once_with(
|
||||
app_name="test-app",
|
||||
user_id="test-user",
|
||||
session_id="test-session-id",
|
||||
)
|
||||
|
||||
async def test_list_artifacts_raises_value_error_when_service_is_none(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test that list_artifacts raises ValueError when no artifact service."""
|
||||
mock_invocation_context.artifact_service = None
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Artifact service is not initialized."
|
||||
):
|
||||
await context.list_artifacts()
|
||||
|
||||
|
||||
class TestContextSaveLoadArtifact:
|
||||
"""Test save_artifact and load_artifact methods in Context."""
|
||||
|
||||
async def test_save_artifact(self, mock_invocation_context):
|
||||
"""Test save_artifact method."""
|
||||
artifact_service = AsyncMock()
|
||||
artifact_service.save_artifact.return_value = 1
|
||||
mock_invocation_context.artifact_service = artifact_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
test_artifact = Part.from_text(text="test content")
|
||||
|
||||
version = await context.save_artifact("test_file.txt", test_artifact)
|
||||
|
||||
artifact_service.save_artifact.assert_called_once_with(
|
||||
app_name="test-app",
|
||||
user_id="test-user",
|
||||
session_id="test-session-id",
|
||||
filename="test_file.txt",
|
||||
artifact=test_artifact,
|
||||
custom_metadata=None,
|
||||
)
|
||||
assert version == 1
|
||||
assert context.actions.artifact_delta["test_file.txt"] == 1
|
||||
|
||||
async def test_load_artifact(self, mock_invocation_context):
|
||||
"""Test load_artifact method."""
|
||||
artifact_service = AsyncMock()
|
||||
test_artifact = Part.from_text(text="test content")
|
||||
artifact_service.load_artifact.return_value = test_artifact
|
||||
mock_invocation_context.artifact_service = artifact_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
result = await context.load_artifact("test_file.txt")
|
||||
|
||||
artifact_service.load_artifact.assert_called_once_with(
|
||||
app_name="test-app",
|
||||
user_id="test-user",
|
||||
session_id="test-session-id",
|
||||
filename="test_file.txt",
|
||||
version=None,
|
||||
)
|
||||
assert result == test_artifact
|
||||
|
||||
async def test_load_artifact_with_version(self, mock_invocation_context):
|
||||
"""Test load_artifact method with specific version."""
|
||||
artifact_service = AsyncMock()
|
||||
test_artifact = Part.from_text(text="test content")
|
||||
artifact_service.load_artifact.return_value = test_artifact
|
||||
mock_invocation_context.artifact_service = artifact_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
result = await context.load_artifact("test_file.txt", version=2)
|
||||
|
||||
artifact_service.load_artifact.assert_called_once_with(
|
||||
app_name="test-app",
|
||||
user_id="test-user",
|
||||
session_id="test-session-id",
|
||||
filename="test_file.txt",
|
||||
version=2,
|
||||
)
|
||||
assert result == test_artifact
|
||||
|
||||
|
||||
class TestContextCredentialMethods:
|
||||
"""Test credential methods in Context."""
|
||||
|
||||
async def test_save_credential_with_service(
|
||||
self, mock_invocation_context, mock_auth_config
|
||||
):
|
||||
"""Test save_credential when credential service is available."""
|
||||
credential_service = AsyncMock()
|
||||
mock_invocation_context.credential_service = credential_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
await context.save_credential(mock_auth_config)
|
||||
|
||||
credential_service.save_credential.assert_called_once_with(
|
||||
mock_auth_config, context
|
||||
)
|
||||
|
||||
async def test_save_credential_no_service(
|
||||
self, mock_invocation_context, mock_auth_config
|
||||
):
|
||||
"""Test save_credential when credential service is not available."""
|
||||
mock_invocation_context.credential_service = None
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Credential service is not initialized"
|
||||
):
|
||||
await context.save_credential(mock_auth_config)
|
||||
|
||||
async def test_load_credential_with_service(
|
||||
self, mock_invocation_context, mock_auth_config, mock_auth_credential
|
||||
):
|
||||
"""Test load_credential when credential service is available."""
|
||||
credential_service = AsyncMock()
|
||||
credential_service.load_credential.return_value = mock_auth_credential
|
||||
mock_invocation_context.credential_service = credential_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
result = await context.load_credential(mock_auth_config)
|
||||
|
||||
credential_service.load_credential.assert_called_once_with(
|
||||
mock_auth_config, context
|
||||
)
|
||||
assert result == mock_auth_credential
|
||||
|
||||
async def test_load_credential_no_service(
|
||||
self, mock_invocation_context, mock_auth_config
|
||||
):
|
||||
"""Test load_credential when credential service is not available."""
|
||||
mock_invocation_context.credential_service = None
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Credential service is not initialized"
|
||||
):
|
||||
await context.load_credential(mock_auth_config)
|
||||
|
||||
|
||||
class TestContextGetAuthResponse:
|
||||
"""Test get_auth_response method in Context."""
|
||||
|
||||
def test_get_auth_response(self, mock_invocation_context, mock_auth_config):
|
||||
"""Test get_auth_response method."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with patch.object(
|
||||
auth_handler, "AuthHandler", autospec=True
|
||||
) as mock_auth_handler:
|
||||
mock_handler_instance = mock_auth_handler.return_value
|
||||
mock_handler_instance.get_auth_response.return_value = "auth-response"
|
||||
|
||||
result = context.get_auth_response(mock_auth_config)
|
||||
|
||||
mock_auth_handler.assert_called_once_with(mock_auth_config)
|
||||
mock_handler_instance.get_auth_response.assert_called_once_with(
|
||||
context.state
|
||||
)
|
||||
assert result == "auth-response"
|
||||
|
||||
|
||||
class TestContextRequestCredential:
|
||||
"""Test request_credential method in Context."""
|
||||
|
||||
def test_request_credential_with_function_call_id(
|
||||
self, mock_invocation_context, mock_auth_config
|
||||
):
|
||||
"""Test request_credential when function_call_id is set."""
|
||||
context = Context(
|
||||
mock_invocation_context,
|
||||
function_call_id="test-function-call-id",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
auth_handler, "AuthHandler", autospec=True
|
||||
) as mock_auth_handler:
|
||||
mock_handler_instance = mock_auth_handler.return_value
|
||||
mock_handler_instance.generate_auth_request.return_value = "auth-request"
|
||||
|
||||
context.request_credential(mock_auth_config)
|
||||
|
||||
mock_auth_handler.assert_called_once_with(mock_auth_config)
|
||||
mock_handler_instance.generate_auth_request.assert_called_once()
|
||||
assert (
|
||||
context.actions.requested_auth_configs["test-function-call-id"]
|
||||
== "auth-request"
|
||||
)
|
||||
|
||||
def test_request_credential_without_function_call_id_raises(
|
||||
self, mock_invocation_context, mock_auth_config
|
||||
):
|
||||
"""Test request_credential raises ValueError when no function_call_id."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="request_credential requires function_call_id",
|
||||
):
|
||||
context.request_credential(mock_auth_config)
|
||||
|
||||
|
||||
class TestContextRequestConfirmation:
|
||||
"""Test request_confirmation method in Context."""
|
||||
|
||||
def test_request_confirmation_with_function_call_id(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test request_confirmation when function_call_id is set."""
|
||||
context = Context(
|
||||
mock_invocation_context,
|
||||
function_call_id="test-function-call-id",
|
||||
)
|
||||
|
||||
context.request_confirmation(
|
||||
hint="Please confirm", payload={"action": "delete"}
|
||||
)
|
||||
|
||||
confirmation = context.actions.requested_tool_confirmations[
|
||||
"test-function-call-id"
|
||||
]
|
||||
assert confirmation.hint == "Please confirm"
|
||||
assert confirmation.payload == {"action": "delete"}
|
||||
|
||||
def test_request_confirmation_with_only_hint(self, mock_invocation_context):
|
||||
"""Test request_confirmation with only hint provided."""
|
||||
context = Context(
|
||||
mock_invocation_context,
|
||||
function_call_id="test-function-call-id",
|
||||
)
|
||||
|
||||
context.request_confirmation(hint="Confirm this action")
|
||||
|
||||
confirmation = context.actions.requested_tool_confirmations[
|
||||
"test-function-call-id"
|
||||
]
|
||||
assert confirmation.hint == "Confirm this action"
|
||||
assert confirmation.payload is None
|
||||
|
||||
def test_request_confirmation_without_function_call_id_raises(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test request_confirmation raises ValueError when no function_call_id."""
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="request_confirmation requires function_call_id",
|
||||
):
|
||||
context.request_confirmation()
|
||||
|
||||
|
||||
class TestContextMemoryMethods:
|
||||
"""Test memory methods in Context."""
|
||||
|
||||
async def test_add_session_to_memory_success(self, mock_invocation_context):
|
||||
"""Test that add_session_to_memory calls the memory service correctly."""
|
||||
memory_service = AsyncMock()
|
||||
mock_invocation_context.memory_service = memory_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
await context.add_session_to_memory()
|
||||
|
||||
memory_service.add_session_to_memory.assert_called_once_with(
|
||||
mock_invocation_context.session
|
||||
)
|
||||
|
||||
async def test_add_session_to_memory_no_service_raises(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test that add_session_to_memory raises ValueError when memory service is None."""
|
||||
mock_invocation_context.memory_service = None
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"Cannot add session to memory: memory service is not available\."
|
||||
),
|
||||
):
|
||||
await context.add_session_to_memory()
|
||||
|
||||
async def test_search_memory_success(self, mock_invocation_context, mocker):
|
||||
"""Test that search_memory calls the memory service correctly."""
|
||||
memory_service = AsyncMock()
|
||||
mock_search_response = mocker.create_autospec(
|
||||
SearchMemoryResponse, instance=True
|
||||
)
|
||||
memory_service.search_memory.return_value = mock_search_response
|
||||
mock_invocation_context.memory_service = memory_service
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
result = await context.search_memory("test query")
|
||||
|
||||
memory_service.search_memory.assert_called_once_with(
|
||||
app_name="test-app",
|
||||
user_id="test-user",
|
||||
query="test query",
|
||||
)
|
||||
assert result == mock_search_response
|
||||
|
||||
async def test_search_memory_no_service_raises(self, mock_invocation_context):
|
||||
"""Test that search_memory raises ValueError when memory service is None."""
|
||||
mock_invocation_context.memory_service = None
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(ValueError, match="Memory service is not available."):
|
||||
await context.search_memory("test query")
|
||||
|
||||
async def test_add_events_to_memory_success(self, mock_invocation_context):
|
||||
"""Test that add_events_to_memory calls the memory service correctly."""
|
||||
memory_service = AsyncMock()
|
||||
mock_invocation_context.memory_service = memory_service
|
||||
test_event = MagicMock()
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
await context.add_events_to_memory(
|
||||
events=[test_event],
|
||||
custom_metadata={"ttl": "6000s"},
|
||||
)
|
||||
|
||||
memory_service.add_events_to_memory.assert_called_once_with(
|
||||
app_name=mock_invocation_context.session.app_name,
|
||||
user_id=mock_invocation_context.session.user_id,
|
||||
session_id=mock_invocation_context.session.id,
|
||||
events=[test_event],
|
||||
custom_metadata={"ttl": "6000s"},
|
||||
)
|
||||
|
||||
async def test_add_events_to_memory_no_service_raises(
|
||||
self, mock_invocation_context
|
||||
):
|
||||
"""Test that add_events_to_memory raises ValueError when no service."""
|
||||
mock_invocation_context.memory_service = None
|
||||
|
||||
context = Context(mock_invocation_context)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Cannot add events to memory: memory service is not available\.",
|
||||
):
|
||||
await context.add_events_to_memory(events=[MagicMock()])
|
||||
Reference in New Issue
Block a user