mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
refactor: Extract helper function for llm request building and response processing
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 860212868
This commit is contained in:
Xiang (Sean) Zhou
committed by
Copybara-Service
Copybara-Service
parent
2380afd287
commit
753084fd46
@@ -54,7 +54,7 @@ class _AgentTransferLlmRequestProcessor(BaseLlmRequestProcessor):
|
||||
)
|
||||
|
||||
llm_request.append_instructions([
|
||||
_build_target_agents_instructions(
|
||||
_build_transfer_instructions(
|
||||
transfer_to_agent_tool.name,
|
||||
invocation_context.agent,
|
||||
transfer_targets,
|
||||
@@ -83,11 +83,24 @@ Agent description: {target_agent.description}
|
||||
line_break = '\n'
|
||||
|
||||
|
||||
def _build_target_agents_instructions(
|
||||
def _build_transfer_instructions(
|
||||
tool_name: str,
|
||||
agent: LlmAgent,
|
||||
target_agents: list[BaseAgent],
|
||||
agent: 'LlmAgent',
|
||||
target_agents: list['BaseAgent'],
|
||||
) -> str:
|
||||
"""Build instructions for agent transfer.
|
||||
|
||||
This function generates the instruction text that guides the LLM on how to
|
||||
use the transfer tool to delegate to other agents.
|
||||
|
||||
Args:
|
||||
tool_name: The name of the transfer tool (e.g., 'transfer_to_agent').
|
||||
agent: The current agent that may initiate transfers.
|
||||
target_agents: List of agents that can be transferred to.
|
||||
|
||||
Returns:
|
||||
Instruction text for the LLM about agent transfers.
|
||||
"""
|
||||
# Build list of available agent names for the NOTE
|
||||
# target_agents already includes parent agent if applicable,
|
||||
# so no need to add it again
|
||||
|
||||
@@ -69,6 +69,42 @@ DEFAULT_TASK_COMPLETION_DELAY = 1.0
|
||||
DEFAULT_ENABLE_CACHE_STATISTICS = False
|
||||
|
||||
|
||||
def _finalize_model_response_event(
|
||||
llm_request: LlmRequest,
|
||||
llm_response: LlmResponse,
|
||||
model_response_event: Event,
|
||||
) -> Event:
|
||||
"""Finalize and build the model response event from LLM response.
|
||||
|
||||
Merges the LLM response data into the model response event and
|
||||
populates function call IDs and long-running tool information.
|
||||
|
||||
Args:
|
||||
llm_request: The original LLM request.
|
||||
llm_response: The LLM response from the model.
|
||||
model_response_event: The base event to populate.
|
||||
|
||||
Returns:
|
||||
The finalized Event with LLM response data merged in.
|
||||
"""
|
||||
finalized_event = Event.model_validate({
|
||||
**model_response_event.model_dump(exclude_none=True),
|
||||
**llm_response.model_dump(exclude_none=True),
|
||||
})
|
||||
|
||||
if finalized_event.content:
|
||||
function_calls = finalized_event.get_function_calls()
|
||||
if function_calls:
|
||||
functions.populate_client_function_call_id(finalized_event)
|
||||
finalized_event.long_running_tool_ids = (
|
||||
functions.get_long_running_function_calls(
|
||||
function_calls, llm_request.tools_dict
|
||||
)
|
||||
)
|
||||
|
||||
return finalized_event
|
||||
|
||||
|
||||
class BaseLlmFlow(ABC):
|
||||
"""A basic flow that calls the LLM in a loop until a final response is generated.
|
||||
|
||||
@@ -941,22 +977,9 @@ class BaseLlmFlow(ABC):
|
||||
llm_response: LlmResponse,
|
||||
model_response_event: Event,
|
||||
) -> Event:
|
||||
model_response_event = Event.model_validate({
|
||||
**model_response_event.model_dump(exclude_none=True),
|
||||
**llm_response.model_dump(exclude_none=True),
|
||||
})
|
||||
|
||||
if model_response_event.content:
|
||||
function_calls = model_response_event.get_function_calls()
|
||||
if function_calls:
|
||||
functions.populate_client_function_call_id(model_response_event)
|
||||
model_response_event.long_running_tool_ids = (
|
||||
functions.get_long_running_function_calls(
|
||||
function_calls, llm_request.tools_dict
|
||||
)
|
||||
)
|
||||
|
||||
return model_response_event
|
||||
return _finalize_model_response_event(
|
||||
llm_request, llm_response, model_response_event
|
||||
)
|
||||
|
||||
async def _handle_control_event_flush(
|
||||
self, invocation_context: InvocationContext, llm_response: LlmResponse
|
||||
|
||||
@@ -29,55 +29,71 @@ from ...utils.output_schema_utils import can_use_output_schema_with_tools
|
||||
from ._base_llm_processor import BaseLlmRequestProcessor
|
||||
|
||||
|
||||
def _build_basic_request(
|
||||
invocation_context: InvocationContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
"""Populate basic LlmRequest fields from agent configuration.
|
||||
|
||||
Sets up model, config, output_schema, and live connect configuration
|
||||
based on the agent and run configuration.
|
||||
|
||||
Args:
|
||||
invocation_context: The invocation context containing agent and run config.
|
||||
llm_request: The LlmRequest to populate.
|
||||
"""
|
||||
agent = invocation_context.agent
|
||||
model = agent.canonical_model
|
||||
llm_request.model = model if isinstance(model, str) else model.model
|
||||
llm_request.config = (
|
||||
agent.generate_content_config.model_copy(deep=True)
|
||||
if agent.generate_content_config
|
||||
else types.GenerateContentConfig()
|
||||
)
|
||||
# Only set output_schema if no tools are specified. as of now, model don't
|
||||
# support output_schema and tools together. we have a workaround to support
|
||||
# both output_schema and tools at the same time. see
|
||||
# _output_schema_processor.py for details
|
||||
if agent.output_schema:
|
||||
if not agent.tools or can_use_output_schema_with_tools(model):
|
||||
llm_request.set_output_schema(agent.output_schema)
|
||||
|
||||
llm_request.live_connect_config.response_modalities = (
|
||||
invocation_context.run_config.response_modalities
|
||||
)
|
||||
llm_request.live_connect_config.speech_config = (
|
||||
invocation_context.run_config.speech_config
|
||||
)
|
||||
llm_request.live_connect_config.output_audio_transcription = (
|
||||
invocation_context.run_config.output_audio_transcription
|
||||
)
|
||||
llm_request.live_connect_config.input_audio_transcription = (
|
||||
invocation_context.run_config.input_audio_transcription
|
||||
)
|
||||
llm_request.live_connect_config.realtime_input_config = (
|
||||
invocation_context.run_config.realtime_input_config
|
||||
)
|
||||
llm_request.live_connect_config.enable_affective_dialog = (
|
||||
invocation_context.run_config.enable_affective_dialog
|
||||
)
|
||||
llm_request.live_connect_config.proactivity = (
|
||||
invocation_context.run_config.proactivity
|
||||
)
|
||||
llm_request.live_connect_config.session_resumption = (
|
||||
invocation_context.run_config.session_resumption
|
||||
)
|
||||
llm_request.live_connect_config.context_window_compression = (
|
||||
invocation_context.run_config.context_window_compression
|
||||
)
|
||||
|
||||
|
||||
class _BasicLlmRequestProcessor(BaseLlmRequestProcessor):
|
||||
|
||||
@override
|
||||
async def run_async(
|
||||
self, invocation_context: InvocationContext, llm_request: LlmRequest
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
agent = invocation_context.agent
|
||||
model = agent.canonical_model
|
||||
llm_request.model = model if isinstance(model, str) else model.model
|
||||
llm_request.config = (
|
||||
agent.generate_content_config.model_copy(deep=True)
|
||||
if agent.generate_content_config
|
||||
else types.GenerateContentConfig()
|
||||
)
|
||||
# Only set output_schema if no tools are specified. as of now, model don't
|
||||
# support output_schema and tools together. we have a workaround to support
|
||||
# both output_schema and tools at the same time. see
|
||||
# _output_schema_processor.py for details
|
||||
if agent.output_schema:
|
||||
if not agent.tools or can_use_output_schema_with_tools(model):
|
||||
llm_request.set_output_schema(agent.output_schema)
|
||||
|
||||
llm_request.live_connect_config.response_modalities = (
|
||||
invocation_context.run_config.response_modalities
|
||||
)
|
||||
llm_request.live_connect_config.speech_config = (
|
||||
invocation_context.run_config.speech_config
|
||||
)
|
||||
llm_request.live_connect_config.output_audio_transcription = (
|
||||
invocation_context.run_config.output_audio_transcription
|
||||
)
|
||||
llm_request.live_connect_config.input_audio_transcription = (
|
||||
invocation_context.run_config.input_audio_transcription
|
||||
)
|
||||
llm_request.live_connect_config.realtime_input_config = (
|
||||
invocation_context.run_config.realtime_input_config
|
||||
)
|
||||
llm_request.live_connect_config.enable_affective_dialog = (
|
||||
invocation_context.run_config.enable_affective_dialog
|
||||
)
|
||||
llm_request.live_connect_config.proactivity = (
|
||||
invocation_context.run_config.proactivity
|
||||
)
|
||||
llm_request.live_connect_config.session_resumption = (
|
||||
invocation_context.run_config.session_resumption
|
||||
)
|
||||
llm_request.live_connect_config.context_window_compression = (
|
||||
invocation_context.run_config.context_window_compression
|
||||
)
|
||||
_build_basic_request(invocation_context, llm_request)
|
||||
|
||||
# TODO: handle tool append here, instead of in BaseTool.process_llm_request.
|
||||
|
||||
|
||||
@@ -28,81 +28,103 @@ from ._base_llm_processor import BaseLlmRequestProcessor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...agents.invocation_context import InvocationContext
|
||||
from ...agents.llm_agent import LlmAgent
|
||||
from ...models.llm_request import LlmRequest
|
||||
|
||||
|
||||
class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
|
||||
"""Handles instructions and global instructions for LLM flow."""
|
||||
async def _process_agent_instruction(
|
||||
agent: 'LlmAgent',
|
||||
invocation_context: 'InvocationContext',
|
||||
) -> str:
|
||||
"""Process agent instruction with state injection.
|
||||
|
||||
async def _process_agent_instruction(
|
||||
self, agent, invocation_context: InvocationContext
|
||||
) -> str:
|
||||
"""Process agent instruction with state injection.
|
||||
Resolves the agent's instruction and injects session state variables
|
||||
unless bypass_state_injection is set.
|
||||
|
||||
Args:
|
||||
agent: The agent with instruction to process
|
||||
invocation_context: The invocation context
|
||||
Args:
|
||||
agent: The agent with instruction to process.
|
||||
invocation_context: The invocation context.
|
||||
|
||||
Returns:
|
||||
The processed instruction text
|
||||
"""
|
||||
raw_si, bypass_state_injection = await agent.canonical_instruction(
|
||||
ReadonlyContext(invocation_context)
|
||||
Returns:
|
||||
The processed instruction text with state variables injected.
|
||||
"""
|
||||
raw_si, bypass_state_injection = await agent.canonical_instruction(
|
||||
ReadonlyContext(invocation_context)
|
||||
)
|
||||
si = raw_si
|
||||
if not bypass_state_injection:
|
||||
si = await instructions_utils.inject_session_state(
|
||||
raw_si, ReadonlyContext(invocation_context)
|
||||
)
|
||||
return si
|
||||
|
||||
|
||||
async def _build_instructions(
|
||||
invocation_context: 'InvocationContext',
|
||||
llm_request: 'LlmRequest',
|
||||
) -> None:
|
||||
"""Build and append instructions to the LLM request.
|
||||
|
||||
Handles global instructions (deprecated), static_instruction, and
|
||||
dynamic instruction based on agent configuration.
|
||||
|
||||
Args:
|
||||
invocation_context: The invocation context.
|
||||
llm_request: The LlmRequest to populate with instructions.
|
||||
"""
|
||||
from ...agents.base_agent import BaseAgent
|
||||
from ...agents.llm_agent import LlmAgent
|
||||
|
||||
agent = invocation_context.agent
|
||||
|
||||
root_agent: BaseAgent = agent.root_agent
|
||||
|
||||
# Handle global instructions (DEPRECATED - use GlobalInstructionPlugin instead)
|
||||
# TODO: Remove this code block when global_instruction field is removed
|
||||
if isinstance(root_agent, LlmAgent) and root_agent.global_instruction:
|
||||
raw_si, bypass_state_injection = (
|
||||
await root_agent.canonical_global_instruction(
|
||||
ReadonlyContext(invocation_context)
|
||||
)
|
||||
)
|
||||
si = raw_si
|
||||
if not bypass_state_injection:
|
||||
si = await instructions_utils.inject_session_state(
|
||||
raw_si, ReadonlyContext(invocation_context)
|
||||
)
|
||||
return si
|
||||
llm_request.append_instructions([si])
|
||||
|
||||
# Handle static_instruction - add via append_instructions
|
||||
if agent.static_instruction:
|
||||
from google.genai import _transformers
|
||||
|
||||
# Convert ContentUnion to Content using genai transformer
|
||||
static_content = _transformers.t_content(agent.static_instruction)
|
||||
llm_request.append_instructions(static_content)
|
||||
|
||||
# Handle instruction based on whether static_instruction exists
|
||||
if agent.instruction and not agent.static_instruction:
|
||||
# Only add to system instructions if no static instruction exists
|
||||
si = await _process_agent_instruction(agent, invocation_context)
|
||||
llm_request.append_instructions([si])
|
||||
elif agent.instruction and agent.static_instruction:
|
||||
# Static instruction exists, so add dynamic instruction to content
|
||||
from google.genai import types
|
||||
|
||||
si = await _process_agent_instruction(agent, invocation_context)
|
||||
# Create user content for dynamic instruction
|
||||
dynamic_content = types.Content(role='user', parts=[types.Part(text=si)])
|
||||
llm_request.contents.append(dynamic_content)
|
||||
|
||||
|
||||
class _InstructionsLlmRequestProcessor(BaseLlmRequestProcessor):
|
||||
"""Handles instructions and global instructions for LLM flow."""
|
||||
|
||||
@override
|
||||
async def run_async(
|
||||
self, invocation_context: InvocationContext, llm_request: LlmRequest
|
||||
) -> AsyncGenerator[Event, None]:
|
||||
from ...agents.base_agent import BaseAgent
|
||||
from ...agents.llm_agent import LlmAgent
|
||||
|
||||
agent = invocation_context.agent
|
||||
|
||||
root_agent: BaseAgent = agent.root_agent
|
||||
|
||||
# Handle global instructions (DEPRECATED - use GlobalInstructionPlugin instead)
|
||||
# TODO: Remove this code block when global_instruction field is removed
|
||||
if isinstance(root_agent, LlmAgent) and root_agent.global_instruction:
|
||||
raw_si, bypass_state_injection = (
|
||||
await root_agent.canonical_global_instruction(
|
||||
ReadonlyContext(invocation_context)
|
||||
)
|
||||
)
|
||||
si = raw_si
|
||||
if not bypass_state_injection:
|
||||
si = await instructions_utils.inject_session_state(
|
||||
raw_si, ReadonlyContext(invocation_context)
|
||||
)
|
||||
llm_request.append_instructions([si])
|
||||
|
||||
# Handle static_instruction - add via append_instructions
|
||||
if agent.static_instruction:
|
||||
from google.genai import _transformers
|
||||
|
||||
# Convert ContentUnion to Content using genai transformer
|
||||
static_content = _transformers.t_content(agent.static_instruction)
|
||||
llm_request.append_instructions(static_content)
|
||||
|
||||
# Handle instruction based on whether static_instruction exists
|
||||
if agent.instruction and not agent.static_instruction:
|
||||
# Only add to system instructions if no static instruction exists
|
||||
si = await self._process_agent_instruction(agent, invocation_context)
|
||||
llm_request.append_instructions([si])
|
||||
elif agent.instruction and agent.static_instruction:
|
||||
# Static instruction exists, so add dynamic instruction to content
|
||||
from google.genai import types
|
||||
|
||||
si = await self._process_agent_instruction(agent, invocation_context)
|
||||
# Create user content for dynamic instruction
|
||||
dynamic_content = types.Content(role='user', parts=[types.Part(text=si)])
|
||||
llm_request.contents.append(dynamic_content)
|
||||
await _build_instructions(invocation_context, llm_request)
|
||||
|
||||
# Maintain async generator behavior
|
||||
return
|
||||
|
||||
Reference in New Issue
Block a user