Callbacks

Callbacks can be registered to execute on specific events during the prompt and response cycle. This allows you to perform actions such as logging, modifying prompts, or triggering additional processes based on the agent's activity.

Action Callbacks

Action callbacks are triggered when an action is invoked within an agent. This allows you to customize the behavior of actions, such as modifying the action's parameters or logging the action's execution. Great for retrieval augmented generation (RAG) workflows.

before_action :set_context

private
def set_context
  # Logic to set the context for the action
  @context_set = true
  prompt_context.instructions = "Context has been set"
end

Generation Callbacks

Generation callbacks are executed during the generation process of an agent. This allows you to modify the prompt, handle responses, or perform additional processing based on the generated content.

after_generation :process_response

private
def process_response
  # Access the generation provider response
  @response_data = generation_provider.response
end

Around Generation Callbacks

Around generation callbacks wrap the entire generation process, allowing you to perform setup and teardown operations. This is useful for timing, caching, logging, or any operation that needs to wrap the generation process.

Basic Around Generation

The around_generation callback wraps the generation with custom logic:

class AgentWithAroundGeneration < ActiveAgent::Base
  layout "agent"

  generate_with :openai, model: "gpt-4o-mini", instructions: "You are a helpful assistant."

  around_generation :track_generation_calls

  def test_action
    prompt_context(message: "Generate a test response")
  end

  def another_action
    prompt_context(message: "Generate another response")
  end

  private

  def track_generation_calls
    # Track how many times generation is called
    self.class.instance_variable_set(:@generation_count, (self.class.instance_variable_get(:@generation_count) || 0) + 1)

    # Store agent instance in context for tracking
    context.agent_instance.instance_variable_set(:@called_at, Time.current)

    # Execute the actual generation
    yield
  end
end

Testing Around Generation

The callback tracks generation calls and timing:

test "around_generation wraps the generation process" do
  VCR.use_cassette("openai_prompt_context_response") do
    # First call - should execute the generation
    generation_result = @agent_class.test_action.generate_now

    assert generation_result.message.content.present?
    assert_equal 1, @agent_class.instance_variable_get(:@generation_count)

    # Check that the agent instance was tracked
    assert generation_result.prompt.agent_instance.instance_variable_get(:@called_at).present?
  end
end

Around Generation with Conditions

You can apply around_generation callbacks conditionally using :only and :except options:

class AgentWithAroundGenerationAndConditions < ActiveAgent::Base
  layout "agent"

  generate_with :openai, model: "gpt-4o-mini", instructions: "You are a helpful assistant."

  around_generation :log_timing, only: :timed_action
  around_generation :cache_llm_request, except: :uncached_action

  def timed_action
    prompt_context(message: "Timed response")
  end

  def cached_action
    prompt_context(message: "Cached response")
  end

  def uncached_action
    prompt_context(message: "Uncached response")
  end

  private

  def log_timing
    start_time = Time.current
    result = yield
    @generation_time = Time.current - start_time
    # Store timing in context for test access through agent_instance
    context.agent_instance.instance_variable_set(:@generation_time, @generation_time)
    result
  end

  def cache_llm_request
    # Simple tracking for the test
    @cached_actions ||= []
    @cached_actions << action_name

    # Store in context for test access
    context.agent_instance.instance_variable_set(:@cached_actions, @cached_actions)

    yield
  end
end

This pattern is useful for:

• Performance monitoring: Track generation times for specific actions
• Caching: Cache LLM responses for expensive operations
• Rate limiting: Implement custom rate limiting logic
• Debugging: Log detailed information about specific generations

On Stream Callbacks

On stream callbacks are triggered during the streaming of responses from an agent. This allows you to handle real-time updates, such as displaying partial responses in a user interface or logging the progress of the response generation.

Streaming Implementation

The streaming agent demonstrates real-time response handling:

class StreamingAgent < ApplicationAgent
  layout "agent"
  generate_with :openai,
    model: "gpt-4.1-nano",
    instructions: "You're a chat agent. Your job is to help users with their questions.",
    stream: true

  on_stream :broadcast_message

  private

  def broadcast_message
    response = generation_provider.response

    # Broadcast the message to the specified channel
    ActionCable.server.broadcast(
      "#{response.message.generation_id}_messages",
      partial: "streaming_agent/message",
      locals: { message: response.message.content, scroll_to: true }
    )
  end
end

Testing Streaming

The streaming functionality broadcasts each chunk as it arrives:

StreamingAgent.with(message: "Stream this message").prompt_context.generate_now

In this test, the agent broadcasts 30 chunks during the streaming response.