LemonAgent.Agent (lemon_agent v0.1.0)

View Source

GenServer implementation of an AI Agent that manages conversation state and streams.

This module ports the TypeScript Agent class to Elixir, providing a stateful process that handles prompts, tool execution, and event broadcasting to subscribers.

Features

  • State Management: Maintains conversation messages, model configuration, and tools
  • Streaming: Consumes LemonAgent.Loop streams and broadcasts events to subscribers
  • Queue System: Supports steering (mid-run interrupts) and follow-up message queues
  • Abort Support: Allows cancellation of running streams
  • Subscriber Management: Auto-cleanup of dead subscribers via process monitoring

Usage

{:ok, agent} = LemonAgent.Agent.start_link(
  initial_state: %{system_prompt: "You are a helpful assistant"},
  convert_to_llm: &my_converter/1
)

# Subscribe to events
unsubscribe = LemonAgent.Agent.subscribe(agent, self())

# Send a prompt
:ok = LemonAgent.Agent.prompt(agent, "Hello!")

# Wait for completion
:ok = LemonAgent.Agent.wait_for_idle(agent)

State

The agent maintains the following state:

  • agent_state - The LemonAgent.Types.AgentState containing messages, model, tools, etc.
  • listeners - List of subscriber PIDs that receive events
  • abort_ref - Reference for abort signaling
  • running_task - The currently running Task or nil
  • steering_queue / follow_up_queue - Message queues for mid-run and post-run messages
  • waiters - Processes waiting for the agent to become idle

Summary

Types

convert_to_llm_fn()
get_api_key_fn()
opts()
queue_mode()
state()
stream_fn()
transform_context_fn()
waiter()

Functions

abort(agent)

Aborts the currently running stream.

append_message(agent, message)

Appends a message to the conversation.

child_spec(init_arg)

Returns a specification to start this module under a supervisor.

clear_all_queues(agent)

Clears all pending steering and follow-up messages.

clear_follow_up_queue(agent)

Clears all pending follow-up messages.

clear_steering_queue(agent)

Clears all pending steering messages.

continue(agent)

Continues from the current context.

follow_up(agent, message, opts \\ [])

Queues a follow-up message to be processed after the agent finishes.

get_follow_up_mode(agent)

Gets the follow-up queue mode.

get_session_id(agent)

Gets the current session ID.

get_state(agent)

Gets the current agent state.

get_steering_mode(agent)

Gets the steering queue mode.

prompt(agent, message)

Sends a prompt to the agent.

replace_messages(agent, messages)

Replaces all messages in the conversation.

reset(agent)

Resets the agent to initial state.

set_follow_up_mode(agent, mode)

Sets the follow-up queue consumption mode.

set_model(agent, model)

Sets the model to use for LLM calls.

set_session_id(agent, session_id)

Sets the session ID for provider caching.

set_steering_mode(agent, mode)

Sets the steering queue consumption mode.

set_system_prompt(agent, prompt)

Sets the system prompt.

set_thinking_level(agent, level)

Sets the thinking/reasoning level.

set_tools(agent, tools)

Sets the available tools.

start_link(opts \\ [])

Starts an Agent GenServer.

steer(agent, message, opts \\ [])

Queues a steering message to interrupt the agent mid-run.

subscribe(agent, subscriber_pid)

Subscribes a process to agent events.

wait_for_idle(agent, opts \\ [])

Waits for the agent to become idle (no running task).

Types

convert_to_llm_fn()

@type convert_to_llm_fn() :: ([LemonAgent.Types.agent_message()] ->
                          [LemonAi.Types.message()]
                          | {:ok, [LemonAi.Types.message()]})

get_api_key_fn()

@type get_api_key_fn() :: (String.t() -> String.t() | nil)

opts()

@type opts() :: [
  initial_state: map(),
  convert_to_llm: convert_to_llm_fn(),
  transform_context: transform_context_fn(),
  stream_fn: stream_fn() | nil,
  steering_mode: queue_mode(),
  follow_up_mode: queue_mode(),
  session_id: String.t(),
  get_api_key: get_api_key_fn(),
  thinking_budgets: map(),
  stream_options: LemonAi.Types.StreamOptions.t(),
  max_tool_turns: pos_integer() | :infinity | nil,
  tool_timeout_ms: pos_integer() | :infinity | nil,
  queue_call_timeout: timeout(),
  name: GenServer.name()
]

queue_mode()

@type queue_mode() :: :all | :one_at_a_time

state()

@type state() :: %{
  agent_state: LemonAgent.Types.AgentState.t(),
  listeners: [{pid(), reference()}],
  abort_ref: reference() | nil,
  running_task: Task.t() | nil,
  steering_queue: [LemonAgent.Types.agent_message()],
  follow_up_queue: [LemonAgent.Types.agent_message()],
  steering_mode: queue_mode(),
  follow_up_mode: queue_mode(),
  convert_to_llm: convert_to_llm_fn(),
  transform_context: transform_context_fn() | nil,
  stream_fn: stream_fn() | nil,
  session_id: String.t() | nil,
  get_api_key: get_api_key_fn() | nil,
  thinking_budgets: map(),
  stream_options: LemonAi.Types.StreamOptions.t(),
  max_tool_turns: pos_integer() | :infinity | nil,
  tool_timeout_ms: pos_integer() | :infinity | nil,
  queue_call_timeout: timeout(),
  waiters: [waiter()]
}

stream_fn()

@type stream_fn() :: (LemonAi.Types.Model.t(),
                LemonAi.Types.Context.t(),
                LemonAi.Types.StreamOptions.t() ->
                  {:ok, LemonAi.EventStream.t()}
                  | LemonAi.EventStream.t()
                  | {:error, term()})

transform_context_fn()

@type transform_context_fn() :: ([LemonAgent.Types.agent_message()],
                           reference()
                           | nil ->
                             {:ok, [LemonAgent.Types.agent_message()]}
                             | [LemonAgent.Types.agent_message()])

waiter()

@type waiter() :: {:call, GenServer.from()} | {:notify, pid(), reference()}

Functions

abort(agent)

@spec abort(GenServer.server()) :: :ok

Aborts the currently running stream.

This sends an abort signal to the running task, which will cause it to terminate as soon as possible. The abort is asynchronous - use wait_for_idle/1 to wait for the task to actually complete.

append_message(agent, message)

@spec append_message(GenServer.server(), LemonAgent.Types.agent_message()) :: :ok

Appends a message to the conversation.

child_spec(init_arg)

Returns a specification to start this module under a supervisor.

See Supervisor.

clear_all_queues(agent)

@spec clear_all_queues(GenServer.server()) :: :ok

Clears all pending steering and follow-up messages.

clear_follow_up_queue(agent)

@spec clear_follow_up_queue(GenServer.server()) :: :ok

Clears all pending follow-up messages.

clear_steering_queue(agent)

@spec clear_steering_queue(GenServer.server()) :: :ok

Clears all pending steering messages.

continue(agent)

@spec continue(GenServer.server()) ::
  :ok | {:error, :already_streaming | :no_messages | :cannot_continue}

Continues from the current context.

Used for retrying after overflow or continuing from existing messages. Returns {:error, :already_streaming} if already processing. Returns {:error, :no_messages} if there are no messages to continue from. Returns {:error, :cannot_continue} if the last message is from the assistant.

follow_up(agent, message, opts \\ [])

@spec follow_up(GenServer.server(), LemonAgent.Types.agent_message(), keyword()) ::
  :ok

Queues a follow-up message to be processed after the agent finishes.

Follow-up messages are delivered only when the agent has no more tool calls and no steering messages. Use this for messages that should wait until the agent completes its current work. Pass system_prompt: prompt to apply a per-turn prompt refresh when the queued message is consumed.

get_follow_up_mode(agent)

@spec get_follow_up_mode(GenServer.server()) :: queue_mode()

Gets the follow-up queue mode.

get_session_id(agent)

@spec get_session_id(GenServer.server()) :: String.t() | nil

Gets the current session ID.

get_state(agent)

@spec get_state(GenServer.server()) :: LemonAgent.Types.AgentState.t()

Gets the current agent state.

get_steering_mode(agent)

@spec get_steering_mode(GenServer.server()) :: queue_mode()

Gets the steering queue mode.

prompt(agent, message)

@spec prompt(
  GenServer.server(),
  String.t()
  | LemonAgent.Types.agent_message()
  | [LemonAgent.Types.agent_message()]
) :: :ok | {:error, :already_streaming}

Sends a prompt to the agent.

The prompt can be:

  • A string (converted to a user message)
  • An agent message struct
  • A list of agent message structs

Returns :ok immediately. Use wait_for_idle/1 to wait for completion. Returns {:error, :already_streaming} if a prompt is already being processed.

Examples

:ok = LemonAgent.Agent.prompt(agent, "Hello!")
:ok = LemonAgent.Agent.prompt(agent, %{role: :user, content: "Hi", timestamp: now})
:ok = LemonAgent.Agent.prompt(agent, [msg1, msg2])

replace_messages(agent, messages)

@spec replace_messages(GenServer.server(), [LemonAgent.Types.agent_message()]) :: :ok

Replaces all messages in the conversation.

reset(agent)

@spec reset(GenServer.server()) :: :ok

Resets the agent to initial state.

Clears all messages, queues, and error state. Does not change configuration like system_prompt, model, or tools.

set_follow_up_mode(agent, mode)

@spec set_follow_up_mode(GenServer.server(), queue_mode()) :: :ok

Sets the follow-up queue consumption mode.

  • :all - Send all follow-up messages at once
  • :one_at_a_time - Send one follow-up message per turn

set_model(agent, model)

@spec set_model(GenServer.server(), LemonAi.Types.Model.t()) :: :ok

Sets the model to use for LLM calls.

set_session_id(agent, session_id)

@spec set_session_id(GenServer.server(), String.t() | nil) :: :ok

Sets the session ID for provider caching.

set_steering_mode(agent, mode)

@spec set_steering_mode(GenServer.server(), queue_mode()) :: :ok

Sets the steering queue consumption mode.

  • :all - Send all steering messages at once
  • :one_at_a_time - Send one steering message per turn

set_system_prompt(agent, prompt)

@spec set_system_prompt(GenServer.server(), String.t()) :: :ok

Sets the system prompt.

set_thinking_level(agent, level)

@spec set_thinking_level(GenServer.server(), LemonAgent.Types.thinking_level()) :: :ok

Sets the thinking/reasoning level.

Valid levels: :off, :minimal, :low, :medium, :high, :xhigh

set_tools(agent, tools)

@spec set_tools(GenServer.server(), [LemonAgent.Types.AgentTool.t()]) :: :ok

Sets the available tools.

start_link(opts \\ [])

@spec start_link(opts()) :: GenServer.on_start()

Starts an Agent GenServer.

Options

  • :initial_state - Map merged into the default AgentState
  • :convert_to_llm - Function to convert agent messages to LLM messages (required)
  • :transform_context - Optional function to transform context before conversion
  • :stream_fn - Custom stream function (defaults to Loop's LemonAi.stream/3)
  • :steering_mode - How to consume steering queue: :all or :one_at_a_time (default)
  • :follow_up_mode - How to consume follow-up queue: :all or :one_at_a_time (default)
  • :session_id - Optional session identifier for provider caching
  • :get_api_key - Function to dynamically resolve API keys
  • :thinking_budgets - Map of thinking level budgets for token-based providers
  • :stream_options - StreamOptions for provider requests (temperature, max_tokens, etc.)
  • :max_tool_turns - Max assistant tool-use turns before terminal fallback
  • :tool_timeout_ms - Optional per-tool task timeout in milliseconds
  • :queue_call_timeout - Timeout for loop queue polling GenServer calls (:infinity or ms, default: 30 minutes)
  • :name - Optional GenServer name

Examples

{:ok, agent} = LemonAgent.Agent.start_link(
  initial_state: %{system_prompt: "You are helpful"},
  convert_to_llm: fn msgs -> Enum.filter(msgs, &llm_compatible?/1) end
)

steer(agent, message, opts \\ [])

@spec steer(GenServer.server(), LemonAgent.Types.agent_message(), keyword()) :: :ok

Queues a steering message to interrupt the agent mid-run.

Steering messages are delivered after the current tool execution completes, skipping remaining tool calls. Use this for "steering" the agent while it's working. Pass system_prompt: prompt to apply a per-turn prompt refresh when the queued message is consumed by an already-running loop.

subscribe(agent, subscriber_pid)

@spec subscribe(GenServer.server(), pid()) :: (-> :ok)

Subscribes a process to agent events.

The subscriber will receive {:agent_event, event} messages for each event emitted by the agent. The subscriber is automatically monitored and will be removed if it exits.

Returns an unsubscribe function that can be called to stop receiving events.

Examples

unsubscribe = LemonAgent.Agent.subscribe(agent, self())

receive do
  {:agent_event, %{type: :message_end} = event} ->
    IO.puts("Got message: #{inspect(event)}")
end

unsubscribe.()

wait_for_idle(agent, opts \\ [])

@spec wait_for_idle(GenServer.server(), keyword() | timeout()) ::
  :ok | {:error, :timeout}

Waits for the agent to become idle (no running task).

If the agent is already idle, returns immediately with :ok. If the agent is streaming, blocks until the current run completes.

Options

  • :timeout - Maximum time to wait in milliseconds (default: :infinity)

Examples

:ok = LemonAgent.Agent.wait_for_idle(agent)
:ok = LemonAgent.Agent.wait_for_idle(agent, timeout: 5000)