diff --git a/pyproject.toml b/pyproject.toml index 98fde2d77..a8dee81f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,10 @@ hint = [ "sentence-transformers>=5.0.0", ] transformers = [ - "transformers>=4.38.2", + "transformers>=4.38.2", +] +ecologits = [ + "ecologits", ] diff --git a/src/agentlab/llm/chat_api.py b/src/agentlab/llm/chat_api.py index 4c642220a..72967dee8 100644 --- a/src/agentlab/llm/chat_api.py +++ b/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 diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index 7b4b64f65..ecda3a463 100644 --- a/src/agentlab/llm/tracking.py +++ b/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,14 +68,27 @@ 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): @@ -52,13 +96,25 @@ def stats(self): 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.""" diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py index 0809c54a6..6fa514b63 100644 --- a/tests/llm/test_tracking.py +++ b/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