agent_tool_test.go

Overview

This file contains unit tests validating the functionality of the agenttool package, which implements tools wrapping sub-agents to be used as callable tools within AI workflows. The tests verify the correctness of agent tool creation, declaration generation, input/output schema validations, interaction with mock language models, and configuration flags such as skipping summarization.

The primary focus is on ensuring that the agenttool correctly:

• Generates function declarations for LLM integration.

• Validates input arguments against defined JSON schemas.

• Runs sub-agents in isolated sessions and collects their outputs.

• Validates the output from sub-agent executions.

• Handles tool runs without explicit schemas.

• Honors configuration flags for summarization behavior.

This test file plays a crucial role in guaranteeing the robustness and correctness of the Agent Tools component of the Tooling System and its interaction with the Agent Framework.

Detailed Breakdown of Tests

Test Functions

1. TestAgentTool_Declaration

• Purpose: Tests that the agenttool correctly generates a function declaration reflecting the agent's input schema.

• Key Points:

  • Creates an agent with a defined input JSON schema (request as a string).

  • Instantiates the agent tool.

  • Asserts that the returned function declaration matches the expected name ("math_agent"), description, and parameter schema.

• Example Usage:

// Create an agent with input schema and get the function declaration.
agent := createAgent(t, inputSchema, nil)
agentTool := agenttool.New(agent, nil)
decl := agentTool.(toolinternal.FunctionTool).Declaration()

• Return: None (test fails on mismatch).

2. TestAgentTool_DeclarationWithoutSchema

• Purpose: Verifies declaration generation when the agent has no input schema.

• Key Points:
  Expected to fallback to a default input schema with a required "request": string property.

3. TestAgentTool_Run_InputValidation

• Purpose: Ensures that input validation for the tool's Run method rejects invalid or incomplete arguments.

• Test Cases:

  • Extra unknown fields.

  • Invalid types for properties.

  • Missing required fields.

• Behavior: Each invalid input case must result in an error from Run.

4. TestAgentTool_Run_OutputValidation

• Purpose: Validates that the tool detects invalid output produced by the sub-agent.

• Setup:
  Uses a mock LLM returning output violating the output schema (e.g., "is_valid" field has invalid type).

• Expectation:
  Run should return an error indicating output validation failure.

5. TestAgentTool_Run_Successful

• Purpose: Tests a successful run of the agent tool with valid input and output schemas.

• Setup:
  Mock LLM returns valid JSON output matching the output schema.

• Assertions:

  • No error in Run.

  • Returned map matches expected output fields and values.

6. TestAgentTool_Run_WithoutSchema

• Purpose: Verifies run behavior when the agent has no input or output schema.

• Behavior:
  The tool returns a default map with result extracted from the first text part of the LLM response.

• Checks:
  Correct parsing of multi-part LLM responses, ignoring subsequent parts.

7. TestAgentTool_Run_EmptyModelResponse

• Purpose: Handles the scenario where the mock LLM returns empty content.

• Expectation:
  The tool returns an empty map without error.

8. TestAgentTool_Run_SkipSummarization

• Purpose: Tests the SkipSummarization configuration flag controlling summarization behavior after tool execution.

• Scenario:

  • Creates two agent tools: one with SkipSummarization=true and one with false.

  • Runs each tool and asserts the SkipSummarization flag in the session actions matches the config.

Helper Functions

createAgent

• Purpose: Creates a new LLM-based agent with optional input and output JSON schemas.

• Parameters:

  • t *testing.T: Testing instance.

  • inputSchema *genai.Schema: Optional input JSON schema.

  • outputSchema *genai.Schema: Optional output JSON schema.

• Returns: An agent.Agent instance.

• Behavior:
  Uses the Gemini model (fake API key) and sets common agent configuration such as name, description, and instruction.

createAgentWithModel

• Purpose: Similar to createAgent but accepts a mock or custom LLM model instead of creating a Gemini model.

• Allows: Injecting mocked LLM behavior for testing.

createToolContext

• Purpose: Prepares a tool.Context for running the agent tool.

• Parameters:

  • t *testing.T: Testing instance.

  • testAgent agent.Agent: The agent to create context for.

• Returns: A tool.Context instance with an in-memory session and invocation context.

• Implementation Details:

  • Creates an in-memory session service.

  • Creates a session with fixed app/user/session IDs.

  • Wraps session in mutable session internal type.

  • Creates an invocation context incorporating the session.

  • Returns the internal tool context for use in tool invocation.

Important Implementation Details and Algorithms

• Input/Output Validation:
  The agent tool relies on JSON schema definitions for input/output validation. Inputs passed to Run are validated against the input schema before execution, and outputs from the sub-agent are validated similarly against the output schema.

• Mock Model Usage:
  Tests employ testutil.MockModel to simulate LLM responses, allowing control over the model output and testing error conditions like invalid outputs or empty responses.

• Session Isolation:
  Each sub-agent invocation runs within a new session context created in createToolContext(), simulating realistic agent invocation lifecycles.

• Skip Summarization Flag:
  The SkipSummarization flag in the agent tool config controls whether automatic summarization of the session events should be skipped after the sub-agent tool execution, affecting downstream event processing.

• Use of go-cmp for Deep Comparison:
  The tests use cmp.Diff from github.com/google/go-cmp/cmp to compare expected and actual outputs or declarations, providing detailed diff output in case of mismatches.

• Type Assertions:
  The tests assert that the created agent tool implements the toolinternal.FunctionTool interface, which exposes the Declaration() and Run() methods critical for function tool integration.

Interaction with Other System Components

• Agent Package:
  The test creates agents using the llmagent package (an LLM-based agent), configuring them with schemas and instructions.

• Session Management:
  Uses in-memory session service from the session package to create and manage sessions for agent runs.

• Tooling System:
  The tests exercise the agenttool package, part of the broader Tooling System, wrapping agents as tools.

• Internal Contexts:
  Uses internal invocation contexts (icontext) to simulate real invocation environments for the tools.

• Mock Model Utility:
  The testutil.MockModel provides controlled LLM responses to simulate different output scenarios.

Visual Diagram: Structure and Workflow of agent_tool_test.go

flowchart TD
A[Test Functions] -->|createAgent| B[Agent Creation Helpers]
A -->|createAgentWithModel| B
A -->|createToolContext| C[Tool Context Creation]
A -->|agenttool.New| D[Agent Tool Instance]
D --> E[FunctionTool Interface]
E --> F["Declaration()"]
E --> G["Run()"]
F --> H[Validate Input Schema]
G --> I[Execute Sub-Agent]
I --> J[Use MockModel for LLM Response]
J --> K[Validate Output Schema]
G --> L[Skip Summarization Flag]
C --> M[In-Memory Session Service]
M --> N[Session Creation]
N --> O[Invocation Context Creation]

Usage Examples from Tests

// Example: Validate agent tool declaration with input schema
func TestAgentTool_Declaration(t *testing.T) {
    inputSchema := &genai.Schema{...}
    agent := createAgent(t, inputSchema, nil)
    agentTool := agenttool.New(agent, nil)
    decl := agentTool.(toolinternal.FunctionTool).Declaration()
    // Assert decl matches expected schema
}

// Example: Run agent tool with input validation error expected
func TestAgentTool_Run_InputValidation(t *testing.T) {
    ...
    _, err := toolImpl.Run(toolCtx, invalidArgs)
    if err == nil {
        t.Fatal("expected error due to invalid input")
    }
}

// Example: Run agent tool successfully with mock model
func TestAgentTool_Run_Successful(t *testing.T) {
    testLLM := &testutil.MockModel{...}
    agent := createAgentWithModel(t, inputSchema, outputSchema, testLLM)
    agentTool := agenttool.New(agent, nil)
    result, err := agentTool.(toolinternal.FunctionTool).Run(toolCtx, map[string]any{"is_magic": true})
    // Assert no error and result matches expected output
}

References to Related Topics

• The tests validate the agenttool package, which is fully described in the Agent Tools subtopic of the Tooling System.

• The agenttool wraps agents created via the LLM Integration and Agents framework, specifically those implemented by the LLM Agent Configuration and Agent Lifecycle and Callbacks.

• Session and context management used in tests relate to the Session Management and Agent Invocation Context topics.

• Input/output validation leverages JSON schema concepts from the LLM tooling system and aligns with Function Tools.

• The mock LLM models are part of test utilities that simulate responses for agents in the context of LLM Integration and Agents.

This documentation provides a comprehensive understanding of the agent_tool_test.go file, covering its purpose, test cases, helper functions, and integration points within the system.