diff --git a/CHANGELOG.md b/CHANGELOG.md index 875eb66..028386b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ - Propagate stores to sub-agents and persist `parent_agent_id`, `agent_module`, and `started_at` metadata for reconstructing conversation trees - Add `Legion.Store.Postgres`, backed by an existing PostgreSQL Ecto repo, with partial upserts, `get/1`, `list/1`, configurable table names, configurable persistence frequency, and an optional `ecto_sql` dependency - Add versioned, idempotent `Legion.Store.Migration.Postgres` migrations with configurable table names and `pg_notify` notifications for inserts and updates; migration versions are tracked in the agents table comment; generated stores expose `__repo__/0` and `__table__/0` for database-backed consumers such as LegionWeb +- Add configurable per-request LLM usage persistence, enabled by default and disabled globally with `config :legion, :track_usage, false`.; each usage map is recursively string-keyed and `Legion.Store.Postgres` stores them as jsonb[]. - Bump the default model from `openai:gpt-4o-mini` to `openai:gpt-5.4` - `Legion.Tools.HumanTool.ask/1` now raises when called under `eval_and_complete` - the turn would end as soon as the code returns, silently discarding the human's answer; the error feeds back to the model, which retries under `eval_and_continue` diff --git a/lib/legion.ex b/lib/legion.ex index a5b499a..19643fb 100644 --- a/lib/legion.ex +++ b/lib/legion.ex @@ -5,9 +5,53 @@ defmodule Legion do |> String.split("") |> Enum.fetch!(1) + use Supervisor + alias Legion.{AgentIndex, AgentServer} alias Legion.Store.Payload + @doc """ + Starts Legion's supervisor. + + Add it to your application's supervision tree after dependencies required by + its configured recovery stores. For a Repo-backed store, place it after your + Repo. + + ## Examples + + defmodule MyApp.Application do + use Application + + def start(_type, _args) do + children = [ + MyApp.Repo, + {Legion, []} + ] + + Supervisor.start_link(children, strategy: :one_for_one, name: MyApp.Supervisor) + end + end + """ + def start_link(opts) when is_list(opts) do + Supervisor.start_link(__MODULE__, nil, name: __MODULE__) + end + + @doc false + def child_spec(opts) do + %{ + id: __MODULE__, + start: {__MODULE__, :start_link, [opts]}, + type: :supervisor + } + end + + @impl Supervisor + def init(_opts) do + children = [{Legion.Recovery, Application.fetch_env(:legion, :recovery)}] + + Supervisor.init(children, strategy: :one_for_one, name: Legion.Supervisor) + end + @doc """ Runs an agent on a single task and returns the result. @@ -46,7 +90,7 @@ defmodule Legion do {:ok, pid} = Legion.start_link(ChatAgent, store: MyApp.AgentStore, agent_id: "user_42:chat_7") {:ok, pid} = Legion.start_link(ChatAgent, agent_id: "user_42:chat_7") # store from app config """ - def start_link(agent_module, opts \\ []) do + def start_link(agent_module, opts \\ []) when is_atom(agent_module) do AgentServer.start_link(agent_module, opts) end diff --git a/lib/legion/agent_server.ex b/lib/legion/agent_server.ex index 411b980..6e1b818 100644 --- a/lib/legion/agent_server.ex +++ b/lib/legion/agent_server.ex @@ -25,6 +25,8 @@ defmodule Legion.AgentServer do :store, :agent_id, :persistence_frequency, + :track_usage, + usage: nil, executor_state: :nonexistent, bindings: [] ] @@ -58,11 +60,12 @@ defmodule Legion.AgentServer do agent_id = agent_id || generate_id() persistence_frequency = Store.persistence_frequency(store) + track_usage = Application.get_env(:legion, :track_usage, true) gen_opts = [name: Legion.AgentIndex.name(agent_id)] config = resolve_config(agent_module, opts) - {{agent_module, config, store, agent_id, persistence_frequency}, gen_opts} + {{agent_module, config, store, agent_id, persistence_frequency, track_usage}, gen_opts} end def call(agent, message, timeout \\ :infinity) do @@ -84,7 +87,7 @@ defmodule Legion.AgentServer do # Server callbacks @impl true - def init({agent_module, config, store, agent_id, persistence_frequency}) do + def init({agent_module, config, store, agent_id, persistence_frequency, track_usage}) do parent_agent_id = Vault.get(:agent_id) mode = Map.get(config, :start_mode, :normal) @@ -104,7 +107,7 @@ defmodule Legion.AgentServer do %{agent: agent_module} ) - {saved_messages, saved_bindings, saved_executor_state} = + {saved_messages, saved_bindings, saved_executor_state, saved_usage} = case store && store.get(agent_id) do {:ok, %Payload{ @@ -112,12 +115,13 @@ defmodule Legion.AgentServer do messages: messages, bindings: bindings, executor_state: executor_state - } + }, + usage: usage }} -> - {messages, bindings, executor_state} + {messages, bindings, executor_state, if(track_usage, do: usage || [], else: nil)} _no_state -> - {[], [], :nonexistent} + {[], [], :nonexistent, if(track_usage, do: [], else: nil)} end state = %__MODULE__{ @@ -128,14 +132,17 @@ defmodule Legion.AgentServer do agent_id: agent_id, persistence_frequency: persistence_frequency, bindings: saved_bindings, - executor_state: saved_executor_state + executor_state: saved_executor_state, + track_usage: track_usage, + usage: saved_usage } {:ok, persist(state, agent_module: state.agent_module, parent_agent_id: parent_agent_id, - started_at: NaiveDateTime.utc_now() + started_at: NaiveDateTime.utc_now(), + usage: state.usage ), {:continue, %{start_mode: mode, executor_state: saved_executor_state}}} end @@ -229,7 +236,7 @@ defmodule Legion.AgentServer do executor_config = Map.put(state.config, :checkpoint, checkpoint) - {status, value, final_messages, final_bindings} = + {status, value, final_messages, final_bindings, turn_usage} = Telemetry.span( [:legion, :agent, :message], %{agent: state.agent_module, message: state.messages |> List.last() |> Map.get(:content)}, @@ -239,7 +246,7 @@ defmodule Legion.AgentServer do initial_bindings = if conversation_scope?, do: state.bindings, else: [] - {status, value, messages, bindings} = + {status, value, messages, bindings, _turn_usage} = result = Executor.run( state.agent_module, @@ -250,15 +257,31 @@ defmodule Legion.AgentServer do ) iterations = Enum.count(messages, &(&1[:role] == "assistant")) - prev_count - {result, %{iterations: iterations, status: status, result: value, bindings: bindings}} + + {result, + %{ + iterations: iterations, + status: status, + result: value, + bindings: bindings + }} end ) kept_bindings = if conversation_scope?, do: final_bindings, else: [] + usage = if state.track_usage, do: state.usage ++ turn_usage + state = - %{state | messages: final_messages, bindings: kept_bindings} - |> persist([:conversation_state, status: :idle]) + %{ + state + | messages: final_messages, + bindings: kept_bindings, + usage: usage + } + + fields = [:conversation_state, status: :idle, usage: usage] + state = persist(state, fields) {{status, value}, state} end @@ -279,7 +302,7 @@ defmodule Legion.AgentServer do %{payload | conversation_state: persisted_conversation_state(checkpoint)} {field, value}, payload - when field in [:agent_module, :parent_agent_id, :status, :started_at] -> + when field in [:agent_module, :parent_agent_id, :status, :started_at, :usage] -> Map.put(payload, field, value) unknown, _payload -> diff --git a/lib/legion/executor.ex b/lib/legion/executor.ex index 91fd630..c38b3a6 100644 --- a/lib/legion/executor.ex +++ b/lib/legion/executor.ex @@ -129,71 +129,102 @@ defmodule Legion.Executor do `bindings` seeds the code-evaluation binding. `executor_state` resumes a step checkpoint when present: `:awaiting_llm` continues from its saved iteration and retry counters, while `:completing` finishes without another LLM request. - Pass `:nonexistent` to start a new loop. + Pass `nil` to start a new loop. `:turn_usage` is the complete, ordered list + of usage maps returned by LLM requests in the current turn. - Returns `{:ok, result, messages, bindings}` or `{:cancel, reason, messages, bindings}`. + Returns `{:ok, result, messages, bindings, turn_usage}` or + `{:cancel, reason, messages, bindings, turn_usage}`. """ def run(agent_module, messages, config, bindings \\ [], executor_state \\ :nonexistent) do config = Map.merge(@default_config, config) case executor_state do :nonexistent -> - loop(agent_module, messages, config, 0, 0, bindings) + loop(agent_module, messages, config, 0, 0, bindings, []) %{phase: :awaiting_llm, iteration: i, retries: r} -> - loop(agent_module, messages, config, i, r, bindings) + loop(agent_module, messages, config, i, r, bindings, []) %{phase: :completing, iteration: _i, retries: _r} -> - {:ok, nil, messages, bindings} + {:ok, nil, messages, bindings, []} end end - defp loop(agent_module, messages, config, iteration, retries, bindings) do + defp loop(agent_module, messages, config, iteration, retries, bindings, turn_usage) do if iteration >= config.max_iterations do - {:cancel, :reached_max_iterations, messages, bindings} + {:cancel, :reached_max_iterations, messages, bindings, turn_usage} else Telemetry.span( [:legion, :iteration], %{agent: agent_module, iteration: iteration}, - fn -> iterate(agent_module, messages, config, iteration, retries, bindings) end + fn -> + iterate(agent_module, messages, config, iteration, retries, bindings, turn_usage) + end ) end end - defp iterate(agent_module, messages, config, iteration, retries, bindings) do + defp iterate(agent_module, messages, config, iteration, retries, bindings, turn_usage) do # credo:disable-for-next-line - try do - with {:ok, action, messages} <- call_llm(agent_module, messages, config, iteration), - :ok <- validate_action_type(agent_module, action) do - result = - handle_action(agent_module, messages, config, action, iteration, retries, bindings) + llm_result = + try do + call_llm(agent_module, messages, config, iteration, turn_usage) + rescue + error -> {:error, error, turn_usage} + end - {result, %{action: action["action"]}} - else - {:error, reason} -> - result = - handle_execution_error( - agent_module, - messages, - config, - reason, - iteration, - retries, - bindings - ) + case llm_result do + {:ok, action, messages, turn_usage} -> + case validate_action_type(agent_module, action) do + :ok -> + result = + handle_action( + agent_module, + messages, + config, + action, + iteration, + retries, + bindings, + turn_usage + ) + + {result, %{action: action["action"]}} - {result, %{action: nil}} - end - rescue - e -> + {:error, reason} -> + result = + handle_execution_error( + agent_module, + messages, + config, + reason, + iteration, + retries, + bindings, + turn_usage + ) + + {result, %{action: nil}} + end + + {:error, reason, turn_usage} -> result = - handle_execution_error(agent_module, messages, config, e, iteration, retries, bindings) + handle_execution_error( + agent_module, + messages, + config, + reason, + iteration, + retries, + bindings, + turn_usage + ) {result, %{action: nil}} end end - defp call_llm(agent_module, messages, config, iteration) do + defp call_llm(agent_module, messages, config, iteration, turn_usage) do Telemetry.span( [:legion, :llm, :request], %{ @@ -205,17 +236,41 @@ defmodule Legion.Executor do fn -> case ReqLLM.generate_object(config.model, messages, action_schema(agent_module, config)) do {:ok, response} -> - action = extract_object(response) - messages = messages ++ [message(:assistant, Jason.encode!(action))] - {{:ok, action, messages}, %{object: action}} + handle_llm_response(response, messages, turn_usage) {:error, reason} -> - {{:error, "LLM request failed: #{inspect(reason)}"}, %{error: reason}} + {{:error, "LLM request failed: #{inspect(reason)}", turn_usage}, %{error: reason}} end end ) end + defp handle_llm_response(response, messages, turn_usage) do + turn_usage = turn_usage ++ [normalize_usage(response.usage)] + + case extract_object(response) do + {:ok, action} when is_map(action) -> + messages = messages ++ [message(:assistant, Jason.encode!(action))] + {{:ok, action, messages, turn_usage}, %{object: action}} + + {:error, reason} -> + {{:error, "LLM response object invalid: #{inspect(reason)}", turn_usage}, + %{error: reason}} + end + end + + defp normalize_usage(usage) when is_map(usage) do + Map.new(usage, fn {key, value} -> + {normalize_usage_key(key), normalize_usage(value)} + end) + end + + defp normalize_usage(usage) when is_list(usage), do: Enum.map(usage, &normalize_usage/1) + defp normalize_usage(usage), do: usage + + defp normalize_usage_key(key) when is_atom(key), do: Atom.to_string(key) + defp normalize_usage_key(key), do: key + defp checkpoint!(config, messages, bindings, executor_state) do case config[:checkpoint] do nil -> @@ -242,12 +297,22 @@ defmodule Legion.Executor do %{"action" => "return", "result" => result}, _i, _r, - bindings + bindings, + turn_usage ), - do: {:ok, result, messages, bindings} + do: {:ok, result, messages, bindings, turn_usage} - defp handle_action(_agent, messages, _config, %{"action" => "done"}, _i, _r, bindings), - do: {:ok, nil, messages, bindings} + defp handle_action( + _agent, + messages, + _config, + %{"action" => "done"}, + _i, + _r, + bindings, + turn_usage + ), + do: {:ok, nil, messages, bindings, turn_usage} defp handle_action( agent, @@ -256,7 +321,8 @@ defmodule Legion.Executor do %{"action" => eval, "code" => code}, i, retries, - bindings + bindings, + turn_usage ) when eval in ["eval_and_continue", "eval_and_complete"] and code != "" do # Tools that must see the answer come back to the model (e.g. HumanTool) @@ -280,15 +346,15 @@ defmodule Legion.Executor do checkpoint!(config, messages, new_bindings, executor_state) if eval == "eval_and_continue", - do: loop(agent, messages, config, i + 1, 0, new_bindings), - else: {:ok, result, messages, new_bindings} + do: loop(agent, messages, config, i + 1, 0, new_bindings, turn_usage), + else: {:ok, result, messages, new_bindings, turn_usage} {:error, error} -> - handle_execution_error(agent, messages, config, error, i, retries, bindings) + handle_execution_error(agent, messages, config, error, i, retries, bindings, turn_usage) end end - defp handle_action(agent, messages, config, action, i, retries, bindings), + defp handle_action(agent, messages, config, action, i, retries, bindings, turn_usage), do: handle_execution_error( agent, @@ -297,7 +363,8 @@ defmodule Legion.Executor do "Unexpected action: #{inspect(action)}", i, retries, - bindings + bindings, + turn_usage ) defp eval_in_span(agent_module, code, config, bindings) do @@ -344,9 +411,18 @@ defmodule Legion.Executor do end end - defp handle_execution_error(agent_module, messages, config, error, iteration, retries, bindings) do + defp handle_execution_error( + agent_module, + messages, + config, + error, + iteration, + retries, + bindings, + turn_usage + ) do if retries >= config.max_retries do - {:cancel, :reached_max_retries, messages, bindings} + {:cancel, :reached_max_retries, messages, bindings, turn_usage} else error_text = error |> format_error() |> truncate_content(config[:max_message_length]) @@ -367,7 +443,7 @@ defmodule Legion.Executor do retries: next_retries }) - loop(agent_module, messages, config, iteration, next_retries, bindings) + loop(agent_module, messages, config, iteration, next_retries, bindings, turn_usage) end end @@ -387,14 +463,16 @@ defmodule Legion.Executor do {:error, "Response missing required 'action' field, got: #{inspect(action)}"} end - defp extract_object(%{object: object}) when is_map(object), do: object + defp extract_object(%{object: object}) when is_map(object), do: {:ok, object} defp extract_object(%{message: %{tool_calls: tool_calls}}) when is_list(tool_calls) do - ReqLLM.ToolCall.find_args(tool_calls, "structured_output") || - raise "LLM response contained no structured object" + case ReqLLM.ToolCall.find_args(tool_calls, "structured_output") do + args when is_map(args) -> {:ok, args} + _ -> {:error, "LLM response contained no structured object"} + end end - defp extract_object(_response), do: raise("LLM response contained no structured object") + defp extract_object(_response), do: {:error, "LLM response contained no structured object"} defp format_result(result, bindings, config) do variable_names = bindings |> config.sandbox.binding_names() |> Enum.map(&"`#{&1}`") diff --git a/lib/legion/recovery.ex b/lib/legion/recovery.ex index db2f0bf..451b723 100644 --- a/lib/legion/recovery.ex +++ b/lib/legion/recovery.ex @@ -2,7 +2,7 @@ defmodule Legion.Recovery do @moduledoc """ Worker which runs the recovery process for interrupted runs. - `Legion.Application` starts this worker only when `:recovery` is configured: + Legion's supervisor starts this worker only when `:recovery` is configured: config :legion, :recovery, stores: [MyApp.AgentStore], diff --git a/lib/legion/store.ex b/lib/legion/store.ex index 7952f3d..f84c5fb 100644 --- a/lib/legion/store.ex +++ b/lib/legion/store.ex @@ -52,6 +52,17 @@ defmodule Legion.Store do started from within an agent's process tree picks up that store unless given an explicit `:store` of its own. + ## Usage tracking + + Legion persists a string-keyed copy of `ReqLLM.Response.usage` for every LLM + request in a conversation, ordered by request. Tracking is enabled by + default. Disable it globally before starting an agent: + + config :legion, :track_usage, false + + The setting is read when an agent starts. Disabled agents neither restore nor + update usage; existing usage in a store is preserved. + ## Identifying a conversation `:agent_id` is the key a conversation is saved under - it names one @@ -74,7 +85,8 @@ defmodule Legion.Store do `:nonexistent` for ordinary snapshots and is a map with `:phase`, `:iteration`, and `:retries` for step checkpoints. `:status` records whether the agent is mid-turn. The payload also carries the agent module, parent conversation, - and start time when those values are known. + and start time when those values are known. Its `:usage` field is the ordered + list of string-keyed LLM usage maps when tracking is enabled. With `binding_scope: :turn`, active bindings are included in step snapshots while the turn is running and cleared from the final snapshot. Bindings with diff --git a/lib/legion/store/migration/postgres/v01.ex b/lib/legion/store/migration/postgres/v01.ex index 8d863be..45a099b 100644 --- a/lib/legion/store/migration/postgres/v01.ex +++ b/lib/legion/store/migration/postgres/v01.ex @@ -13,6 +13,7 @@ defmodule Legion.Store.Migration.Postgres.V01 do add :status, :text, null: false, default: "idle" add :started_at, :naive_datetime_usec add :conversation_state, :binary + add :usage, {:array, :map}, null: false, default: [] add :inserted_at, :naive_datetime_usec, null: false, diff --git a/lib/legion/store/payload.ex b/lib/legion/store/payload.ex index 4c106b0..da4d718 100644 --- a/lib/legion/store/payload.ex +++ b/lib/legion/store/payload.ex @@ -18,7 +18,8 @@ defmodule Legion.Store.Payload do :parent_agent_id, :status, :started_at, - :conversation_state + :conversation_state, + :usage ] @type status :: :idle | :running @@ -32,7 +33,7 @@ defmodule Legion.Store.Payload do @type state :: %{ messages: [map()], bindings: keyword(), - executor_state: :nonexistent | executor_state() + executor_state: executor_state() | :nonexistent } @type t :: %__MODULE__{ @@ -41,6 +42,7 @@ defmodule Legion.Store.Payload do parent_agent_id: Legion.Store.agent_id() | nil, status: status() | nil, started_at: NaiveDateTime.t() | nil, - conversation_state: state() | nil + conversation_state: state() | nil, + usage: [map()] | nil } end diff --git a/lib/legion/store/postgres.ex b/lib/legion/store/postgres.ex index 70d1ed7..cba255c 100644 --- a/lib/legion/store/postgres.ex +++ b/lib/legion/store/postgres.ex @@ -47,19 +47,19 @@ defmodule Legion.Store.Postgres do per conversation, upserted on every save. Step snapshots therefore require no additional migration. - `save/1` performs partial upserts, so the same row carries the - conversation state and identity: `agent_module` (in `inspect/1` form, e.g. - `"MyApp.ResearchAgent"`), `parent_agent_id` linking a sub-agent to the - conversation that spawned it, and `started_at` as a UTC `NaiveDateTime` - stored with microsecond precision. Omitted payload fields preserve their - existing values. + `save/1` performs partial upserts, so a row carries the conversation state + and identity. Omitted payload fields preserve their existing values. The + `updated_at` timestamp is automatically set to the current UTC time on every save. + Only `agent_id` and `inserted_at` are never updated. The row's `status` flips to `'running'` when a turn starts and back to `'idle'` when it completes. Step writes update only the conversation state, leaving the running status unchanged. - `list/1` and `get/1` read persisted - conversations back from the same table. + Usage is stored as a `jsonb[]`: each element contains one complete, + string-keyed LLM usage map. + + `list/1` and `get/1` read persisted conversations back from the same table. The migration also installs a trigger that `pg_notify`s the table's channel (the table name) with the `agent_id` on every insert or update, so @@ -92,6 +92,7 @@ defmodule Legion.Store.Postgres do field :status, :string field :started_at, :naive_datetime_usec field :conversation_state, :binary + field :usage, {:array, :map} field :inserted_at, :naive_datetime_usec field :updated_at, :naive_datetime_usec end @@ -173,7 +174,8 @@ defmodule Legion.Store.Postgres do parent_agent_id: record.parent_agent_id, status: decode_status(record.status), started_at: record.started_at, - conversation_state: decode_conversation_state(record.conversation_state) + conversation_state: decode_conversation_state(record.conversation_state), + usage: record.usage } end diff --git a/mix.exs b/mix.exs index 7dc7a2a..5068beb 100644 --- a/mix.exs +++ b/mix.exs @@ -46,7 +46,6 @@ defmodule Legion.MixProject do def application do [ - mod: {Legion.Application, []}, extra_applications: [:logger] ] end diff --git a/test/integration/step_persistence_test.exs b/test/integration/step_persistence_test.exs index ef01eff..402e6f2 100644 --- a/test/integration/step_persistence_test.exs +++ b/test/integration/step_persistence_test.exs @@ -88,7 +88,8 @@ defmodule Legion.Integration.StepPersistenceTest do id: "test", model: "test", context: nil, - object: %{"action" => action, "code" => code, "result" => result} + object: %{"action" => action, "code" => code, "result" => result}, + usage: %{turn_usage: 0} }} end end diff --git a/test/legion/agent_server_test.exs b/test/legion/agent_server_test.exs index 07e722d..17bc6ae 100644 --- a/test/legion/agent_server_test.exs +++ b/test/legion/agent_server_test.exs @@ -41,33 +41,29 @@ defmodule Legion.AgentServerTest do @moduletag capture_log: true - defp llm_response(result) do - {:ok, - %ReqLLM.Response{ - id: "test", - model: "test", - context: nil, - object: %{"action" => "return", "code" => "", "result" => result} - }} + defp llm_response(result, turn_usage \\ 0) do + llm_object(%{"action" => "return", "code" => "", "result" => result}, turn_usage) end - defp llm_eval_response(code) do - {:ok, - %ReqLLM.Response{ - id: "test", - model: "test", - context: nil, - object: %{"action" => "eval_and_complete", "code" => code, "result" => ""} - }} + defp llm_eval_response(code, turn_usage \\ 0) do + llm_object(%{"action" => "eval_and_complete", "code" => code, "result" => ""}, turn_usage) + end + + defp llm_eval_continue_response(code, turn_usage \\ 0) do + llm_object( + %{"action" => "eval_and_continue", "code" => code, "result" => ""}, + turn_usage + ) end - defp llm_eval_continue_response(code) do + defp llm_object(object, turn_usage) do {:ok, %ReqLLM.Response{ id: "test", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => code, "result" => ""} + object: object, + usage: %{turn_usage: turn_usage} }} end @@ -593,7 +589,8 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, started_at: started_at, status: nil, - conversation_state: nil + conversation_state: nil, + usage: [] } = started assert is_struct(started_at, NaiveDateTime) @@ -621,6 +618,73 @@ defmodule Legion.AgentServerTest do refute Enum.any?(messages, &(&1.role == "system")) end + test "accumulates string-keyed usage across turns" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> llm_response("first", 7) + 2 -> llm_response("second", 11) + end + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-turns") + assert {:ok, "first"} = Legion.call(pid, "first turn") + assert {:ok, "second"} = Legion.call(pid, "second turn") + + assert {:ok, payload} = MemoryStore.get("usage-turns") + assert Map.get(payload, :usage) == [%{"turn_usage" => 7}, %{"turn_usage" => 11}] + end + + test "restored conversations add only new invocation usage" do + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "usage-restore", + usage: [%{turn_usage: 100}], + conversation_state: %{messages: [], bindings: [], executor_state: nil} + }) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("new work", 20) + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-restore") + assert {:ok, "new work"} = Legion.call(pid, "continue") + + assert {:ok, %Payload{usage: [%{turn_usage: 100}, %{"turn_usage" => 20}]}} = + MemoryStore.get("usage-restore") + end + + test "does not update usage when globally disabled" do + previous = Application.get_env(:legion, :track_usage, :unset) + Application.put_env(:legion, :track_usage, false) + + on_exit(fn -> + if previous == :unset, + do: Application.delete_env(:legion, :track_usage), + else: Application.put_env(:legion, :track_usage, previous) + end) + + assert :ok = + MemoryStore.save(%Payload{ + agent_id: "usage-disabled", + usage: [%{turn_usage: 100}], + conversation_state: %{messages: [], bindings: [], executor_state: nil} + }) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + llm_response("new work", 20) + end) + + {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "usage-disabled") + assert {:ok, "new work"} = Legion.call(pid, "continue") + + assert {:ok, %Payload{usage: [%{turn_usage: 100}]}} = + MemoryStore.get("usage-disabled") + end + test "saves a snapshot before the caller receives its reply" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> llm_response("Paris") @@ -992,7 +1056,7 @@ defmodule Legion.AgentServerTest do test "resume/2 returns the recorded process while it is alive" do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-live") - assert Legion.running?(pid) + assert Process.alive?(pid) assert {:ok, ^pid} = Legion.resume("resume-live", store: MemoryStore) end @@ -1012,11 +1076,11 @@ defmodule Legion.AgentServerTest do {:ok, pid} = Legion.start_link(MathAgent, store: MemoryStore, agent_id: "resume-dead") {:ok, _} = Legion.call(pid, "What is the capital of France?") GenServer.stop(pid) - refute Legion.running?(pid) + refute Process.alive?(pid) {:ok, revived} = Legion.resume("resume-dead", store: MemoryStore) - assert Legion.running?(revived) + assert Process.alive?(revived) assert [ %{role: "system"}, @@ -1031,6 +1095,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-awaiting-llm", agent_module: MathAgent, status: :running, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1061,6 +1126,7 @@ defmodule Legion.AgentServerTest do agent_id: "resume-completing", agent_module: MathAgent, status: :running, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1118,6 +1184,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :running, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1189,6 +1256,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: nil, agent_module: MathAgent, status: :idle, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], @@ -1206,6 +1274,7 @@ defmodule Legion.AgentServerTest do parent_agent_id: "recover-parent", agent_module: MathAgent, status: :running, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "compute"}], bindings: [x: 42], diff --git a/test/legion/executor_test.exs b/test/legion/executor_test.exs index dd70e1b..be63d12 100644 --- a/test/legion/executor_test.exs +++ b/test/legion/executor_test.exs @@ -72,8 +72,15 @@ defmodule Legion.ExecutorTest do @moduletag capture_log: true - defp response(object) do - {:ok, %ReqLLM.Response{id: "test", model: "test", context: nil, object: object}} + defp response(object, turn_usage \\ 0) do + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: object, + usage: %{turn_usage: turn_usage} + }} end defp executor_messages(message) do @@ -83,7 +90,98 @@ defmodule Legion.ExecutorTest do ] end - describe "run/4" do + describe "run/3-5" do + test "returns recursively string-keyed usage from a single LLM request" do + usage = %{ + input_tokens: 12, + output_tokens: 5, + turn_usage: 17, + tool_usage: %{web_search: 1} + } + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "return", "code" => "", "result" => "42"}, + usage: usage + }} + end) + + assert {:ok, "42", _messages, [], + [ + %{ + "input_tokens" => 12, + "output_tokens" => 5, + "turn_usage" => 17, + "tool_usage" => %{"web_search" => 1} + } + ]} = Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "returns usage list from a single LLM request" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + response(%{"action" => "return", "code" => "", "result" => "42"}, 17) + end) + + assert {:ok, "42", _messages, [], [%{"turn_usage" => 17}]} = + Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "preserves provider usage values while stringifying keys" do + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + {:ok, + %ReqLLM.Response{ + id: "test", + model: "test", + context: nil, + object: %{"action" => "return", "code" => "", "result" => "42"}, + usage: %{input_tokens: 12, output_tokens: 5} + }} + end) + + assert {:ok, "42", _messages, [], [%{"input_tokens" => 12, "output_tokens" => 5}]} = + Legion.Executor.run(MathAgent, executor_messages("what is 42?"), %{}) + end + + test "preserves usage order across a multi-response turn" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> response(%{"action" => "eval_and_continue", "code" => "x = 10", "result" => ""}, 7) + 2 -> response(%{"action" => "return", "code" => "", "result" => "done"}, 11) + end + end) + + assert {:ok, "done", _messages, _bindings, [%{"turn_usage" => 7}, %{"turn_usage" => 11}]} = + Legion.Executor.run( + MathAgent, + executor_messages("compute"), + %{} + ) + end + + test "retains usage from an invalid response while retrying" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> response(nil, 7) + 2 -> response(%{"action" => "return", "code" => "", "result" => "recovered"}, 11) + end + end) + + assert {:ok, "recovered", _messages, [], [%{"turn_usage" => 7}, %{"turn_usage" => 11}]} = + Legion.Executor.run(MathAgent, executor_messages("recover"), %{}) + end + test "returns result for return action" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{"action" => "return", "code" => "", "result" => "42"}) @@ -216,6 +314,22 @@ defmodule Legion.ExecutorTest do assert {:ok, "recovered"} = Legion.execute(MathAgent, "retry me") end + test "raised LLM exception triggers retry without adding usage" do + call_count = :counters.new(1, [:atomics]) + + stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> + :counters.add(call_count, 1, 1) + + case :counters.get(call_count, 1) do + 1 -> raise "provider exploded" + 2 -> response(%{"action" => "return", "code" => "", "result" => "recovered"}, 11) + end + end) + + assert {:ok, "recovered", _messages, [], [%{"turn_usage" => 11}]} = + Legion.Executor.run(MathAgent, executor_messages("retry raised error"), %{}) + end + test "third-party tool module without extra_allowed_modules/0 does not crash eval" do stub(ReqLLM, :generate_object, fn _model, _messages, _schema -> response(%{ @@ -275,7 +389,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, 20, _messages, _bindings} = + assert {:ok, 20, _messages, _bindings, _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("compute"), @@ -326,7 +440,7 @@ defmodule Legion.ExecutorTest do :ok end - assert {:ok, "recovered", _messages, []} = + assert {:ok, "recovered", _messages, [], _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("recover"), @@ -351,7 +465,7 @@ defmodule Legion.ExecutorTest do response(%{"action" => "return", "code" => "", "result" => "done"}) end) - assert {:ok, "done", _messages, []} = + assert {:ok, "done", _messages, [], _turn_usage} = Legion.Executor.run( MathAgent, executor_messages("finish"), diff --git a/test/legion/parallel_and_pipeline_test.exs b/test/legion/parallel_and_pipeline_test.exs index 588052f..75cfe59 100644 --- a/test/legion/parallel_and_pipeline_test.exs +++ b/test/legion/parallel_and_pipeline_test.exs @@ -20,7 +20,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "test", model: "test", context: nil, - object: %{"action" => "return", "code" => "", "result" => result} + object: %{"action" => "return", "code" => "", "result" => result}, + usage: %{turn_usage: 0} }} end @@ -54,7 +55,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "cancel", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => "return 1 + 1", "result" => ""} + object: %{"action" => "eval_and_continue", "code" => "return 1 + 1", "result" => ""}, + usage: %{turn_usage: 0} }} end end) @@ -134,7 +136,8 @@ defmodule Legion.ParallelAndPipelineTest do id: "cancel", model: "test", context: nil, - object: %{"action" => "eval_and_continue", "code" => "return 1 + 1", "result" => ""} + object: %{"action" => "eval_and_continue", "code" => "return 1 + 1", "result" => ""}, + usage: %{turn_usage: 0} }} end) diff --git a/test/legion/recovery_test.exs b/test/legion/recovery_test.exs index 592ae74..478c147 100644 --- a/test/legion/recovery_test.exs +++ b/test/legion/recovery_test.exs @@ -243,6 +243,7 @@ defmodule Legion.RecoveryTest do parent_agent_id: nil, agent_module: RecoveryAgent, status: :running, + usage: [], conversation_state: %{ messages: [%{role: "user", type: :user, content: "recover me"}], bindings: [], @@ -257,7 +258,8 @@ defmodule Legion.RecoveryTest do id: "test", model: "test", context: nil, - object: %{"action" => "return", "code" => "", "result" => result} + object: %{"action" => "return", "code" => "", "result" => result}, + usage: %{turn_usage: 0} }} end end diff --git a/test/legion/store/postgres_db_test.exs b/test/legion/store/postgres_db_test.exs index 92928e1..9ac3369 100644 --- a/test/legion/store/postgres_db_test.exs +++ b/test/legion/store/postgres_db_test.exs @@ -14,6 +14,36 @@ defmodule Legion.Store.PostgresDbTest do :ok end + test "stores usage as a jsonb array" do + payload = %Payload{ + agent_id: "usage-jsonb", + usage: [ + %{input_tokens: 12, output_tokens: 5, turn_usage: 17, tool_usage: %{web_search: 1}}, + %{input_tokens: 7, output_tokens: 3, turn_usage: 10} + ] + } + + assert :ok = Store.save(payload) + + assert {:ok, + %Payload{ + usage: [ + %{ + "input_tokens" => 12, + "output_tokens" => 5, + "turn_usage" => 17, + "tool_usage" => %{"web_search" => 1} + }, + %{"input_tokens" => 7, "output_tokens" => 3, "turn_usage" => 10} + ] + }} = Store.get("usage-jsonb") + + assert %{rows: [["jsonb[]"]]} = + Repo.query!("SELECT pg_typeof(usage)::text FROM legion_agents WHERE agent_id = $1", [ + "usage-jsonb" + ]) + end + test "save/1 fully inserts every payload field" do payload = %Payload{ agent_id: "user_42", @@ -25,11 +55,14 @@ defmodule Legion.Store.PostgresDbTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{turn_usage: 100}] } + expected_payload = %{payload | usage: [%{"turn_usage" => 100}]} + assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("user_42") + assert {:ok, ^expected_payload} = Store.get("user_42") end test "save/1 partially inserts only the supplied payload fields" do @@ -44,7 +77,7 @@ defmodule Legion.Store.PostgresDbTest do assert :ok = Store.save(payload) assert {:ok, stored} = Store.get("state-only") - assert stored == %{payload | status: :idle} + assert stored == %{payload | status: :idle, usage: []} end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -58,7 +91,8 @@ defmodule Legion.Store.PostgresDbTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{turn_usage: 100}] } assert :ok = Store.save(initial) @@ -79,7 +113,8 @@ defmodule Legion.Store.PostgresDbTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{"turn_usage" => 100}] }} = Store.get("user_42") %{rows: [[updated_at]]} = diff --git a/test/legion/store/postgres_test.exs b/test/legion/store/postgres_test.exs index e255baa..2c0225b 100644 --- a/test/legion/store/postgres_test.exs +++ b/test/legion/store/postgres_test.exs @@ -37,6 +37,7 @@ defmodule Legion.Store.PostgresTest do status: "idle", started_at: nil, conversation_state: nil, + usage: [], inserted_at: nil, updated_at: nil } @@ -74,7 +75,8 @@ defmodule Legion.Store.PostgresTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{turn_usage: 100}] } assert :ok = Store.save(payload) @@ -91,7 +93,7 @@ defmodule Legion.Store.PostgresTest do } } - expected_payload = %{payload | status: :idle} + expected_payload = %{payload | status: :idle, usage: []} assert :ok = Store.save(payload) assert {:ok, ^expected_payload} = Store.get("state-only") end @@ -110,7 +112,9 @@ defmodule Legion.Store.PostgresTest do } assert :ok = Store.save(payload) - assert {:ok, ^payload} = Store.get("step-state") + + expected_payload = %{payload | usage: []} + assert {:ok, ^expected_payload} = Store.get("step-state") end test "save/1 partial upsert preserves omitted fields and advances updated_at" do @@ -124,7 +128,8 @@ defmodule Legion.Store.PostgresTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{turn_usage: 100}] } assert :ok = Store.save(initial) @@ -142,7 +147,8 @@ defmodule Legion.Store.PostgresTest do messages: [%{role: "user", content: "hi"}], bindings: [x: 42], executor_state: :nonexistent - } + }, + usage: [%{turn_usage: 100}] }} = Store.get("user_42") assert NaiveDateTime.compare(FakeRepo.run("user_42").updated_at, previous_updated_at) == :gt diff --git a/test/legion/supervisor_test.exs b/test/legion/supervisor_test.exs new file mode 100644 index 0000000..5ebe26f --- /dev/null +++ b/test/legion/supervisor_test.exs @@ -0,0 +1,76 @@ +defmodule Legion.SupervisorTest do + use ExUnit.Case, async: false + + defmodule FakeRepo do + use GenServer + + def start_link(_opts), do: GenServer.start_link(__MODULE__, :ok, name: __MODULE__) + + @impl GenServer + def init(:ok), do: {:ok, :ready} + end + + defmodule RecoveryStore do + def list(_limit) do + send( + Process.whereis(:legion_supervisor_test), + {:recovery_scanned, Process.whereis(FakeRepo)} + ) + + [] + end + end + + defmodule TestAgent do + @moduledoc "Agent used to verify Legion's public startup function." + use Legion.Agent + end + + setup do + previous = Application.fetch_env(:legion, :recovery) + Process.register(self(), :legion_supervisor_test) + + on_exit(fn -> + case previous do + {:ok, config} -> Application.put_env(:legion, :recovery, config) + :error -> Application.delete_env(:legion, :recovery) + end + end) + end + + test "does not auto-start a Legion supervisor" do + assert [] = Application.spec(:legion, :mod) + end + + test "uses start_link when embedded as a child" do + assert %{id: Legion, start: {Legion, :start_link, [[]]}, type: :supervisor} = + Legion.child_spec([]) + end + + test "starts an agent when start_link receives an agent module" do + assert {:ok, pid} = Legion.start_link(TestAgent) + assert is_binary(Legion.get_agent_id(pid)) + end + + test "adds recovery worker with configured options" do + config = [stores: [RecoveryStore], store_scan_limit: 3, concurrent_request_limit: 2] + Application.put_env(:legion, :recovery, config) + + assert {:ok, + {_supervisor_flags, + [%{id: Legion.Recovery, start: {Legion.Recovery, :start_link, [{:ok, ^config}]}}]}} = + Legion.init([]) + end + + test "starts recovery after a client repo" do + Application.put_env(:legion, :recovery, stores: [RecoveryStore], store_scan_limit: 1) + + start_supervised!(%{ + id: :client_supervisor, + start: {Supervisor, :start_link, [[FakeRepo, {Legion, []}], [strategy: :one_for_one]]} + }) + + assert_receive {:recovery_scanned, repo_pid} + assert is_pid(repo_pid) + end +end