Skip to content

Commit 4255743

Browse files
committed
fix: support instructions provider for agents - handle sync cases and empty instructions
1 parent 41660e5 commit 4255743

File tree

2 files changed

+135
-4
lines changed

2 files changed

+135
-4
lines changed

typescript-sdk/integrations/adk-middleware/src/adk_middleware/adk_agent.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import time
77
import json
88
import asyncio
9+
import inspect
910
from datetime import datetime
1011

1112
from ag_ui.core import (
@@ -705,11 +706,25 @@ async def _start_background_execution(
705706

706707
if callable(current_instruction):
707708
# Handle instructions provider
708-
async def instruction_provider_wrapper(*args, **kwargs):
709-
original_instructions = await current_instruction(*args, **kwargs)
710-
return f"{original_instructions}\n\n{system_content}"
709+
if inspect.iscoroutinefunction(current_instruction):
710+
# Async instruction provider
711+
async def instruction_provider_wrapper_async(*args, **kwargs):
712+
instructions = system_content
713+
original_instructions = await current_instruction(*args, **kwargs) or ''
714+
if original_instructions:
715+
instructions = f"{original_instructions}\n\n{instructions}"
716+
return instructions
717+
new_instruction = instruction_provider_wrapper_async
718+
else:
719+
# Sync instruction provider
720+
def instruction_provider_wrapper_sync(*args, **kwargs):
721+
instructions = system_content
722+
original_instructions = current_instruction(*args, **kwargs) or ''
723+
if original_instructions:
724+
instructions = f"{original_instructions}\n\n{instructions}"
725+
return instructions
726+
new_instruction = instruction_provider_wrapper_sync
711727

712-
new_instruction = instruction_provider_wrapper
713728
logger.debug(
714729
f"Will wrap callable InstructionProvider and append SystemMessage: '{system_content[:100]}...'")
715730
else:

typescript-sdk/integrations/adk-middleware/tests/test_adk_agent.py

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,122 @@ async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
300300
assert agent_instruction == expected_instruction
301301
assert received_context is test_context
302302

303+
@pytest.mark.asyncio
304+
async def test_system_message_appended_to_instruction_provider_with_none(self):
305+
"""Test that SystemMessage as first message gets appended to agent instructions
306+
when they are set via instruction provider."""
307+
# Create an agent with initial instructions, but return None
308+
async def instruction_provider(context) -> str:
309+
return None
310+
311+
mock_agent = Agent(
312+
name="test_agent",
313+
instruction=instruction_provider
314+
)
315+
316+
adk_agent = ADKAgent(adk_agent=mock_agent, app_name="test_app", user_id="test_user")
317+
318+
# Create input with SystemMessage as first message
319+
system_input = RunAgentInput(
320+
thread_id="test_thread",
321+
run_id="test_run",
322+
messages=[
323+
SystemMessage(id="sys_1", role="system", content="Be very concise in responses."),
324+
UserMessage(id="msg_1", role="user", content="Hello")
325+
],
326+
context=[],
327+
state={},
328+
tools=[],
329+
forwarded_props={}
330+
)
331+
332+
# Mock the background execution to capture the modified agent
333+
captured_agent = None
334+
original_run_background = adk_agent._run_adk_in_background
335+
336+
async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
337+
nonlocal captured_agent
338+
captured_agent = adk_agent
339+
# Just put a completion event in the queue and return
340+
await event_queue.put(None)
341+
342+
with patch.object(adk_agent, '_run_adk_in_background', side_effect=mock_run_background):
343+
# Start execution to trigger agent modification
344+
execution = await adk_agent._start_background_execution(system_input)
345+
346+
# Wait briefly for the background task to start
347+
await asyncio.sleep(0.01)
348+
349+
# Verify the agent's instruction was wrapped correctly
350+
assert captured_agent is not None
351+
assert callable(captured_agent.instruction) is True
352+
353+
# No empty new lines should be added before the instructions
354+
expected_instruction = "Be very concise in responses."
355+
agent_instruction = await captured_agent.instruction({})
356+
assert agent_instruction == expected_instruction
357+
358+
@pytest.mark.asyncio
359+
async def test_system_message_appended_to_sync_instruction_provider(self):
360+
"""Test that SystemMessage as first message gets appended to agent instructions
361+
when they are set via sync instruction provider."""
362+
# Create an agent with initial instructions
363+
received_context = None
364+
365+
def instruction_provider(context) -> str:
366+
nonlocal received_context
367+
received_context = context
368+
return "You are a helpful assistant."
369+
370+
mock_agent = Agent(
371+
name="test_agent",
372+
instruction=instruction_provider
373+
)
374+
375+
adk_agent = ADKAgent(adk_agent=mock_agent, app_name="test_app", user_id="test_user")
376+
377+
# Create input with SystemMessage as first message
378+
system_input = RunAgentInput(
379+
thread_id="test_thread",
380+
run_id="test_run",
381+
messages=[
382+
SystemMessage(id="sys_1", role="system", content="Be very concise in responses."),
383+
UserMessage(id="msg_1", role="user", content="Hello")
384+
],
385+
context=[],
386+
state={},
387+
tools=[],
388+
forwarded_props={}
389+
)
390+
391+
# Mock the background execution to capture the modified agent
392+
captured_agent = None
393+
original_run_background = adk_agent._run_adk_in_background
394+
395+
async def mock_run_background(input, adk_agent, user_id, app_name, event_queue):
396+
nonlocal captured_agent
397+
captured_agent = adk_agent
398+
# Just put a completion event in the queue and return
399+
await event_queue.put(None)
400+
401+
with patch.object(adk_agent, '_run_adk_in_background', side_effect=mock_run_background):
402+
# Start execution to trigger agent modification
403+
execution = await adk_agent._start_background_execution(system_input)
404+
405+
# Wait briefly for the background task to start
406+
await asyncio.sleep(0.01)
407+
408+
# Verify agent was captured
409+
assert captured_agent is not None
410+
assert callable(captured_agent.instruction)
411+
412+
# Test that the context object received in instruction provider is the same
413+
test_context = {"test": "value"}
414+
expected_instruction = "You are a helpful assistant.\n\nBe very concise in responses."
415+
agent_instruction = captured_agent.instruction(test_context) # Note: no await for sync function
416+
assert agent_instruction == expected_instruction
417+
assert received_context is test_context
418+
303419
@pytest.mark.asyncio
304420
async def test_system_message_not_first_ignored(self):
305421
"""Test that SystemMessage not as first message is ignored."""

0 commit comments

Comments
 (0)