LLM Integration and Agents

Overview

The LLM Integration and Agents module provides specialized AI agents designed to interact seamlessly with large language models (LLMs). These agents serve as intelligent intermediaries that encapsulate the logic to communicate with LLM backends, manage invocation contexts, orchestrate tool usage, process instructions, and handle callbacks. This module solves the problem of integrating complex LLM interactions into an extensible agent framework that supports dynamic instruction injection, tool invocation, sub-agent delegation, and rich lifecycle hooks.

By abstracting LLM communication and agent behavior into configurable components, this module enables the creation of AI agents that can perform varied tasks such as answering queries, invoking external tools, maintaining state, and transferring control to sub-agents, all driven by the capabilities of underlying LLMs like Google's Gemini models.

Core Concepts

LLMAgent

The centerpiece of this module is the LLMAgent, a concrete implementation of the generic agent.Agent interface, which integrates an LLM model with a set of tools, instructions, and callback mechanisms.

Instruction Injection

Instructions given to the LLM can be dynamically generated or templated. This module supports:

This flexible instruction handling allows agents to adapt their prompts dynamically based on user context or session state.

Callbacks and Hooks

The agent defines four main callback types to hook into the invocation lifecycle:

These callbacks provide extensibility points for logging, caching, auditing, validation, error handling, and result enhancement.

How It Works

Agent Construction

Agents are instantiated via the constructor llmagent.New(cfg Config). The configuration includes the agent name, description, LLM model, instructions, tools, and callbacks.

During construction:

Invocation Flow

When the agent is invoked via its run method, it performs the following core steps:

  1. Invocation Context Setup: It wraps the invocation context with internal types, attaching session, artifacts, memory, and other context data.

  2. Flow Execution: It creates an llminternal.Flow object that orchestrates the interaction with the LLM and tools, configured with:

    • The LLM model.

    • Request and response processors.

    • The before/after model and tool callbacks.

  3. Running the Flow: It runs the flow as a generator yielding session.Event objects representing LLM responses, tool calls, or other agent events.

  4. Output State Saving: As events are produced, the agent conditionally saves output text to the session state if configured (OutputKey).

Tool Execution

The LLM-generated content may include function calls to tools. When a tool call is detected:

Agent Transfer

Agents can delegate control to sub-agents by issuing special function calls (transfer_to_agent). This enables hierarchical and cooperative agent compositions, where complex tasks are routed dynamically.

Output Persistence

If the agent has an OutputKey configured, it extracts the text content from the final LLM response event and saves it to the session state under that key. This allows subsequent agents or tools to reference prior outputs.

Interaction with Other Modules

Important Design Patterns and Approaches

Key Code References and Illustrations

Agent Construction and Configuration

func New(cfg Config) (agent.Agent, error) {
    // Wrap callbacks into internal types
    // Initialize llmAgent with model, tools, instruction, callbacks
    // Create base agent with agent.New, setting run function
    // Return the constructed LLMAgent
}

The constructor accepts a rich configuration structure (Config) that includes agent naming, description, sub-agents, LLM model, instructions, tools, schemas, and callbacks.

Agent Run Method

func (a *llmAgent) run(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
    // Wrap context with internal invocation context
    // Create llminternal.Flow with model, callbacks, processors
    // Execute flow.Run to generate event stream
    // Save output to state if configured
    // Yield events to caller
}

This method encapsulates the core flow of interacting with the LLM and tools, producing events representing the agent's ongoing actions.

Output State Saving

func (a *llmAgent) maybeSaveOutputToState(event *session.Event) {
    if event.Author != a.Name() { return }
    if a.OutputKey != "" && !event.Partial && event.Content != nil {
        // Concatenate text parts excluding thoughts
        // Save result under OutputKey in event.Actions.StateDelta
    }
}

This function ensures that the agent's final textual output is persisted into session state, enabling stateful agent coordination.

Callback Types

type BeforeModelCallback func(ctx agent.CallbackContext, llmRequest *model.LLMRequest) (*model.LLMResponse, error)
type AfterModelCallback func(ctx agent.CallbackContext, llmResponse *model.LLMResponse, llmResponseError error) (*model.LLMResponse, error)
type BeforeToolCallback func(ctx tool.Context, tool tool.Tool, args map[string]any) (map[string]any, error)
type AfterToolCallback func(ctx tool.Context, tool tool.Tool, args, result map[string]any, err error) (map[string]any, error)

These callback signatures define the interception points for model and tool invocation, allowing early return, error handling, and result modification.

Visualization of Agent Invocation Flow

sequenceDiagram
participant User as User Input
participant Agent as LLMAgent
participant LLM as LLM Model
participant Tools as Tooling System
participant Session as Session Management
User->>Agent: Send user message
Agent->>Session: Append message to session state
Agent->>Agent: Run BeforeAgentCallbacks
Agent->>Agent: Prepare instruction (static or via InstructionProvider)
Agent->>LLM: Send LLMRequest with instruction and user input
LLM-->>Agent: Return LLMResponse (streamed)
alt LLMResponse contains function call
Agent->>Agent: Run BeforeToolCallbacks
Agent->>Tools: Invoke tool with arguments
Tools-->>Agent: Tool result or error
Agent->>Agent: Run AfterToolCallbacks
Agent->>LLM: Continue LLM interaction with tool result
else LLMResponse is final
Agent->>Agent: Run AfterModelCallbacks
Agent->>Agent: Run AfterAgentCallbacks
Agent->>Session: Append LLM response event, save output to state if configured
Agent-->>User: Return response events
end

This sequence diagram illustrates the interaction among the user, agent, LLM, tools, and session during a single invocation cycle.

Subtopics