feat(tracking): add ecologits carbon/energy metrics to LLM tracking

pyproject.toml

@@ -59,7 +59,10 @@ hint = [
     "sentence-transformers>=5.0.0",
 ]
 transformers = [
-    "transformers>=4.38.2", 
+    "transformers>=4.38.2",
+]
+ecologits = [
+    "ecologits",
 ]
 
 

src/agentlab/llm/chat_api.py

@@ -288,6 +288,7 @@ def __init__(
         )
 
     def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict:
+        tracking._ensure_ecologits_init()
         # Initialize retry tracking attributes
         self.retries = 0
         self.success = False
@@ -328,11 +329,12 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
         input_tokens = completion.usage.prompt_tokens
         output_tokens = completion.usage.completion_tokens
         cost = input_tokens * self.input_cost + output_tokens * self.output_cost
+        energy_kwh, gwp_kgco2eq = tracking._extract_ecologits_impacts(completion)
 
         if hasattr(tracking.TRACKER, "instance") and isinstance(
             tracking.TRACKER.instance, tracking.LLMTracker
         ):
-            tracking.TRACKER.instance(input_tokens, output_tokens, cost)
+            tracking.TRACKER.instance(input_tokens, output_tokens, cost, energy_kwh, gwp_kgco2eq)
 
         if n_samples == 1:
             res = AIMessage(completion.choices[0].message.content)
@@ -530,6 +532,7 @@ def __init__(
             self.output_cost = 0.0
 
     def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict:
+        tracking._ensure_ecologits_init()
         # Convert OpenAI format to Anthropic format
         system_message = None
         anthropic_messages = []
@@ -574,11 +577,15 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                     + cache_write_tokens * cache_write_cost
                 )
 
+                energy_kwh, gwp_kgco2eq = tracking._extract_ecologits_impacts(response)
+
                 # Track usage if available
                 if hasattr(tracking.TRACKER, "instance") and isinstance(
                     tracking.TRACKER.instance, tracking.LLMTracker
                 ):
-                    tracking.TRACKER.instance(new_input_tokens, output_tokens, cost)
+                    tracking.TRACKER.instance(
+                        new_input_tokens, output_tokens, cost, energy_kwh, gwp_kgco2eq
+                    )
 
                 return AIMessage(response.content[0].text)
 
@@ -681,6 +688,7 @@ def __init__(
             self.output_cost = 0.0
 
     def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict:
+        tracking._ensure_ecologits_init()
         from litellm import completion as litellm_completion
 
         # Initialize retry tracking attributes
@@ -719,11 +727,12 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
         input_tokens = completion.usage.prompt_tokens
         output_tokens = completion.usage.completion_tokens
         cost = input_tokens * self.input_cost + output_tokens * self.output_cost
+        energy_kwh, gwp_kgco2eq = tracking._extract_ecologits_impacts(completion)
 
         if hasattr(tracking.TRACKER, "instance") and isinstance(
             tracking.TRACKER.instance, tracking.LLMTracker
         ):
-            tracking.TRACKER.instance(input_tokens, output_tokens, cost)
+            tracking.TRACKER.instance(input_tokens, output_tokens, cost, energy_kwh, gwp_kgco2eq)
 
         if n_samples == 1:
             res_text = completion.choices[0].message.content

src/agentlab/llm/tracking.py

@@ -17,10 +17,41 @@
 else:
     bedrock_anthropic_callback = None
     openai_info = None
+
+_ecologits_available = importlib.util.find_spec("ecologits") is not None
+if _ecologits_available:
+    from ecologits import EcoLogits
+
 from litellm import completion_cost, get_model_info
 
 TRACKER = threading.local()
 
+_ecologits_initialized = False
+
+
+def _ensure_ecologits_init():
+    """Initialize ecologits if available and not yet initialized."""
+    global _ecologits_initialized
+    if _ecologits_available and not _ecologits_initialized:
+        EcoLogits.init()
+        _ecologits_initialized = True
+
+
+def _extract_ecologits_impacts(response):
+    """Extract energy and GWP impacts from an ecologits-enriched response.
+
+    Returns (energy_kwh, gwp_kgco2eq) or (0.0, 0.0) if not available.
+    """
+    if not _ecologits_available:
+        return 0.0, 0.0
+    impacts = getattr(response, "impacts", None)
+    if impacts is None:
+        return 0.0, 0.0
+    energy_kwh = float(getattr(impacts.energy, "value", 0.0))
+    gwp_kgco2eq = float(getattr(impacts.gwp, "value", 0.0))
+    return energy_kwh, gwp_kgco2eq
+
+
 ANTHROPIC_CACHE_PRICING_FACTOR = {
     "cache_read_tokens": 0.1,  # Cost for 5 min ephemeral cache. See Pricing Here: https://docs.anthropic.com/en/docs/about-claude/pricing#model-pricing
     "cache_write_tokens": 1.25,
@@ -37,28 +68,53 @@ def __init__(self, suffix=""):
         self.input_tokens = 0
         self.output_tokens = 0
         self.cost = 0.0
+        self.energy_kwh = 0.0
+        self.gwp_kgco2eq = 0.0
         self.input_tokens_key = "input_tokens_" + suffix if suffix else "input_tokens"
         self.output_tokens_key = "output_tokens_" + suffix if suffix else "output_tokens"
         self.cost_key = "cost_" + suffix if suffix else "cost"
-
-    def __call__(self, input_tokens: int, output_tokens: int, cost: float):
+        self.energy_kwh_key = "energy_kwh_" + suffix if suffix else "energy_kwh"
+        self.gwp_kgco2eq_key = "gwp_kgco2eq_" + suffix if suffix else "gwp_kgco2eq"
+
+    def __call__(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cost: float,
+        energy_kwh: float = 0.0,
+        gwp_kgco2eq: float = 0.0,
+    ):
         self.input_tokens += input_tokens
         self.output_tokens += output_tokens
         self.cost += cost
+        self.energy_kwh += energy_kwh
+        self.gwp_kgco2eq += gwp_kgco2eq
 
     @property
     def stats(self):
         return {
             self.input_tokens_key: self.input_tokens,
             self.output_tokens_key: self.output_tokens,
             self.cost_key: self.cost,
+            self.energy_kwh_key: self.energy_kwh,
+            self.gwp_kgco2eq_key: self.gwp_kgco2eq,
         }
 
     def add_tracker(self, tracker: "LLMTracker"):
-        self(tracker.input_tokens, tracker.output_tokens, tracker.cost)
+        self(
+            tracker.input_tokens,
+            tracker.output_tokens,
+            tracker.cost,
+            tracker.energy_kwh,
+            tracker.gwp_kgco2eq,
+        )
 
     def __repr__(self):
-        return f"LLMTracker(input_tokens={self.input_tokens}, output_tokens={self.output_tokens}, cost={self.cost})"
+        return (
+            f"LLMTracker(input_tokens={self.input_tokens}, "
+            f"output_tokens={self.output_tokens}, cost={self.cost}, "
+            f"energy_kwh={self.energy_kwh}, gwp_kgco2eq={self.gwp_kgco2eq})"
+        )
 
 
 @contextmanager
@@ -196,6 +252,7 @@ def init_pricing_tracker(self, pricing_api=None):
 
     def __call__(self, *args, **kwargs):
         """Call the API and update the pricing tracker."""
+        _ensure_ecologits_init()
         # 'self' here calls ._call_api() method of the subclass
         response = self._call_api(*args, **kwargs)
         usage = dict(getattr(response, "usage", {}))
@@ -206,8 +263,10 @@ def __call__(self, *args, **kwargs):
         usage = {f"usage_{k}": v for k, v in usage.items() if isinstance(v, (int, float))}
         usage |= {"n_api_calls": 1}
         usage |= {"effective_cost": self.get_effective_cost(response)}
+        energy_kwh, gwp_kgco2eq = _extract_ecologits_impacts(response)
+        usage |= {"energy_kwh": energy_kwh, "gwp_kgco2eq": gwp_kgco2eq}
         self.stats.increment_stats_dict(usage)
-        self.update_pricing_tracker(response)
+        self.update_pricing_tracker(response, energy_kwh, gwp_kgco2eq)
         return self._parse_response(response)
 
     def fetch_pricing_information_from_provider(self) -> Optional[dict]:
@@ -249,14 +308,16 @@ def set_pricing_attributes(self) -> None:
                 self.input_cost = 0.0
                 self.output_cost = 0.0
 
-    def update_pricing_tracker(self, raw_response) -> None:
+    def update_pricing_tracker(
+        self, raw_response, energy_kwh: float = 0.0, gwp_kgco2eq: float = 0.0
+    ) -> None:
         """Update the pricing tracker with the input and output tokens and cost."""
 
         input_tokens, output_tokens = self.get_tokens_counts_from_response(raw_response)
         cost = input_tokens * self.input_cost + output_tokens * self.output_cost
 
         if hasattr(TRACKER, "instance") and isinstance(TRACKER.instance, LLMTracker):
-            TRACKER.instance(input_tokens, output_tokens, cost)
+            TRACKER.instance(input_tokens, output_tokens, cost, energy_kwh, gwp_kgco2eq)
 
     def get_tokens_counts_from_response(self, response) -> tuple:
         """Get the input and output tokens counts from the response, provider-agnostic."""

tests/llm/test_tracking.py

@@ -21,6 +21,8 @@ def test_get_action_decorator():
         "input_tokens": 1,
         "output_tokens": 1,
         "cost": 1.0,
+        "energy_kwh": 0.0,
+        "gwp_kgco2eq": 0.0,
     }
 
 
@@ -175,3 +177,71 @@ def test_openrouter_chat_model():
         answer = chat_model(messages)
     assert "5" in answer.get("content")
     assert tracker.stats["cost"] > 0
+
+
+# -- ecologits integration tests ------------------------------------------------
+
+
+def call_llm_with_ecologits():
+    """Simulate an LLM call that reports ecologits impacts."""
+    if hasattr(tracking.TRACKER, "instance") and isinstance(
+        tracking.TRACKER.instance, tracking.LLMTracker
+    ):
+        tracking.TRACKER.instance(10, 5, 0.001, energy_kwh=0.0003, gwp_kgco2eq=0.00012)
+    return "action", {"stats": {}}
+
+
+def test_tracker_ecologits_fields():
+    """LLMTracker accumulates energy and carbon when provided."""
+    with tracking.set_tracker() as tracker:
+        call_llm_with_ecologits()
+        call_llm_with_ecologits()
+
+    assert tracker.stats["energy_kwh"] == pytest.approx(0.0006)
+    assert tracker.stats["gwp_kgco2eq"] == pytest.approx(0.00024)
+    assert tracker.stats["input_tokens"] == 20
+    assert tracker.stats["output_tokens"] == 10
+    assert tracker.stats["cost"] == pytest.approx(0.002)
+
+
+def test_tracker_ecologits_defaults():
+    """Energy and carbon default to 0.0 when not provided."""
+    with tracking.set_tracker() as tracker:
+        call_llm()
+
+    assert tracker.stats["energy_kwh"] == 0.0
+    assert tracker.stats["gwp_kgco2eq"] == 0.0
+    assert tracker.stats["cost"] == 1.0
+
+
+def test_tracker_ecologits_nested():
+    """Nested trackers accumulate ecologits impacts correctly."""
+    with tracking.set_tracker() as outer:
+        call_llm_with_ecologits()
+        with tracking.set_tracker() as inner:
+            call_llm_with_ecologits()
+
+    assert inner.stats["energy_kwh"] == pytest.approx(0.0003)
+    assert inner.stats["gwp_kgco2eq"] == pytest.approx(0.00012)
+    assert outer.stats["energy_kwh"] == pytest.approx(0.0006)
+    assert outer.stats["gwp_kgco2eq"] == pytest.approx(0.00024)
+
+
+def test_extract_ecologits_impacts_no_impacts():
+    """_extract_ecologits_impacts returns zeros when response has no impacts."""
+
+    class MockResponse:
+        pass
+
+    energy, gwp = tracking._extract_ecologits_impacts(MockResponse())
+    assert energy == 0.0
+    assert gwp == 0.0
+
+
+def test_tracker_repr_includes_ecologits():
+    """LLMTracker repr includes energy and carbon fields."""
+    tracker = tracking.LLMTracker()
+    tracker(10, 5, 0.5, energy_kwh=0.001, gwp_kgco2eq=0.0004)
+    repr_str = repr(tracker)
+    assert "energy_kwh=0.001" in repr_str
+    assert "gwp_kgco2eq=0.0004" in repr_str