Energy Middleware
The EnergyMiddleware tracks token usage, estimates energy consumption, and calculates CO2 emissions for agent calls.
- class jamanota.middleware.EnergyMiddleware[source]
Bases:
AgentMiddlewareMiddleware that tracks token usage, estimates energy consumption, and calculates associated CO2 emissions for agent model calls.
This middleware maintains a thread-safe list of datapoints and supports nested agent calls using a prompt ID stack.
- after_agent(state: AgentState, runtime: Runtime) dict[str, Any] | None[source]
Hook executed after an agent finishes.
Removes the latest prompt ID from the stack.
- Parameters:
state (AgentState) – Current agent state.
runtime (Runtime) – Runtime context.
- Returns:
Optional state updates (unused).
- Return type:
dict[str, Any] | None
- after_model(state: AgentState, runtime: Runtime) dict[str, Any] | None[source]
Hook executed after a model generates a response.
Extracts token usage, estimates energy and emissions, and records a datapoint for a model call.
- Parameters:
state (AgentState) – Current agent state containing messages.
runtime (Runtime) – Runtime context.
- Returns:
Optional state updates (unused).
- Return type:
dict[str, Any] | None
- before_agent(state: AgentState, runtime: Runtime) dict[str, Any] | None[source]
Hook executed before an agent runs.
Ensures that a prompt ID is assigned and propagated through nested agent calls.
- Parameters:
state (AgentState) – Current agent state.
runtime (Runtime) – Runtime context.
- Returns:
Optional state updates (unused).
- Return type:
dict[str, Any] | None
- get_prompt_count() int[source]
Number of top-level prompts seen so far.
- Returns:
Count of unique top-level prompts.
- Return type:
int
- get_report() list[EnergyDataPoint][source]
Retrieve a copy of all collected datapoints.
- Returns:
A copy of the recorded datapoints.
- Return type:
list[EnergyDataPoint]
- get_summary(group_by: Literal['model_name', 'agent_name'], last_n_prompts: int | None = None, last_n_hours: float | None = None) list[EnergyGroupSummary][source]
Returns aggregated energy/token summaries grouped by model or agent.
- Parameters:
group_by – “model_name” or “agent_name”
last_n_prompts – if set, only include data from the last N top-level prompts
last_n_hours – if set, only include data from the last N hours
- get_total_co2() float[source]
Get the total estimated CO2 emissions of all collected datapoints in kilograms.
- Returns:
Total estimated CO2 emissions in kilograms.
- Return type:
float
- get_total_energy() float[source]
Get the total estimated energy consumption of all collected datapoints in Joules.
- Returns:
Total estimated energy in Joules.
- Return type:
float
- get_total_input_tokens() int[source]
Get the total number of input tokens across all collected datapoints.
- Returns:
Total number of input tokens.
- Return type:
int