mirror of
https://github.com/EvolutionAPI/adk-python.git
synced 2025-12-18 11:22:22 -06:00
Moves unittests to root folder and adds github action to run unit tests. (#72)
* Move unit tests to root package. * Adds deps to "test" extra, and mark two broken tests in tests/unittests/auth/test_auth_handler.py * Adds github workflow * minor fix in lite_llm.py for python 3.9. * format pyproject.toml
This commit is contained in:
Jack Sun
GitHub
14
tests/unittests/agents/__init__.py
Normal file
14
tests/unittests/agents/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
407
tests/unittests/agents/test_base_agent.py
Normal file
407
tests/unittests/agents/test_base_agent.py
Normal file
@@ -0,0 +1,407 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Testings for the BaseAgent."""
|
||||
|
||||
from typing import AsyncGenerator
|
||||
from typing import Optional
|
||||
from typing import Union
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.events import Event
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.genai import types
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
def _before_agent_callback_noop(callback_context: CallbackContext) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _before_agent_callback_bypass_agent(
|
||||
callback_context: CallbackContext,
|
||||
) -> types.Content:
|
||||
return types.Content(parts=[types.Part(text='agent run is bypassed.')])
|
||||
|
||||
|
||||
def _after_agent_callback_noop(callback_context: CallbackContext) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _after_agent_callback_append_agent_reply(
|
||||
callback_context: CallbackContext,
|
||||
) -> types.Content:
|
||||
return types.Content(
|
||||
parts=[types.Part(text='Agent reply from after agent callback.')]
|
||||
)
|
||||
|
||||
|
||||
class _IncompleteAgent(BaseAgent):
|
||||
pass
|
||||
|
||||
|
||||
class _TestingAgent(BaseAgent):
|
||||
|
||||
@override
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
branch=ctx.branch,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(parts=[types.Part(text='Hello, world!')]),
|
||||
)
|
||||
|
||||
@override
|
||||
async def _run_live_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
branch=ctx.branch,
|
||||
content=types.Content(parts=[types.Part(text='Hello, live!')]),
|
||||
)
|
||||
|
||||
|
||||
def _create_parent_invocation_context(
|
||||
test_name: str, agent: BaseAgent, branch: Optional[str] = None
|
||||
) -> InvocationContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
return InvocationContext(
|
||||
invocation_id=f'{test_name}_invocation_id',
|
||||
branch=branch,
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
|
||||
def test_invalid_agent_name():
|
||||
with pytest.raises(ValueError):
|
||||
_ = _TestingAgent(name='not an identifier')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(request: pytest.FixtureRequest):
|
||||
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
|
||||
events = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].author == agent.name
|
||||
assert events[0].content.parts[0].text == 'Hello, world!'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_with_branch(request: pytest.FixtureRequest):
|
||||
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent, branch='parent_branch'
|
||||
)
|
||||
|
||||
events = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].author == agent.name
|
||||
assert events[0].content.parts[0].text == 'Hello, world!'
|
||||
assert events[0].branch.endswith(agent.name)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_before_agent_callback_noop(
|
||||
request: pytest.FixtureRequest,
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> Union[types.Content, None]:
|
||||
# Arrange
|
||||
agent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
before_agent_callback=_before_agent_callback_noop,
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
|
||||
spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
|
||||
|
||||
# Act
|
||||
_ = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
# Assert
|
||||
spy_before_agent_callback.assert_called_once()
|
||||
_, kwargs = spy_before_agent_callback.call_args
|
||||
assert 'callback_context' in kwargs
|
||||
assert isinstance(kwargs['callback_context'], CallbackContext)
|
||||
|
||||
spy_run_async_impl.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_before_agent_callback_bypass_agent(
|
||||
request: pytest.FixtureRequest,
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
):
|
||||
# Arrange
|
||||
agent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
before_agent_callback=_before_agent_callback_bypass_agent,
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
spy_run_async_impl = mocker.spy(agent, BaseAgent._run_async_impl.__name__)
|
||||
spy_before_agent_callback = mocker.spy(agent, 'before_agent_callback')
|
||||
|
||||
# Act
|
||||
events = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
# Assert
|
||||
spy_before_agent_callback.assert_called_once()
|
||||
spy_run_async_impl.assert_not_called()
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].content.parts[0].text == 'agent run is bypassed.'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_after_agent_callback_noop(
|
||||
request: pytest.FixtureRequest,
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
):
|
||||
# Arrange
|
||||
agent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
after_agent_callback=_after_agent_callback_noop,
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
spy_after_agent_callback = mocker.spy(agent, 'after_agent_callback')
|
||||
|
||||
# Act
|
||||
events = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
# Assert
|
||||
spy_after_agent_callback.assert_called_once()
|
||||
_, kwargs = spy_after_agent_callback.call_args
|
||||
assert 'callback_context' in kwargs
|
||||
assert isinstance(kwargs['callback_context'], CallbackContext)
|
||||
assert len(events) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_after_agent_callback_append_reply(
|
||||
request: pytest.FixtureRequest,
|
||||
):
|
||||
# Arrange
|
||||
agent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
after_agent_callback=_after_agent_callback_append_agent_reply,
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
|
||||
# Act
|
||||
events = [e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
# Assert
|
||||
assert len(events) == 2
|
||||
assert events[1].author == agent.name
|
||||
assert (
|
||||
events[1].content.parts[0].text
|
||||
== 'Agent reply from after agent callback.'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_incomplete_agent(request: pytest.FixtureRequest):
|
||||
agent = _IncompleteAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
[e async for e in agent.run_async(parent_ctx)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_live(request: pytest.FixtureRequest):
|
||||
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
|
||||
events = [e async for e in agent.run_live(parent_ctx)]
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].author == agent.name
|
||||
assert events[0].content.parts[0].text == 'Hello, live!'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_live_with_branch(request: pytest.FixtureRequest):
|
||||
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent, branch='parent_branch'
|
||||
)
|
||||
|
||||
events = [e async for e in agent.run_live(parent_ctx)]
|
||||
|
||||
assert len(events) == 1
|
||||
assert events[0].author == agent.name
|
||||
assert events[0].content.parts[0].text == 'Hello, live!'
|
||||
assert events[0].branch.endswith(agent.name)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_live_incomplete_agent(request: pytest.FixtureRequest):
|
||||
agent = _IncompleteAgent(name=f'{request.function.__name__}_test_agent')
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, agent
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
[e async for e in agent.run_live(parent_ctx)]
|
||||
|
||||
|
||||
def test_set_parent_agent_for_sub_agents(request: pytest.FixtureRequest):
|
||||
sub_agents: list[BaseAgent] = [
|
||||
_TestingAgent(name=f'{request.function.__name__}_sub_agent_1'),
|
||||
_TestingAgent(name=f'{request.function.__name__}_sub_agent_2'),
|
||||
]
|
||||
parent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent',
|
||||
sub_agents=sub_agents,
|
||||
)
|
||||
|
||||
for sub_agent in sub_agents:
|
||||
assert sub_agent.parent_agent == parent
|
||||
|
||||
|
||||
def test_find_agent(request: pytest.FixtureRequest):
|
||||
grand_sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_1'
|
||||
)
|
||||
grand_sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_2'
|
||||
)
|
||||
sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_1',
|
||||
sub_agents=[grand_sub_agent_1],
|
||||
)
|
||||
sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_2',
|
||||
sub_agents=[grand_sub_agent_2],
|
||||
)
|
||||
parent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent',
|
||||
sub_agents=[sub_agent_1, sub_agent_2],
|
||||
)
|
||||
|
||||
assert parent.find_agent(parent.name) == parent
|
||||
assert parent.find_agent(sub_agent_1.name) == sub_agent_1
|
||||
assert parent.find_agent(sub_agent_2.name) == sub_agent_2
|
||||
assert parent.find_agent(grand_sub_agent_1.name) == grand_sub_agent_1
|
||||
assert parent.find_agent(grand_sub_agent_2.name) == grand_sub_agent_2
|
||||
assert sub_agent_1.find_agent(grand_sub_agent_1.name) == grand_sub_agent_1
|
||||
assert sub_agent_1.find_agent(grand_sub_agent_2.name) is None
|
||||
assert sub_agent_2.find_agent(grand_sub_agent_1.name) is None
|
||||
assert sub_agent_2.find_agent(sub_agent_2.name) == sub_agent_2
|
||||
assert parent.find_agent('not_exist') is None
|
||||
|
||||
|
||||
def test_find_sub_agent(request: pytest.FixtureRequest):
|
||||
grand_sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_1'
|
||||
)
|
||||
grand_sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_2'
|
||||
)
|
||||
sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_1',
|
||||
sub_agents=[grand_sub_agent_1],
|
||||
)
|
||||
sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_2',
|
||||
sub_agents=[grand_sub_agent_2],
|
||||
)
|
||||
parent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent',
|
||||
sub_agents=[sub_agent_1, sub_agent_2],
|
||||
)
|
||||
|
||||
assert parent.find_sub_agent(sub_agent_1.name) == sub_agent_1
|
||||
assert parent.find_sub_agent(sub_agent_2.name) == sub_agent_2
|
||||
assert parent.find_sub_agent(grand_sub_agent_1.name) == grand_sub_agent_1
|
||||
assert parent.find_sub_agent(grand_sub_agent_2.name) == grand_sub_agent_2
|
||||
assert sub_agent_1.find_sub_agent(grand_sub_agent_1.name) == grand_sub_agent_1
|
||||
assert sub_agent_1.find_sub_agent(grand_sub_agent_2.name) is None
|
||||
assert sub_agent_2.find_sub_agent(grand_sub_agent_1.name) is None
|
||||
assert sub_agent_2.find_sub_agent(grand_sub_agent_2.name) == grand_sub_agent_2
|
||||
assert parent.find_sub_agent(parent.name) is None
|
||||
assert parent.find_sub_agent('not_exist') is None
|
||||
|
||||
|
||||
def test_root_agent(request: pytest.FixtureRequest):
|
||||
grand_sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_1'
|
||||
)
|
||||
grand_sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}__grand_sub_agent_2'
|
||||
)
|
||||
sub_agent_1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_1',
|
||||
sub_agents=[grand_sub_agent_1],
|
||||
)
|
||||
sub_agent_2 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_sub_agent_2',
|
||||
sub_agents=[grand_sub_agent_2],
|
||||
)
|
||||
parent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent',
|
||||
sub_agents=[sub_agent_1, sub_agent_2],
|
||||
)
|
||||
|
||||
assert parent.root_agent == parent
|
||||
assert sub_agent_1.root_agent == parent
|
||||
assert sub_agent_2.root_agent == parent
|
||||
assert grand_sub_agent_1.root_agent == parent
|
||||
assert grand_sub_agent_2.root_agent == parent
|
||||
|
||||
|
||||
def test_set_parent_agent_for_sub_agent_twice(
|
||||
request: pytest.FixtureRequest,
|
||||
):
|
||||
sub_agent = _TestingAgent(name=f'{request.function.__name__}_sub_agent')
|
||||
_ = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent_1',
|
||||
sub_agents=[sub_agent],
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
_ = _TestingAgent(
|
||||
name=f'{request.function.__name__}_parent_2',
|
||||
sub_agents=[sub_agent],
|
||||
)
|
||||
191
tests/unittests/agents/test_langgraph_agent.py
Normal file
191
tests/unittests/agents/test_langgraph_agent.py
Normal file
@@ -0,0 +1,191 @@
|
||||
# Copyright 2025 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 MagicMock
|
||||
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.langgraph_agent import LangGraphAgent
|
||||
from google.adk.events import Event
|
||||
from google.genai import types
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langchain_core.messages import SystemMessage
|
||||
from langgraph.graph.graph import CompiledGraph
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"checkpointer_value, events_list, expected_messages",
|
||||
[
|
||||
(
|
||||
MagicMock(),
|
||||
[
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="user",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="test prompt")],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="root_agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[types.Part.from_text(text="(some delegation)")],
|
||||
),
|
||||
),
|
||||
],
|
||||
[
|
||||
SystemMessage(content="test system prompt"),
|
||||
HumanMessage(content="test prompt"),
|
||||
],
|
||||
),
|
||||
(
|
||||
None,
|
||||
[
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="user",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="user prompt 1")],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="root_agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
types.Part.from_text(text="root agent response")
|
||||
],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="weather_agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
types.Part.from_text(text="weather agent response")
|
||||
],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="user",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="user prompt 2")],
|
||||
),
|
||||
),
|
||||
],
|
||||
[
|
||||
SystemMessage(content="test system prompt"),
|
||||
HumanMessage(content="user prompt 1"),
|
||||
AIMessage(content="weather agent response"),
|
||||
HumanMessage(content="user prompt 2"),
|
||||
],
|
||||
),
|
||||
(
|
||||
MagicMock(),
|
||||
[
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="user",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="user prompt 1")],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="root_agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
types.Part.from_text(text="root agent response")
|
||||
],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="weather_agent",
|
||||
content=types.Content(
|
||||
role="model",
|
||||
parts=[
|
||||
types.Part.from_text(text="weather agent response")
|
||||
],
|
||||
),
|
||||
),
|
||||
Event(
|
||||
invocation_id="test_invocation_id",
|
||||
author="user",
|
||||
content=types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text="user prompt 2")],
|
||||
),
|
||||
),
|
||||
],
|
||||
[
|
||||
SystemMessage(content="test system prompt"),
|
||||
HumanMessage(content="user prompt 2"),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_langgraph_agent(
|
||||
checkpointer_value, events_list, expected_messages
|
||||
):
|
||||
mock_graph = MagicMock(spec=CompiledGraph)
|
||||
mock_graph_state = MagicMock()
|
||||
mock_graph_state.values = {}
|
||||
mock_graph.get_state.return_value = mock_graph_state
|
||||
|
||||
mock_graph.checkpointer = checkpointer_value
|
||||
mock_graph.invoke.return_value = {
|
||||
"messages": [AIMessage(content="test response")]
|
||||
}
|
||||
|
||||
mock_parent_context = MagicMock(spec=InvocationContext)
|
||||
mock_session = MagicMock()
|
||||
mock_parent_context.session = mock_session
|
||||
mock_parent_context.branch = "parent_agent"
|
||||
mock_parent_context.end_invocation = False
|
||||
mock_session.events = events_list
|
||||
mock_parent_context.invocation_id = "test_invocation_id"
|
||||
mock_parent_context.model_copy.return_value = mock_parent_context
|
||||
|
||||
weather_agent = LangGraphAgent(
|
||||
name="weather_agent",
|
||||
description="A agent that answers weather questions",
|
||||
instruction="test system prompt",
|
||||
graph=mock_graph,
|
||||
)
|
||||
|
||||
result_event = None
|
||||
async for event in weather_agent.run_async(mock_parent_context):
|
||||
result_event = event
|
||||
|
||||
assert result_event.author == "weather_agent"
|
||||
assert result_event.content.parts[0].text == "test response"
|
||||
|
||||
mock_graph.invoke.assert_called_once()
|
||||
mock_graph.invoke.assert_called_with(
|
||||
{"messages": expected_messages},
|
||||
{"configurable": {"thread_id": mock_session.id}},
|
||||
)
|
||||
138
tests/unittests/agents/test_llm_agent_callbacks.py
Normal file
138
tests/unittests/agents/test_llm_agent_callbacks.py
Normal file
@@ -0,0 +1,138 @@
|
||||
# Copyright 2025 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 typing import Any
|
||||
from typing import Optional
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.agents.llm_agent import Agent
|
||||
from google.adk.models import LlmRequest
|
||||
from google.adk.models import LlmResponse
|
||||
from google.genai import types
|
||||
from pydantic import BaseModel
|
||||
import pytest
|
||||
|
||||
from .. import utils
|
||||
|
||||
|
||||
class MockBeforeModelCallback(BaseModel):
|
||||
mock_response: str
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
callback_context: CallbackContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> LlmResponse:
|
||||
return LlmResponse(
|
||||
content=utils.ModelContent(
|
||||
[types.Part.from_text(text=self.mock_response)]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MockAfterModelCallback(BaseModel):
|
||||
mock_response: str
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
callback_context: CallbackContext,
|
||||
llm_response: LlmResponse,
|
||||
) -> LlmResponse:
|
||||
return LlmResponse(
|
||||
content=utils.ModelContent(
|
||||
[types.Part.from_text(text=self.mock_response)]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def noop_callback(**kwargs) -> Optional[LlmResponse]:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_before_model_callback():
|
||||
responses = ['model_response']
|
||||
mock_model = utils.MockModel.create(responses=responses)
|
||||
agent = Agent(
|
||||
name='root_agent',
|
||||
model=mock_model,
|
||||
before_model_callback=MockBeforeModelCallback(
|
||||
mock_response='before_model_callback'
|
||||
),
|
||||
)
|
||||
|
||||
runner = utils.TestInMemoryRunner(agent)
|
||||
assert utils.simplify_events(
|
||||
await runner.run_async_with_new_session('test')
|
||||
) == [
|
||||
('root_agent', 'before_model_callback'),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_before_model_callback_noop():
|
||||
responses = ['model_response']
|
||||
mock_model = utils.MockModel.create(responses=responses)
|
||||
agent = Agent(
|
||||
name='root_agent',
|
||||
model=mock_model,
|
||||
before_model_callback=noop_callback,
|
||||
)
|
||||
|
||||
runner = utils.TestInMemoryRunner(agent)
|
||||
assert utils.simplify_events(
|
||||
await runner.run_async_with_new_session('test')
|
||||
) == [
|
||||
('root_agent', 'model_response'),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_before_model_callback_end():
|
||||
responses = ['model_response']
|
||||
mock_model = utils.MockModel.create(responses=responses)
|
||||
agent = Agent(
|
||||
name='root_agent',
|
||||
model=mock_model,
|
||||
before_model_callback=MockBeforeModelCallback(
|
||||
mock_response='before_model_callback',
|
||||
),
|
||||
)
|
||||
|
||||
runner = utils.TestInMemoryRunner(agent)
|
||||
assert utils.simplify_events(
|
||||
await runner.run_async_with_new_session('test')
|
||||
) == [
|
||||
('root_agent', 'before_model_callback'),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_after_model_callback():
|
||||
responses = ['model_response']
|
||||
mock_model = utils.MockModel.create(responses=responses)
|
||||
agent = Agent(
|
||||
name='root_agent',
|
||||
model=mock_model,
|
||||
after_model_callback=MockAfterModelCallback(
|
||||
mock_response='after_model_callback'
|
||||
),
|
||||
)
|
||||
|
||||
runner = utils.TestInMemoryRunner(agent)
|
||||
assert utils.simplify_events(
|
||||
await runner.run_async_with_new_session('test')
|
||||
) == [
|
||||
('root_agent', 'after_model_callback'),
|
||||
]
|
||||
231
tests/unittests/agents/test_llm_agent_fields.py
Normal file
231
tests/unittests/agents/test_llm_agent_fields.py
Normal file
@@ -0,0 +1,231 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Unit tests for canonical_xxx fields in LlmAgent."""
|
||||
|
||||
from typing import Any
|
||||
from typing import Optional
|
||||
|
||||
from google.adk.agents.callback_context import CallbackContext
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.llm_agent import LlmAgent
|
||||
from google.adk.agents.loop_agent import LoopAgent
|
||||
from google.adk.agents.readonly_context import ReadonlyContext
|
||||
from google.adk.models.llm_request import LlmRequest
|
||||
from google.adk.models.registry import LLMRegistry
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.genai import types
|
||||
from pydantic import BaseModel
|
||||
import pytest
|
||||
|
||||
|
||||
def _create_readonly_context(
|
||||
agent: LlmAgent, state: Optional[dict[str, Any]] = None
|
||||
) -> ReadonlyContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = session_service.create_session(
|
||||
app_name='test_app', user_id='test_user', state=state
|
||||
)
|
||||
invocation_context = InvocationContext(
|
||||
invocation_id='test_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
return ReadonlyContext(invocation_context)
|
||||
|
||||
|
||||
def test_canonical_model_empty():
|
||||
agent = LlmAgent(name='test_agent')
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_ = agent.canonical_model
|
||||
|
||||
|
||||
def test_canonical_model_str():
|
||||
agent = LlmAgent(name='test_agent', model='gemini-pro')
|
||||
|
||||
assert agent.canonical_model.model == 'gemini-pro'
|
||||
|
||||
|
||||
def test_canonical_model_llm():
|
||||
llm = LLMRegistry.new_llm('gemini-pro')
|
||||
agent = LlmAgent(name='test_agent', model=llm)
|
||||
|
||||
assert agent.canonical_model == llm
|
||||
|
||||
|
||||
def test_canonical_model_inherit():
|
||||
sub_agent = LlmAgent(name='sub_agent')
|
||||
parent_agent = LlmAgent(
|
||||
name='parent_agent', model='gemini-pro', sub_agents=[sub_agent]
|
||||
)
|
||||
|
||||
assert sub_agent.canonical_model == parent_agent.canonical_model
|
||||
|
||||
|
||||
def test_canonical_instruction_str():
|
||||
agent = LlmAgent(name='test_agent', instruction='instruction')
|
||||
ctx = _create_readonly_context(agent)
|
||||
|
||||
assert agent.canonical_instruction(ctx) == 'instruction'
|
||||
|
||||
|
||||
def test_canonical_instruction():
|
||||
def _instruction_provider(ctx: ReadonlyContext) -> str:
|
||||
return f'instruction: {ctx.state["state_var"]}'
|
||||
|
||||
agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
|
||||
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
|
||||
|
||||
assert agent.canonical_instruction(ctx) == 'instruction: state_value'
|
||||
|
||||
|
||||
def test_canonical_global_instruction_str():
|
||||
agent = LlmAgent(name='test_agent', global_instruction='global instruction')
|
||||
ctx = _create_readonly_context(agent)
|
||||
|
||||
assert agent.canonical_global_instruction(ctx) == 'global instruction'
|
||||
|
||||
|
||||
def test_canonical_global_instruction():
|
||||
def _global_instruction_provider(ctx: ReadonlyContext) -> str:
|
||||
return f'global instruction: {ctx.state["state_var"]}'
|
||||
|
||||
agent = LlmAgent(
|
||||
name='test_agent', global_instruction=_global_instruction_provider
|
||||
)
|
||||
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
|
||||
|
||||
assert (
|
||||
agent.canonical_global_instruction(ctx)
|
||||
== 'global instruction: state_value'
|
||||
)
|
||||
|
||||
|
||||
def test_output_schema_will_disable_transfer(caplog: pytest.LogCaptureFixture):
|
||||
with caplog.at_level('WARNING'):
|
||||
|
||||
class Schema(BaseModel):
|
||||
pass
|
||||
|
||||
agent = LlmAgent(
|
||||
name='test_agent',
|
||||
output_schema=Schema,
|
||||
)
|
||||
|
||||
# Transfer is automatically disabled
|
||||
assert agent.disallow_transfer_to_parent
|
||||
assert agent.disallow_transfer_to_peers
|
||||
assert (
|
||||
'output_schema cannot co-exist with agent transfer configurations.'
|
||||
in caplog.text
|
||||
)
|
||||
|
||||
|
||||
def test_output_schema_with_sub_agents_will_throw():
|
||||
class Schema(BaseModel):
|
||||
pass
|
||||
|
||||
sub_agent = LlmAgent(
|
||||
name='sub_agent',
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
output_schema=Schema,
|
||||
sub_agents=[sub_agent],
|
||||
)
|
||||
|
||||
|
||||
def test_output_schema_with_tools_will_throw():
|
||||
class Schema(BaseModel):
|
||||
pass
|
||||
|
||||
def _a_tool():
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
output_schema=Schema,
|
||||
tools=[_a_tool],
|
||||
)
|
||||
|
||||
|
||||
def test_before_model_callback():
|
||||
def _before_model_callback(
|
||||
callback_context: CallbackContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
agent = LlmAgent(
|
||||
name='test_agent', before_model_callback=_before_model_callback
|
||||
)
|
||||
|
||||
# TODO: add more logic assertions later.
|
||||
assert agent.before_model_callback is not None
|
||||
|
||||
|
||||
def test_validate_generate_content_config_thinking_config_throw():
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
thinking_config=types.ThinkingConfig()
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_tools_throw():
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
tools=[types.Tool(function_declarations=[])]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_system_instruction_throw():
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
system_instruction='system instruction'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_validate_generate_content_config_response_schema_throw():
|
||||
class Schema(BaseModel):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_ = LlmAgent(
|
||||
name='test_agent',
|
||||
generate_content_config=types.GenerateContentConfig(
|
||||
response_schema=Schema
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_allow_transfer_by_default():
|
||||
sub_agent = LlmAgent(name='sub_agent')
|
||||
agent = LlmAgent(name='test_agent', sub_agents=[sub_agent])
|
||||
|
||||
assert not agent.disallow_transfer_to_parent
|
||||
assert not agent.disallow_transfer_to_peers
|
||||
136
tests/unittests/agents/test_loop_agent.py
Normal file
136
tests/unittests/agents/test_loop_agent.py
Normal file
@@ -0,0 +1,136 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Testings for the SequentialAgent."""
|
||||
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.loop_agent import LoopAgent
|
||||
from google.adk.events import Event
|
||||
from google.adk.events import EventActions
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.genai import types
|
||||
import pytest
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
class _TestingAgent(BaseAgent):
|
||||
|
||||
@override
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, async {self.name}!')]
|
||||
),
|
||||
)
|
||||
|
||||
@override
|
||||
async def _run_live_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, live {self.name}!')]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _TestingAgentWithEscalateAction(BaseAgent):
|
||||
|
||||
@override
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, async {self.name}!')]
|
||||
),
|
||||
actions=EventActions(escalate=True),
|
||||
)
|
||||
|
||||
|
||||
def _create_parent_invocation_context(
|
||||
test_name: str, agent: BaseAgent
|
||||
) -> InvocationContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
return InvocationContext(
|
||||
invocation_id=f'{test_name}_invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(request: pytest.FixtureRequest):
|
||||
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
|
||||
loop_agent = LoopAgent(
|
||||
name=f'{request.function.__name__}_test_loop_agent',
|
||||
max_iterations=2,
|
||||
sub_agents=[
|
||||
agent,
|
||||
],
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, loop_agent
|
||||
)
|
||||
events = [e async for e in loop_agent.run_async(parent_ctx)]
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].author == agent.name
|
||||
assert events[1].author == agent.name
|
||||
assert events[0].content.parts[0].text == f'Hello, async {agent.name}!'
|
||||
assert events[1].content.parts[0].text == f'Hello, async {agent.name}!'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_with_escalate_action(request: pytest.FixtureRequest):
|
||||
non_escalating_agent = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_non_escalating_agent'
|
||||
)
|
||||
escalating_agent = _TestingAgentWithEscalateAction(
|
||||
name=f'{request.function.__name__}_test_escalating_agent'
|
||||
)
|
||||
loop_agent = LoopAgent(
|
||||
name=f'{request.function.__name__}_test_loop_agent',
|
||||
sub_agents=[non_escalating_agent, escalating_agent],
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, loop_agent
|
||||
)
|
||||
events = [e async for e in loop_agent.run_async(parent_ctx)]
|
||||
|
||||
# Only two events are generated because the sub escalating_agent escalates.
|
||||
assert len(events) == 2
|
||||
assert events[0].author == non_escalating_agent.name
|
||||
assert events[1].author == escalating_agent.name
|
||||
assert events[0].content.parts[0].text == (
|
||||
f'Hello, async {non_escalating_agent.name}!'
|
||||
)
|
||||
assert events[1].content.parts[0].text == (
|
||||
f'Hello, async {escalating_agent.name}!'
|
||||
)
|
||||
92
tests/unittests/agents/test_parallel_agent.py
Normal file
92
tests/unittests/agents/test_parallel_agent.py
Normal file
@@ -0,0 +1,92 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Tests for the ParallelAgent."""
|
||||
|
||||
import asyncio
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.parallel_agent import ParallelAgent
|
||||
from google.adk.events import Event
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.genai import types
|
||||
import pytest
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
class _TestingAgent(BaseAgent):
|
||||
|
||||
delay: float = 0
|
||||
"""The delay before the agent generates an event."""
|
||||
|
||||
@override
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
await asyncio.sleep(self.delay)
|
||||
yield Event(
|
||||
author=self.name,
|
||||
branch=ctx.branch,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, async {self.name}!')]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _create_parent_invocation_context(
|
||||
test_name: str, agent: BaseAgent
|
||||
) -> InvocationContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
return InvocationContext(
|
||||
invocation_id=f'{test_name}_invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(request: pytest.FixtureRequest):
|
||||
agent1 = _TestingAgent(
|
||||
name=f'{request.function.__name__}_test_agent_1',
|
||||
delay=0.5,
|
||||
)
|
||||
agent2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
|
||||
parallel_agent = ParallelAgent(
|
||||
name=f'{request.function.__name__}_test_parallel_agent',
|
||||
sub_agents=[
|
||||
agent1,
|
||||
agent2,
|
||||
],
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, parallel_agent
|
||||
)
|
||||
events = [e async for e in parallel_agent.run_async(parent_ctx)]
|
||||
|
||||
assert len(events) == 2
|
||||
# agent2 generates an event first, then agent1. Because they run in parallel
|
||||
# and agent1 has a delay.
|
||||
assert events[0].author == agent2.name
|
||||
assert events[1].author == agent1.name
|
||||
assert events[0].branch.endswith(agent2.name)
|
||||
assert events[1].branch.endswith(agent1.name)
|
||||
assert events[0].content.parts[0].text == f'Hello, async {agent2.name}!'
|
||||
assert events[1].content.parts[0].text == f'Hello, async {agent1.name}!'
|
||||
114
tests/unittests/agents/test_sequential_agent.py
Normal file
114
tests/unittests/agents/test_sequential_agent.py
Normal file
@@ -0,0 +1,114 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
"""Testings for the SequentialAgent."""
|
||||
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.agents.sequential_agent import SequentialAgent
|
||||
from google.adk.events import Event
|
||||
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
||||
from google.genai import types
|
||||
import pytest
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
class _TestingAgent(BaseAgent):
|
||||
|
||||
@override
|
||||
async def _run_async_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, async {self.name}!')]
|
||||
),
|
||||
)
|
||||
|
||||
@override
|
||||
async def _run_live_impl(
|
||||
self, ctx: InvocationContext
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
yield Event(
|
||||
author=self.name,
|
||||
invocation_id=ctx.invocation_id,
|
||||
content=types.Content(
|
||||
parts=[types.Part(text=f'Hello, live {self.name}!')]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _create_parent_invocation_context(
|
||||
test_name: str, agent: BaseAgent
|
||||
) -> InvocationContext:
|
||||
session_service = InMemorySessionService()
|
||||
session = session_service.create_session(
|
||||
app_name='test_app', user_id='test_user'
|
||||
)
|
||||
return InvocationContext(
|
||||
invocation_id=f'{test_name}_invocation_id',
|
||||
agent=agent,
|
||||
session=session,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(request: pytest.FixtureRequest):
|
||||
agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
|
||||
agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
|
||||
sequential_agent = SequentialAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
sub_agents=[
|
||||
agent_1,
|
||||
agent_2,
|
||||
],
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, sequential_agent
|
||||
)
|
||||
events = [e async for e in sequential_agent.run_async(parent_ctx)]
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].author == agent_1.name
|
||||
assert events[1].author == agent_2.name
|
||||
assert events[0].content.parts[0].text == f'Hello, async {agent_1.name}!'
|
||||
assert events[1].content.parts[0].text == f'Hello, async {agent_2.name}!'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_live(request: pytest.FixtureRequest):
|
||||
agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')
|
||||
agent_2 = _TestingAgent(name=f'{request.function.__name__}_test_agent_2')
|
||||
sequential_agent = SequentialAgent(
|
||||
name=f'{request.function.__name__}_test_agent',
|
||||
sub_agents=[
|
||||
agent_1,
|
||||
agent_2,
|
||||
],
|
||||
)
|
||||
parent_ctx = _create_parent_invocation_context(
|
||||
request.function.__name__, sequential_agent
|
||||
)
|
||||
events = [e async for e in sequential_agent.run_live(parent_ctx)]
|
||||
|
||||
assert len(events) == 2
|
||||
assert events[0].author == agent_1.name
|
||||
assert events[1].author == agent_2.name
|
||||
assert events[0].content.parts[0].text == f'Hello, live {agent_1.name}!'
|
||||
assert events[1].content.parts[0].text == f'Hello, live {agent_2.name}!'
|
||||
Reference in New Issue
Block a user