agent.go

Overview

The agent.go file defines the core framework for creating, configuring, and running AI agents within the system. It establishes the foundational Agent interface and provides an implementation that supports hierarchical composition of agents, lifecycle callbacks, event streaming, and integration with session state, artifacts, and memory. This file is a central part of the AI Agent Framework, enabling custom agent behaviors with before- and after-execution hooks, delegation to sub-agents, and asynchronous event generation.

Key responsibilities include:

Defining the Agent interface and related types.
Implementing the base agent struct with lifecycle management.
Managing invocation context and callback execution.
Supporting agent composition through sub-agents.
Integrating tightly with sessions, artifacts, and memory services.
Providing utility functions for callback execution and event authoring.

This file underpins higher-level agent types such as LLM agents, workflow agents, and remote agents by establishing a consistent execution model and lifecycle.

Detailed Explanations

Agent Interface

type Agent interface {
    Name() string
    Description() string
    Run(InvocationContext) iter.Seq2[*session.Event, error]
    SubAgents() []Agent
    internal() *agent
}

Purpose: The core contract for all agents, defining identity, behavior, and composition.
Methods:
- Name() string: Returns the unique name of the agent.
- Description() string: Returns a brief capability description used by LLMs or tooling.
- Run(ctx InvocationContext) iter.Seq2[*session.Event, error]: Executes the agent logic asynchronously, yielding a stream of events or errors.
- SubAgents() []Agent: Lists child agents for delegation.
- internal() *agent: Provides internal access to the concrete implementation (used within the framework).

Agent Creation: `New`

func New(cfg Config) (Agent, error)

Purpose: Constructs a new custom agent based on supplied configuration.
Parameters:
- cfg Config: Configuration settings including name, description, sub-agents, callbacks, and run function.
Returns:
- An Agent instance implementing the interface.
- An error if validation fails (e.g., duplicate sub-agent names).
Functionality:
- Validates that sub-agents are unique within the agent tree.
- Initializes the internal state with the type set to TypeCustomAgent.
- Sets the callbacks and the core run function.
Usage Example:

agent, err := agent.New(agent.Config{
    Name:        "example_agent",
    Description: "An example AI agent",
    Run: func(ctx agent.InvocationContext) iter.Seq2[*session.Event, error] {
        return func(yield func(*session.Event, error) bool) {
            yield(&session.Event{
                Author: "example_agent",
                LLMResponse: model.LLMResponse{
                    Content: genai.NewContentFromText("Hello from example agent", genai.RoleModel),
                },
            }, nil)
        }
    },
})

Config Struct

type Config struct {
    Name                 string
    Description          string
    SubAgents            []Agent
    BeforeAgentCallbacks []BeforeAgentCallback
    Run                  func(InvocationContext) iter.Seq2[*session.Event, error]
    AfterAgentCallbacks  []AfterAgentCallback
}

Name: Unique identifier for the agent; must not be empty or "user".
Description: One-line capability description for use by LLMs or tooling.
SubAgents: Child agents for delegation; parent-child relationships are managed automatically.
BeforeAgentCallbacks: Callbacks executed sequentially before the main run; can short-circuit execution.
Run: Core function defining agent behavior, which yields session events.
AfterAgentCallbacks: Callbacks executed after the main run; can append events or modify state.

Lifecycle Callback Types

type BeforeAgentCallback func(CallbackContext) (*genai.Content, error)
type AfterAgentCallback func(CallbackContext) (*genai.Content, error)

BeforeAgentCallback: Invoked before the agent run; may return content or error to skip main execution.
AfterAgentCallback: Invoked after the agent run; may return additional content or error to append events.

Callbacks receive a CallbackContext that provides access to session state, artifacts, invocation metadata, and methods to modify state.

Artifacts and Memory Interfaces

type Artifacts interface {
    Save(ctx context.Context, name string, data *genai.Part) (*artifact.SaveResponse, error)
    List(context.Context) (*artifact.ListResponse, error)
    Load(ctx context.Context, name string) (*artifact.LoadResponse, error)
    LoadVersion(ctx context.Context, name string, version int) (*artifact.LoadResponse, error)
}

type Memory interface {
    AddSession(context.Context, session.Session) error
    Search(ctx context.Context, query string) (*memory.SearchResponse, error)
}

Artifacts: Provides artifact storage and retrieval capabilities during an agent invocation.
Memory: Provides access to user-scoped memory across sessions.

These interfaces allow agents to persist and retrieve data relevant to the session or user.

Internal `agent` Struct and Methods

type agent struct {
    agentinternal.State

    name, description string
    subAgents         []Agent

    beforeAgentCallbacks []BeforeAgentCallback
    run                  func(InvocationContext) iter.Seq2[*session.Event, error]
    afterAgentCallbacks  []AfterAgentCallback
}

Implements the Agent interface.
Holds configuration, lifecycle callbacks, and state.
Methods:
- Name(), Description(), SubAgents() simply return configured values.
- Run(ctx InvocationContext) iter.Seq2[*session.Event, error] manages the full lifecycle:
  - Runs before callbacks and may short-circuit if they produce content or error.
  - Executes the main run function, streaming events.
  - Runs after callbacks to append events/state.
  - Sets event authors if missing.
- internal() returns itself for internal use.

Agent Run Method Workflow

The Run method orchestrates the invocation lifecycle as follows:

Creates a wrapped InvocationContext (invocationContext) that embeds the original context plus accessors for artifacts, memory, and session.
Executes runBeforeAgentCallbacks:
- Iterates over beforeAgentCallbacks.
- If any callback returns non-nil content or error, yields an event constructed from that content/error and ends invocation early.
If invocation is not ended, runs the core run function, yielding events.
Automatically assigns the event author to the agent's name if not set.
If invocation is still active, executes runAfterAgentCallbacks:
- Runs after callbacks sequentially.
- If any returns content or error, yields an event with that content/error.
Yields events via the returned iterator until completion or early termination.

Callback Execution Helpers

runBeforeAgentCallbacks(ctx InvocationContext) (*session.Event, error)
- Runs all before callbacks.
- If a callback returns content or error, constructs a new event with that content, marks invocation ended, and returns it.
- If state delta modifications exist without content, creates an event with those changes.
- Otherwise returns nil.
runAfterAgentCallbacks(ctx InvocationContext) (*session.Event, error)
- Runs all after callbacks.
- If a callback returns content or error, creates a corresponding event.
- If state delta exists, creates an event with state changes.
- Returns nil if no content or state was produced.

Callback Context

type callbackContext struct {
    context.Context
    invocationContext InvocationContext
    actions           *session.EventActions
}

Implements the CallbackContext interface.
Provides access to:
- Agent name.
- Read-only and mutable session state.
- Artifacts.
- Invocation metadata (InvocationID, UserContent, AppName, Branch, SessionID, UserID).
Allows callbacks to modify state via actions.StateDelta, which is applied to the session at event creation.

Invocation Context Internal Wrapper

type invocationContext struct {
    context.Context

    agent     Agent
    artifacts Artifacts
    memory    Memory
    session   session.Session

    invocationID  string
    branch        string
    userContent   *genai.Content
    runConfig     *RunConfig
    endInvocation bool
}

Wraps the original invocation context, exposing all necessary runtime data and services.
Provides methods for accessing:
- Agent, Artifacts, Memory, Session.
- Invocation metadata like ID, branch, user content, run configuration.
Supports EndInvocation() to mark early termination.
Supports Ended() to query termination status.

Event Authoring Logic

func getAuthorForEvent(ctx InvocationContext, event *session.Event) string {
    if event.LLMResponse.Content != nil && event.LLMResponse.Content.Role == genai.RoleUser {
        return genai.RoleUser
    }
    return ctx.Agent().Name()
}

Assigns the author of an event based on the content role.
If the content role is user, the author is set to "user".
Otherwise, author is set to the invoking agent's name.
Ensures clear attribution of events in multi-agent sessions.

Important Implementation Details and Algorithms

Uniqueness Enforcement: The New constructor enforces uniqueness of sub-agents by name, preventing ambiguous delegation trees.
Event Streaming Model: The agent's Run method returns an iterator that yields events asynchronously, enabling multi-step conversations, tool invocations, or partial results.
Callback Short-Circuiting: Before-agent callbacks can prevent the main run from executing by returning content or errors, allowing preemptive responses or validations.
State Delta Aggregation: Callbacks accumulate state changes in a delta map, which is included in generated events to update session state atomically.
Invocation Branching: Each invocation carries a branch string identifying the agent call path; this supports nested and parallel agent invocations with isolated contexts.
Integration with Session & Artifact Services: The invocation context exposes services for managing session state, artifacts, and memory, enabling agents to persist and retrieve data transparently.

Interactions with Other System Components

Session Management: Uses session.Session to read and modify session state and to create events representing agent outputs or state changes.
Artifact Management: Uses the Artifacts interface to save, list, and load artifacts related to the current session.
Memory Service: Accesses user-scoped memory via the Memory interface for knowledge retrieval or persistence across sessions.
Agent Invocation Context: Implements and wraps the InvocationContext interface defined in Agent Invocation Context, providing uniform access to invocation metadata and services.
Lifecycle Callbacks: Implements lifecycle hooks as described in Agent Lifecycle and Callbacks, enabling extensibility of agent behavior.
LLM Integration: Agents created here serve as base for specialized LLM agents (covered in LLM Integration and Agents).
Agent Composition: Supports hierarchical agent trees and delegation, which is leveraged by Agent Workflow Management and other orchestrators.
Agent Execution Runner: Works with the execution runner (Agent Execution Runner) which drives agent invocation cycles within sessions.

Visual Diagram: Agent Structure and Lifecycle Flow

classDiagram
class Agent {
+Name()
+Description()
+Run()
+SubAgents()
+internal()
}
class agent {
-name: string
-description: string
-subAgents: []Agent
-beforeAgentCallbacks: []BeforeAgentCallback
-run: func(InvocationContext) iter.Seq2[*session.Event, error]
-afterAgentCallbacks: []AfterAgentCallback
+Name()
+Description()
+Run()
+SubAgents()
+internal()
}
Agent <|.. agent
agent o-- "0..*" Agent : subAgents
agent : +runBeforeAgentCallbacks()
agent : +runAfterAgentCallbacks()

flowchart TD
Start[Start Run Invocation]
BeforeCB[Run BeforeAgentCallbacks]
CheckBefore{Before callback returns content/error?}
SkipRun[Skip main agent Run]
RunMain[Run agent's Run function]
AfterCB[Run AfterAgentCallbacks]
EmitEvent[Yield Event]
End[End Invocation]
Start --> BeforeCB --> CheckBefore
CheckBefore -- Yes --> EmitEvent --> SkipRun --> End
CheckBefore -- No --> RunMain --> AfterCB --> EmitEvent --> End

The class diagram shows the Agent interface and its concrete implementation agent with key properties and methods.
The flowchart illustrates the lifecycle flow inside the Run method, highlighting before callbacks that can short-circuit, the main run execution, after callbacks, and event emission.

Summary of Key Types and Interfaces

Type	Description
`Agent`	Interface for AI agents defining identity, behavior, and sub-agents.
`agent`	Concrete implementation of `Agent` supporting lifecycle and callbacks.
`Config`	Configuration struct for creating custom agents.
`InvocationContext`	Context interface providing runtime data/services during invocation.
`CallbackContext`	Interface passed to lifecycle callbacks for state and metadata access.
`BeforeAgentCallback`	Function type for pre-run lifecycle callbacks.
`AfterAgentCallback`	Function type for post-run lifecycle callbacks.
`Artifacts`	Interface for managing session artifacts.
`Memory`	Interface for accessing user-scoped memory.

For further details on invocation context and lifecycle callbacks, see the subtopics Agent Invocation Context and Agent Lifecycle and Callbacks. For usage patterns involving session management and artifact integration, refer to Session Management and Artifact Management.