agent_test.go

Overview

The agent_test.go file contains comprehensive unit tests and supporting mock implementations designed to validate the behavior of the Loop Agent implemented in the loopagent package. Its primary purpose is to verify the correct execution semantics of the Loop Agent when orchestrating multiple sub-agents over repeated iterations, including edge cases like infinite loops, early termination via escalation, and usage of function-call tools.

This test file exercises the looping logic by simulating various scenarios with different configurations of sub-agents, iteration limits, and function call behaviors, ensuring that the Loop Agent properly yields expected session events and respects termination signals. It also provides mock implementations of custom agents and LLM agents with function call tools to facilitate controlled testing.

The tests interact with other core components such as the session package (for event management), the runner package (to execute agents within sessions), and the agent interfaces, thus validating integration points within the broader Agent Workflow Management and Loop Agent contexts.

Detailed Explanations

Test Functions

TestNewLoopAgent

• Purpose:
  Tests the creation and runtime behavior of the Loop Agent with various sub-agent configurations and iteration limits.

• Parameters:
  Uses an internal struct args with:

  • maxIterations uint: Number of iterations the Loop Agent should perform (0 means infinite).

  • subAgents []agent.Agent: Slice of sub-agents to be executed within the loop.

• Test Cases:
  Multiple subtests cover scenarios including:

  • Infinite looping with one custom agent.

  • Looping with a single iteration limit.

  • Looping with multiple sub-agents.

  • Looping with sub-agents that escalate via function calls, optionally skipping summarization.

• Execution Flow:
  Within each subtest:

  1. Creates a Loop Agent instance with the specified config.

  2. Creates an in-memory session service to store events.

  3. Uses the runner package to execute the Loop Agent within a session.

  4. Feeds a user input prompt into the agent runner.

  5. Collects yielded events until iteration limits or escalation cause termination.

  6. Compares collected events against expected events, ignoring non-deterministic metadata fields (IDs, timestamps).

• Verification:
  Uses the cmp package with options to ignore volatile fields, comparing the content of each event including Author and LLMResponse.

• Example Usage:

loopAgent, err := loopagent.New(loopagent.Config{
	MaxIterations: 1,
	AgentConfig: agent.Config{
		Name:      "test_agent",
		SubAgents: []agent.Agent{newCustomAgent(t, 0)},
	},
})
// Run the agent and collect events...

Helper Functions and Types

newCustomAgent(t *testing.T, id int) agent.Agent

• Purpose:
  Creates a simple mock custom agent that returns a static greeting message identified by its id.

• Parameters:

  • t *testing.T: Testing context for error reporting.

  • id int: Identifier used for naming and message content.

• Returns:
  An agent.Agent instance with a Run method implemented by customAgent.Run.

• Behavior:
  The agent yields a single session.Event with content "hello <id>" as the model response.

customAgent struct and Run method

• Fields:

  • id int: Identifier for message content.

  • callCounter int: Tracks number of times Run is invoked.

• Method:

  • Run(agent.InvocationContext) iter.Seq2[*session.Event, error]
    Implements an iterator function yielding one event containing a text message ("hello ") from the agent.

• Example:

func (a *customAgent) Run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
	return func(yield func(*session.Event, error) bool) {
		a.callCounter++
		yield(&session.Event{
			LLMResponse: model.LLMResponse{
				Content: genai.NewContentFromText(fmt.Sprintf("hello %v", a.id), genai.RoleModel),
			},
		}, nil)
	}
}

Function Tools for Escalation Testing

• exampleFunctionThatEscalates(ctx tool.Context, EmptyArgs) (map[string]string, error)
  Sets ctx.Actions().Escalate = true and SkipSummarization = false to simulate escalation.

• exampleFunctionThatEscalatesAndSkips(ctx tool.Context, EmptyArgs) (map[string]string, error)
  Sets Escalate = true and SkipSummarization = true.

These are used to create function tools that the LLM agent can call to trigger escalation events during tests.

newLmmAgentWithFunctionCall(t *testing.T, id int, skipSummarization bool) agent.Agent

• Purpose:
  Creates a mock LLM-based agent that uses a function tool to simulate escalation behavior during execution.

• Parameters:

  • t *testing.T: Testing context.

  • id int: Identifier for the agent instance.

  • skipSummarization bool: Whether the function tool skips summarization.

• Implementation Details:

  • Constructs a function tool wrapping one of the escalation functions above.

  • Creates an llmagent.LLMAgent configured with a FakeLLM model that emits a function call or text response depending on input.

  • Returns the configured agent instance.

FakeLLM struct and GenerateContent method

• Purpose:
  A mock implementation of the model.LLM interface to simulate LLM responses for testing.

• Fields:

  • id int: Identifier for message content.

  • callCounter int: Tracks number of calls.

  • skipSummarization bool: Influences response behavior.

• Method:

  • GenerateContent(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error]
    Yields either a function call content or a text content depending on the input request contents length and internal state.

• Usage:
  Simulates an LLM response that triggers function calls causing escalation in the Loop Agent test scenarios.

Important Implementation Details

• The tests use an in-memory session service to capture and analyze events without external dependencies.

• The agent runner executes the Loop Agent, simulating real invocation contexts with user inputs.

• Tests assert that the Loop Agent correctly:

  • Runs all sub-agents sequentially per iteration.

  • Honors maximum iteration limits.

  • Yields events with expected authors and content.

  • Recognizes escalation actions which cause early termination of loops.

• Usage of function tools to escalate allows testing complex interaction flows within the Loop Agent.

• The yield pattern is used extensively for asynchronous event streaming, matching the agent interface contract.

• The cmp package with custom options ensures event comparisons ignore metadata that naturally varies (timestamps, IDs).

Interaction with Other System Components

• Loop Agent (loopagent package):
  The file tests the core Loop Agent implementation which manages repeated execution of sub-agents.

• Sub-agents (agent.Agent):
  Utilizes mock implementations (customAgent and LLM agents) to simulate sub-agent behavior.

• Session Management (session package):
  Events yielded by agents represent session state changes and are captured for verification.

• Agent Execution Runner (runner package):
  Coordinates agent execution within sessions, handling context and event streaming.

• LLM Integration and Tools:
  The mock LLM agents and function tools demonstrate interaction with the LLM framework and function tool system (functiontool).

• Tooling System (tool package):
  Function tools are created to simulate agent function calls affecting control flow (escalation).

Visual Diagram: Structure and Workflow in agent_test.go

flowchart TD
TestNewLoopAgent[TestNewLoopAgent]
newCustomAgent[newCustomAgent]
customAgent_Run[customAgent.Run]
newLmmAgentWithFunctionCall[newLmmAgentWithFunctionCall]
FakeLLM_GenerateContent[FakeLLM.GenerateContent]
exampleFunctionThatEscalates[exampleFunctionThatEscalates]
exampleFunctionThatEscalatesAndSkips[exampleFunctionThatEscalatesAndSkips]
TestNewLoopAgent --> newCustomAgent
newCustomAgent --> customAgent_Run
TestNewLoopAgent --> newLmmAgentWithFunctionCall
newLmmAgentWithFunctionCall --> FakeLLM_GenerateContent
newLmmAgentWithFunctionCall --> exampleFunctionThatEscalates
newLmmAgentWithFunctionCall --> exampleFunctionThatEscalatesAndSkips

• Explanation:
  This flowchart shows key functions and their relationships:

  • TestNewLoopAgent uses newCustomAgent and newLmmAgentWithFunctionCall to create sub-agents.

  • newCustomAgent returns an agent whose Run method is implemented by customAgent.Run.

  • newLmmAgentWithFunctionCall creates an LLM agent which uses FakeLLM.GenerateContent and one of the escalation functions as function tools.

  • The escalation functions influence agent behavior triggering loop termination conditions.

References to Relevant Topics

• The Loop Agent under test is part of the Loop Agent and Agent Workflow Management topics, which define workflow patterns and sub-agent orchestration.

• The use of agent.Agent interfaces, agent.InvocationContext, and streaming event yields aligns with Agent Lifecycle and Callbacks and Agent Execution Runner.

• Mock LLM agents and function tools leverage concepts described in LLM Integration and Agents and Function Tools.

• Event capturing and session state relate to Session Management.

• The runner package integration exemplifies coordination described in Agent Execution Runner.

This documentation provides a detailed technical overview of the testing code validating the looping workflow agent capabilities, the mock agents used for testing, and their interplay with core system components.