Await canonical instructions in test_llm_agent_fields.py

PiperOrigin-RevId: 758372807
This commit is contained in:
Selcuk Gun 2025-05-13 14:16:57 -07:00 committed by Copybara-Service
parent 82b8e304ec
commit 303af440ee

View File

@ -75,41 +75,45 @@ def test_canonical_model_inherit():
assert sub_agent.canonical_model == parent_agent.canonical_model
def test_canonical_instruction_str():
async def test_canonical_instruction_str():
agent = LlmAgent(name='test_agent', instruction='instruction')
ctx = _create_readonly_context(agent)
assert agent.canonical_instruction(ctx) == 'instruction'
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction'
def test_canonical_instruction():
async 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'
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction: state_value'
def test_async_canonical_instruction():
async def test_async_canonical_instruction():
async 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'
canonical_instruction = await agent.canonical_instruction(ctx)
assert canonical_instruction == 'instruction: state_value'
def test_canonical_global_instruction_str():
async 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'
canonical_global_instruction = await agent.canonical_global_instruction(ctx)
assert canonical_global_instruction == 'global instruction'
def test_canonical_global_instruction():
async def test_canonical_global_instruction():
def _global_instruction_provider(ctx: ReadonlyContext) -> str:
return f'global instruction: {ctx.state["state_var"]}'
@ -118,13 +122,11 @@ def test_canonical_global_instruction():
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
assert (
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_async_canonical_global_instruction():
async def test_async_canonical_global_instruction():
async def _global_instruction_provider(ctx: ReadonlyContext) -> str:
return f'global instruction: {ctx.state["state_var"]}'
@ -133,10 +135,8 @@ def test_async_canonical_global_instruction():
)
ctx = _create_readonly_context(agent, state={'state_var': 'state_value'})
assert (
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):