feat! Update session service interface to be async.

Also keep the sync version in the InMemorySessionService as create_session_sync() as a temporary migration option.

PiperOrigin-RevId: 759252188
This commit is contained in:
Google Team Member
2025-05-15 12:23:33 -07:00
committed by Copybara-Service
parent 5b3204c356
commit 1804ca39a6
23 changed files with 268 additions and 264 deletions

View File

@@ -15,7 +15,7 @@
"""Unit tests for canonical_xxx fields in LlmAgent."""
from typing import Any
from typing import Optional, cast
from typing import Optional
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.invocation_context import InvocationContext
@@ -30,11 +30,11 @@ from pydantic import BaseModel
import pytest
async def _create_readonly_context(
def _create_readonly_context(
agent: LlmAgent, state: Optional[dict[str, Any]] = None
) -> ReadonlyContext:
session_service = InMemorySessionService()
session = await session_service.create_session(
session = session_service.create_session(
app_name='test_app', user_id='test_user', state=state
)
invocation_context = InvocationContext(
@@ -77,7 +77,7 @@ def test_canonical_model_inherit():
async def test_canonical_instruction_str():
agent = LlmAgent(name='test_agent', instruction='instruction')
ctx = await _create_readonly_context(agent)
ctx = _create_readonly_context(agent)
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction'
@@ -88,9 +88,7 @@ async def test_canonical_instruction():
return f'instruction: {ctx.state["state_var"]}'
agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
ctx = await _create_readonly_context(
agent, state={'state_var': 'state_value'}
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction: state_value'
@@ -101,9 +99,7 @@ async def test_async_canonical_instruction():
return f'instruction: {ctx.state["state_var"]}'
agent = LlmAgent(name='test_agent', instruction=_instruction_provider)
ctx = await _create_readonly_context(
agent, state={'state_var': 'state_value'}
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction: state_value'
@@ -111,10 +107,10 @@ async def test_async_canonical_instruction():
async def test_canonical_global_instruction_str():
agent = LlmAgent(name='test_agent', global_instruction='global instruction')
ctx = await _create_readonly_context(agent)
ctx = _create_readonly_context(agent)
canonical_instruction = await agent.canonical_global_instruction(ctx)
assert canonical_instruction == 'global instruction'
canonical_global_instruction = await agent.canonical_global_instruction(ctx)
assert canonical_global_instruction == 'global instruction'
async def test_canonical_global_instruction():
@@ -124,9 +120,7 @@ async def test_canonical_global_instruction():
agent = LlmAgent(
name='test_agent', global_instruction=_global_instruction_provider
)
ctx = await _create_readonly_context(
agent, state={'state_var': 'state_value'}
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
canonical_global_instruction = await agent.canonical_global_instruction(ctx)
assert canonical_global_instruction == 'global instruction: state_value'
@@ -139,14 +133,10 @@ async def test_async_canonical_global_instruction():
agent = LlmAgent(
name='test_agent', global_instruction=_global_instruction_provider
)
ctx = await _create_readonly_context(
agent, state={'state_var': 'state_value'}
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
assert (
await agent.canonical_global_instruction(ctx)
== 'global instruction: state_value'
)
canonical_global_instruction = await agent.canonical_global_instruction(ctx)
assert canonical_global_instruction == 'global instruction: state_value'
def test_output_schema_will_disable_transfer(caplog: pytest.LogCaptureFixture):