AI Agent Framework

The AI Agent Framework forms the foundational layer for creating, managing, and executing AI agents within the system. It provides a standardized interface and core implementations that enable agents to process inputs, invoke sub-agents, interact with session state, manage artifacts, and respond asynchronously through event streams. This framework abstracts the complexities involved in agent lifecycle management, invocation context propagation, callback handling, and multi-agent compositions.

Core Concepts and Purpose

At its core, the AI Agent Framework defines the Agent interface, which is the primary contract for all AI agents in the system. This interface and its supporting types enable developers to:

The framework solves the problem of standardizing how AI agents are constructed and invoked, ensuring consistent integration with sessions, artifacts, memory, and tools. It also promotes composability by allowing agents to delegate tasks to sub-agents with isolated contexts.

Agent Interface and Implementation

The Agent Interface

The central interface is Agent (defined in agent.go), which requires the following:

Agents are created via the New constructor, which accepts a Config struct specifying their name, description, sub-agents, callbacks, and the core run function defining the agent's behavior.

Agent Configuration and Lifecycle Callbacks

The Config struct includes:

This callback mechanism allows injection of custom logic around agent execution, such as pre-processing user input, validation, or post-processing output.

Event Streaming Model

The Run method returns an iterator (iter.Seq2[*session.Event, error]) that yields multiple session events asynchronously. This design supports multi-step invocations where an agent can emit partial results, invoke tools or LLMs multiple times, and transfer control gracefully.

Example of a simplified run sequence:

for event, err := range agent.Run(ctx) {
    if err != nil {
        // handle error
    }
    // process event (e.g., send to client, update session)
}

The framework automatically sets the event author to the agent's name if not specified.

Invocation Context

Agents receive an InvocationContext (detailed below) encapsulating all runtime data and services needed during a single invocation call. This context includes session state, artifacts, memory, user content, and flags for controlling invocation flow.

Invocation Context

The InvocationContext interface (defined in context.go) represents all contextual information for an agent invocation. It abstracts:

Agents call EndInvocation() on the context to signal that no further steps should be executed.

Agent Lifecycle and Callbacks

The framework manages the agent lifecycle in three main stages:

  1. BeforeAgentCallbacks: Executed sequentially before the agent's main Run logic. If any callback returns content or an error, the agent run is skipped, and the callback output is emitted as an event.

  2. Run: The core agent logic runs, yielding events as it processes the invocation.

  3. AfterAgentCallbacks: Executed sequentially after the main run completes. If any callback returns new content or error, it generates and yields an additional event.

Callbacks receive a CallbackContext which allows them to access the session state, artifacts, and invocation metadata. They can also modify the session state by returning state deltas that are applied as part of the event actions.

This lifecycle control enables flexible insertion of custom logic before and after agent execution, supporting use cases like validation, logging, or augmenting responses.

Sub-Agent Composition and Invocation Branching

Agents can be composed into trees by specifying sub-agents in their configuration. The framework automatically manages parent-child relationships and allows agents to delegate parts of their tasks to sub-agents.

Each agent invocation carries a branch string identifying its position in the agent hierarchy. This branching mechanism is especially important for agents that run sub-agents in parallel, ensuring isolated conversation histories per sub-agent.

Branches are represented as dot-separated strings (e.g., agent1.agent2.agent3), supporting nested sub-agent calls and context isolation.

Interaction with Sessions, Artifacts, and Memory

The framework tightly integrates with the session and artifact management systems:

This integration allows agents to maintain continuity, enrich responses with stored data, and share knowledge across multiple invocations.

Agent Loading and Management

The framework provides a Loader interface (in loader.go) to manage collections of agents. A loader can:

Two main implementations are provided:

This abstraction supports dynamic agent discovery and runtime flexibility.

Example: Agent Run Flow

The following flowchart illustrates the step-by-step execution of an agent invocation within the framework, highlighting callback execution and event yielding:

flowchart TD
Start[Start Invocation]
BeforeCB[Run BeforeAgentCallbacks]
CheckBefore{Callback returned content or error?}
RunAgent[Run Agent Run Function]
AfterCB[Run AfterAgentCallbacks]
EmitEvent[Yield Event]
End[End Invocation]
Start --> BeforeCB --> CheckBefore
CheckBefore -- Yes --> EmitEvent --> End
CheckBefore -- No --> RunAgent --> AfterCB --> EmitEvent --> End

Code Snippet: Agent Creation and Run Invocation

// Create a new custom agent with before and after callbacks.
agent, err := agent.New(agent.Config{
    Name:        "example_agent",
    Description: "An example AI agent",
    Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
        // Agent logic here...
        return func(yield func(*session.Event, error) bool) {
            // Yield an initial event.
            yield(&session.Event{
                Author: "example_agent",
                LLMResponse: model.LLMResponse{
                    Content: genai.NewContentFromText("Hello from example agent", genai.RoleModel),
                },
            }, nil)
        }
    },
    BeforeAgentCallbacks: []agent.BeforeAgentCallback{
        func(ctx agent.CallbackContext) (*genai.Content, error) {
            // Optional pre-run logic
            return nil, nil
        },
    },
    AfterAgentCallbacks: []agent.AfterAgentCallback{
        func(ctx agent.CallbackContext) (*genai.Content, error) {
            // Optional post-run logic
            return nil, nil
        },
    },
})

// Running the agent with an invocation context.
for event, err := range agent.Run(invocationCtx) {
    if err != nil {
        // Handle error
    }
    // Process event (e.g., send to client)
}

Interactions with Other Modules

The AI Agent Framework interacts closely with several other parts of the system:

Summary of Key Types

Type

Description

Agent

Interface defining agent behavior and composition.

InvocationContext

Provides context and services for a single agent call.

BeforeAgentCallback

Callback invoked before the main agent run.

AfterAgentCallback

Callback invoked after the main agent run.

Loader

Interface for loading and managing agents.

Artifacts

Interface for artifact operations within a session.

Memory

Interface for user-scoped memory access.

This framework enables modular, composable AI agents with rich lifecycle management and contextual awareness, serving as the backbone for all subsequent specialized agent implementations within the system such as LLM agents, workflow agents, and remote agents.

For detailed usage of invocation context and lifecycle callbacks, see the subtopics: