state_agent_test.go

Overview

This file contains comprehensive unit and integration tests for the LLM agent implementation, focusing on session lifecycle management, agent callbacks, and tool invocation with callback chains. It provides mock implementations, test helpers, and multiple test cases that verify the correctness of agent behavior within a session-driven environment.

The tests validate how the LLM agent:

The file thus serves as a critical verification suite ensuring robust integration of LLM agents with session management, callback processing, and tool execution.

Key Components

FakeLLM

FakeLLM is a mock implementation of the model.LLM interface used to simulate LLM responses in tests.

Usage:
Injected as the LLM model in agent tests to control and predict LLM output during agent invocation.

assertSessionValues

A helper function to verify the presence or absence of keys in both the callback context session state and the underlying session service state.

Callbacks for Agent Lifecycle

These are implementations of agent lifecycle callbacks wired into the tests to simulate state mutation and validation at different phases:

Usage:
These callbacks are registered in the agent configuration during tests to validate correct state propagation and callback chaining (Agent Lifecycle and Callbacks).

TestAgentSessionLifecycle

This test function verifies the full lifecycle of an agent session:

It validates that session state propagates correctly through all lifecycle callbacks and that the agent integrates properly with the session service (Session Management, Agent Execution Runner).

Tool Implementations

Three example tools simulate external functionalities callable by the agent:

Each tool receives a typed argument struct and returns a typed result struct, facilitating JSON serialization and integration with function tools (Function Tools).

Before Tool Callbacks

Callbacks run before tool execution to enable auditing, security, and validation:

These callbacks demonstrate how to enforce policies and observe tooling activity before execution (Agent Lifecycle and Callbacks).

After Tool Callbacks

Callbacks run after tool execution to enhance or process results asynchronously:

These callbacks modify the results returned from tools, enabling richer responses or additional processing steps (Agent Lifecycle and Callbacks).

Helper Function: collectToolResults

Consumes the event stream from a runner execution and collects the tool function call responses into a slice of generic maps for inspection.

Used in the TestToolCallbacksAgent to gather and verify the output of tool calls.

TestToolCallbacksAgent

This complex test exercises the interaction of the agent with tools and their associated callbacks:

This test validates the full pipeline of tool invocation with layered callbacks and state management (Tooling System, LLM Integration and Agents).

Interaction with Other Parts of the System

Important Implementation Details and Algorithms

Usage Examples

Creating and Running a Simple Agent Session (From TestAgentSessionLifecycle)

ctx := context.Background()
testSessionService = session.InMemoryService()

fakeLLM := &FakeLLM{
    GenerateContentFunc: func(ctx context.Context, req *model.LLMRequest, stream bool) (model.LLMResponse, error) {
        return model.LLMResponse{
            Content: genai.NewContentFromText("test model response", genai.RoleModel),
        }, nil
    },
}

rootAgent, err := llmagent.New(llmagent.Config{
    Name: "root_agent",
    Instruction: "Test instruction",
    Model: fakeLLM,
    BeforeAgentCallbacks: []agent.BeforeAgentCallback{beforeAgentCallback(t)},
    BeforeModelCallbacks: []llmagent.BeforeModelCallback{beforeModelCallback(t)},
    AfterModelCallbacks: []llmagent.AfterModelCallback{afterModelCallback(t)},
    AfterAgentCallbacks: []agent.AfterAgentCallback{afterAgentCallback(t)},
})
if err != nil {
    t.Fatal(err)
}

r, err := runner.New(runner.Config{
    AppName: "test_app",
    Agent: rootAgent,
    SessionService: testSessionService,
})
if err != nil {
    t.Fatal(err)
}

createResp, err := testSessionService.Create(ctx, &session.CreateRequest{AppName: "test_app", UserID: "test_user"})
if err != nil {
    t.Fatal(err)
}
sessionID := createResp.Session.ID()

userContent := genai.NewContentFromText("Hello agent", genai.RoleUser)
eventStream := r.Run(ctx, "test_user", sessionID, userContent, agent.RunConfig{})

for _, err := range eventStream {
    if err != nil {
        t.Fatal(err)
    }
}

// Validate session state keys set by callbacks...

Testing Tool Callbacks (From TestToolCallbacksAgent)

fakeLLM := &FakeLLM{
    GenerateContentFunc: func(ctx context.Context, req *model.LLMRequest, stream bool) (model.LLMResponse, error) {
        // Map user queries to tool function calls
        switch userText {
        case "weather in London":
            return model.LLMResponse{Content: genai.NewContentFromFunctionCall("get_weather", map[string]any{"location": "London"}, genai.RoleModel)}, nil
        // other cases...
        }
    },
}

getWeatherTool, _ := functiontool.New(functiontool.Config{Name: "get_weather"}, GetWeather)
calculateTool, _ := functiontool.New(functiontool.Config{Name: "calculate"}, Calculate)
logActivityTool, _ := functiontool.New(functiontool.Config{Name: "log_activity"}, LogActivity)

agentConfig := llmagent.Config{
    Name: "tool_callbacks_agent",
    Model: fakeLLM,
    Tools: []tool.Tool{getWeatherTool, calculateTool, logActivityTool},
    BeforeToolCallbacks: []llmagent.BeforeToolCallback{
        beforeToolAuditCallback,
        beforeToolSecurityCallback,
        beforeToolValidationCallback,
    },
    AfterToolCallbacks: []llmagent.AfterToolCallback{
        afterToolEnhancementCallback,
        afterToolAsyncCallback,
    },
}

rootAgent, err := llmagent.New(agentConfig)
// Setup runner, create session, run agent, collect and verify results...

Mermaid Diagram: Structure and Workflow of state_agent_test.go

flowchart TD
A[TestAgentSessionLifecycle] --> B[FakeLLM]
A --> C[Callbacks: beforeAgent, beforeModel, afterModel, afterAgent]
A --> D["Session Service (InMemory)"]
A --> E[Runner]
F[TestToolCallbacksAgent] --> G[FakeLLM with Query-to-Tool Mapping]
F --> H[Tools: GetWeather, Calculate, LogActivity]
F --> I[BeforeToolCallbacks: Audit, Security, Validation]
F --> J[AfterToolCallbacks: Enhancement, Async]
F --> D
B -->|Simulates| LLM
H -->|Simulates| External Tools
I -->|Pre-process| H
J -->|Post-process| H
E -->|Runs| A
E -->|Runs| F
D -->|Persists State| A
D -->|Persists State| F

Detailed Descriptions of Functions and Types

FakeLLM

assertSessionValues

beforeAgentCallback

beforeModelCallback

afterModelCallback

afterAgentCallback

TestAgentSessionLifecycle

Tool Implementations

Before Tool Callbacks

After Tool Callbacks

collectToolResults

TestToolCallbacksAgent

References to Related Topics

This documentation provides a detailed understanding of the test suite validating agent session lifecycle and tool callback integrations in the LLM agent framework.