Added energy tracking features!

app.py

@@ -1470,6 +1470,57 @@ def validate_azure_openai_key(api_key):
                         f"${total_cost:.4f}",
                         f"{remote_total:,} total tokens",
                     )
+
+            # Energy usage for both protocols
+            if ("remote_only_energy" in output) and ("minion_local_energy" in output) and ("minion_remote_energy" in output):
+                st.header("Energy Usage")
+                st.markdown("This is the amount of energy saved by running part of the computation locally. We estimate remote energy consumption based on the number of input/output tokens as well as assumptions about the remote model architecture.")
+
+                minion_local_energy = output["minion_local_energy"]
+                minion_remote_energy = output["minion_remote_energy"]
+                minion_energy = minion_local_energy + minion_remote_energy
+
+                remote_only_energy = output["remote_only_energy"]
+
+                energy_savings = remote_only_energy - minion_energy
+
+                c1, c2, c3 = st.columns(3)
+                c1.metric(
+                    f"Minion(s) Energy Consumption",
+                    f"{minion_energy:.2f} J",
+                    None
+                )
+                c2.metric(
+                    f"Remote-Only Energy Consumption",
+                    f"{remote_only_energy:.2f} J",
+                    None
+                )
+                c3.metric(
+                    f"Total Energy Savings!",
+                    f"{energy_savings:.2f} J",
+                    None
+                )
+                # Convert to long format DataFrame for explicit ordering
+                df = pd.DataFrame(
+                    {
+                        "Protocol": [
+                            f"Minion(s)",
+                            f"Minion(s)",
+                            f"Remote-Only",
+                        ],
+                        "Model Type": [
+                            "Remote",
+                            "Local",
+                            "Remote",
+                        ],
+                        "Energy": [
+                            minion_remote_energy,
+                            minion_local_energy,
+                            remote_only_energy,
+                        ],
+                    }
+                )
+                st.bar_chart(df, x="Protocol", y="Energy", color="Model Type")
 
             # Display meta information for minions protocol
             if "meta" in output:

minions/minion.py

@@ -6,6 +6,10 @@
 
 from minions.clients import OpenAIClient, TogetherClient
 
+from minions.usage import Usage
+from minions.utils.energy_tracking import PowerMonitor
+from minions.utils.energy_tracking import cloud_inference_energy_estimate, better_cloud_inference_energy_estimate
+
 from minions.prompts.minion import (
     SUPERVISOR_CONVERSATION_PROMPT,
     SUPERVISOR_FINAL_PROMPT,
@@ -16,7 +20,6 @@
     WORKER_PRIVACY_SHIELD_PROMPT,
     REFORMAT_QUERY_PROMPT,
 )
-from minions.usage import Usage
 
 
 def _escape_newlines_in_strings(json_str: str) -> str:
@@ -101,6 +104,10 @@ def __call__(
             Dict containing final_answer, conversation histories, and usage statistics
         """
 
+        # Setup and start power monitor
+        monitor = PowerMonitor(mode="auto", interval=1.0)
+        monitor.start()
+
         if max_rounds is None:
             max_rounds = self.max_rounds
 
@@ -138,6 +145,9 @@ def __call__(
         remote_usage = Usage()
         local_usage = Usage()
 
+        remote_usage_tokens = []
+        local_usage_tokens = []
+
         worker_messages = []
         supervisor_messages = []
 
@@ -161,6 +171,8 @@ def __call__(
                 messages=[{"role": "user", "content": reformat_query_task}]
             )
             local_usage += usage
+            local_usage_tokens.append((usage.prompt_tokens, usage.completion_tokens))
+
             pii_reformatted_task = reformatted_task[0]
 
             # Log the reformatted task
@@ -218,6 +230,8 @@ def __call__(
             )
 
         remote_usage += supervisor_usage
+        remote_usage_tokens.append((supervisor_usage.prompt_tokens, supervisor_usage.completion_tokens))
+
         supervisor_messages.append(
             {"role": "assistant", "content": supervisor_response[0]}
         )
@@ -256,6 +270,7 @@ def __call__(
             )
 
             local_usage += worker_usage
+            local_usage_tokens.append((worker_usage.prompt_tokens, worker_usage.completion_tokens))
 
             if is_privacy:
                 if self.callback:
@@ -273,6 +288,7 @@ def __call__(
                     messages=[{"role": "user", "content": worker_privacy_shield_prompt}]
                 )
                 local_usage += worker_usage
+                local_usage_tokens.append((worker_usage.prompt_tokens, worker_usage.completion_tokens))
 
                 worker_messages.append(
                     {"role": "assistant", "content": worker_response[0]}
@@ -323,6 +339,7 @@ def __call__(
                 )
 
                 remote_usage += usage
+                remote_usage_tokens.append((usage.prompt_tokens, usage.completion_tokens))
 
                 supervisor_messages.append(
                     {"role": "assistant", "content": step_by_step_response[0]}
@@ -360,6 +377,7 @@ def __call__(
                 )
 
             remote_usage += supervisor_usage
+            remote_usage_tokens.append((supervisor_usage.prompt_tokens, supervisor_usage.completion_tokens))
             supervisor_messages.append(
                 {"role": "assistant", "content": supervisor_response[0]}
             )
@@ -408,11 +426,65 @@ def __call__(
         with open(log_path, "w", encoding="utf-8") as f:
             json.dump(conversation_log, f, indent=2, ensure_ascii=False)
 
+        # Stop tracking power
+        monitor.stop()
+
+        # Estimate energy consumption over entire conversation/query
+        use_better_estimate = True
+
+        # Estimate local minion energy consumption
+        final_estimates = monitor.get_final_estimates()
+        minion_local_energy = float(final_estimates["Measured Energy"][:-2])
+
+        # Estimate remote energy consumption (minion and remote-only)
+        minion_remote_energy = None
+        remote_only_energy = None
+
+        if use_better_estimate:
+            minion_remote_energy = 0
+            for (input_tokens, output_tokens) in remote_usage_tokens:
+                estimate = better_cloud_inference_energy_estimate(
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens
+                )
+                minion_remote_energy += estimate["total_energy_joules"]
+
+            remote_only_energy = minion_remote_energy
+            for (input_tokens, output_tokens) in local_usage_tokens:
+                estimate = better_cloud_inference_energy_estimate(
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens
+                )
+                remote_only_energy += estimate["total_energy_joules"]
+
+        else:
+            # Local/remote input/output tokens
+            local_input_tokens = local_usage.prompt_tokens
+            local_output_tokens = local_usage.completion_tokens
+            remote_input_tokens = remote_usage.prompt_tokens
+            remote_output_tokens = remote_usage.completion_tokens
+
+            total_input_tokens = local_input_tokens + remote_input_tokens
+            total_output_tokens = local_output_tokens + remote_output_tokens
+
+            # Estimate remote-only energy consumption (remote processes all input/output tokens)
+            _, remote_only_energy, _ = cloud_inference_energy_estimate(
+                tokens=total_output_tokens+total_input_tokens
+            )
+
+            # Estimate minion energy consumption (including both remote and local energy consumption)
+            _, minion_remote_energy, _ = cloud_inference_energy_estimate(
+                tokens=remote_output_tokens+remote_input_tokens,
+            )
+
         return {
             "final_answer": final_answer,
             "supervisor_messages": supervisor_messages,
             "worker_messages": worker_messages,
             "remote_usage": remote_usage,
             "local_usage": local_usage,
             "log_file": log_path,
+            "remote_only_energy": remote_only_energy,
+            "minion_local_energy": minion_local_energy,
+            "minion_remote_energy": minion_remote_energy
         }

minions/minions.py

@@ -8,6 +8,8 @@
 import numpy as np
 
 from minions.usage import Usage
+from minions.utils.energy_tracking import PowerMonitor
+from minions.utils.energy_tracking import cloud_inference_energy_estimate
 
 from minions.prompts.minions import (
     WORKER_PROMPT_TEMPLATE,
@@ -258,6 +260,10 @@ def __call__(
             Dict containing final_answer and conversation histories
         """
 
+        # Setup and start power monitor
+        monitor = PowerMonitor(mode="auto", interval=1.0)
+        monitor.start()
+
         self.max_rounds = max_rounds or self.max_rounds
 
         # Initialize usage tracking
@@ -732,10 +738,37 @@ def filter_fn(job: Job) -> bool:
                 f"Exhausted all rounds without finding a final answer. Returning the last synthesized response."
             )
             final_answer = "No answer found."
+
+        # Stop tracking power
+        monitor.stop()
+
+        # Local/remote input/output tokens
+        local_input_tokens = local_usage.prompt_tokens
+        local_output_tokens = local_usage.completion_tokens
+        remote_input_tokens = remote_usage.prompt_tokens
+        remote_output_tokens = remote_usage.completion_tokens
+
+        total_input_tokens = local_input_tokens + remote_input_tokens
+        total_output_tokens = local_output_tokens + remote_output_tokens
+
+        # Estimate remote-only energy consumption (remote processes all input/output tokens)
+        _, remote_only_energy, _ = cloud_inference_energy_estimate(
+            tokens=total_output_tokens
+        )
+
+        # Estimate minion energy consumption (including both remote and local energy consumption)
+        final_estimates = monitor.get_final_estimates()
+        minion_local_energy = float(final_estimates["Measured Energy"][:-2])
+        _, minion_remote_energy, _ = cloud_inference_energy_estimate(
+            tokens=remote_output_tokens,
+        )
 
         return {
             "final_answer": final_answer,
             "meta": meta,
             "local_usage": local_usage,
             "remote_usage": remote_usage,
+            "remote_only_energy": remote_only_energy,
+            "minion_local_energy": minion_local_energy,
+            "minion_remote_energy": minion_remote_energy
         }