feat(llm): track, propagate, and log LLM token usage across evolution…

openevolve/database.py

@@ -2587,8 +2587,9 @@ def log_prompt(
         self,
         program_id: str,
         template_key: str,
-        prompt: Dict[str, str],
+        prompt: Dict[str, Any],
         responses: Optional[List[str]] = None,
+        token_usage: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
         Log a prompt for a program.
@@ -2599,6 +2600,7 @@ def log_prompt(
         template_key: Key for the prompt template
         prompt: Prompts in the format {template_key: { 'system': str, 'user': str }}.
         responses: Optional list of responses to the prompt, if available.
+        token_usage: Optional token usage dictionary for this LLM call.
         """
 
         if not self.config.log_prompts:
@@ -2607,6 +2609,8 @@ def log_prompt(
         if responses is None:
             responses = []
         prompt["responses"] = responses
+        if token_usage is not None:
+            prompt["token_usage"] = token_usage
 
         if self.prompts_by_program is None:
             self.prompts_by_program = {}

openevolve/llm/ensemble.py

@@ -71,17 +71,23 @@ def __init__(self, models_cfg: List[LLMModelConfig]):
             )
             logger._ensemble_logged = True
 
+        self.last_usage: Optional[Dict[str, Any]] = None
+
     async def generate(self, prompt: str, **kwargs) -> str:
         """Generate text using a randomly selected model based on weights"""
         model = self._sample_model()
-        return await model.generate(prompt, **kwargs)
+        res = await model.generate(prompt, **kwargs)
+        self.last_usage = getattr(model, "last_usage", None)
+        return res
 
     async def generate_with_context(
         self, system_message: str, messages: List[Dict[str, str]], **kwargs
     ) -> str:
         """Generate text using a system message and conversational context"""
         model = self._sample_model()
-        return await model.generate_with_context(system_message, messages, **kwargs)
+        res = await model.generate_with_context(system_message, messages, **kwargs)
+        self.last_usage = getattr(model, "last_usage", None)
+        return res
 
     def _sample_model(self) -> LLMInterface:
         """Sample a model from the ensemble based on weights"""

openevolve/llm/openai.py

@@ -63,6 +63,7 @@ def __init__(
         self.api_key = model_cfg.api_key
         self.random_seed = getattr(model_cfg, "random_seed", None)
         self.reasoning_effort = getattr(model_cfg, "reasoning_effort", None)
+        self.last_usage: Optional[Dict[str, Any]] = None
 
         # Manual mode: enabled via llm.manual_mode in config.yaml
         self.manual_mode = (getattr(model_cfg, "manual_mode", False) is True)
@@ -225,6 +226,28 @@ async def _call_api(self, params: Dict[str, Any]) -> str:
         logger = logging.getLogger(__name__)
         logger.debug(f"API parameters: {params}")
         logger.debug(f"API response: {response.choices[0].message.content}")
+
+        usage = getattr(response, "usage", None)
+        if usage is not None:
+            prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
+            completion_tokens = getattr(usage, "completion_tokens", 0) or 0
+            total_tokens = getattr(usage, "total_tokens", 0) or (prompt_tokens + completion_tokens)
+            self.last_usage = {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "model": str(self.model),
+            }
+            logger.info(
+                f"LLM Token Usage ({self.model}): "
+                f"prompt_tokens={prompt_tokens}, "
+                f"completion_tokens={completion_tokens}, "
+                f"total_tokens={total_tokens}"
+            )
+        else:
+            self.last_usage = None
+            logger.info(f"LLM Token Usage ({self.model}): usage info not returned by API")
+
         return response.choices[0].message.content
 
     async def _manual_wait_for_answer(

openevolve/process_parallel.py

@@ -34,6 +34,7 @@ class SerializableResult:
     iteration: int = 0
     error: Optional[str] = None
     target_island: Optional[int] = None  # Island where child should be placed
+    token_usage: Optional[Dict[str, Any]] = None
 
 
 def _worker_init(config_dict: dict, evaluation_file: str, parent_env: dict = None) -> None:
@@ -203,13 +204,16 @@ def _run_iteration_worker(
                     messages=[{"role": "user", "content": prompt["user"]}],
                 )
             )
+            token_usage = getattr(_worker_llm_ensemble, "last_usage", None)
         except Exception as e:
             logger.error(f"LLM generation failed: {e}")
             return SerializableResult(error=f"LLM generation failed: {str(e)}", iteration=iteration)
 
         # Check for None response
         if llm_response is None:
-            return SerializableResult(error="LLM returned None response", iteration=iteration)
+            return SerializableResult(
+                error="LLM returned None response", iteration=iteration, token_usage=token_usage
+            )
 
         # Parse response based on evolution mode
         if _worker_config.diff_based_evolution:
@@ -224,7 +228,9 @@ def _run_iteration_worker(
             diff_blocks = extract_diffs(llm_response, _worker_config.diff_pattern)
             if not diff_blocks:
                 return SerializableResult(
-                    error="No valid diffs found in response", iteration=iteration
+                    error="No valid diffs found in response",
+                    iteration=iteration,
+                    token_usage=token_usage,
                 )
 
             if _worker_config.prompt.programs_as_changes_description:
@@ -235,7 +241,9 @@ def _run_iteration_worker(
                         changes_description_text=parent_changes_desc,
                     )
                 except Exception as e:
-                    return SerializableResult(error=str(e), iteration=iteration)
+                    return SerializableResult(
+                        error=str(e), iteration=iteration, token_usage=token_usage
+                    )
 
                 child_code, _ = apply_diff_blocks(parent.code, code_blocks)
                 child_changes_desc, desc_applied = apply_diff_blocks(
@@ -251,6 +259,7 @@ def _run_iteration_worker(
                     return SerializableResult(
                         error="changes_description was not updated or empty, program is discarded",
                         iteration=iteration,
+                        token_usage=token_usage,
                     )
 
                 changes_summary = format_diff_summary(
@@ -272,7 +281,9 @@ def _run_iteration_worker(
             new_code = parse_full_rewrite(llm_response, _worker_config.language)
             if not new_code:
                 return SerializableResult(
-                    error=f"No valid code found in response", iteration=iteration
+                    error=f"No valid code found in response",
+                    iteration=iteration,
+                    token_usage=token_usage,
                 )
 
             child_code = new_code
@@ -283,6 +294,7 @@ def _run_iteration_worker(
             return SerializableResult(
                 error=f"Generated code exceeds maximum length ({len(child_code)} > {_worker_config.max_code_length})",
                 iteration=iteration,
+                token_usage=token_usage,
             )
 
         # Evaluate the child program
@@ -308,6 +320,7 @@ def _run_iteration_worker(
                 "changes": changes_summary,
                 "parent_metrics": parent.metrics,
                 "island": parent_island,
+                "token_usage": token_usage,
             },
         )
 
@@ -325,6 +338,7 @@ def _run_iteration_worker(
             artifacts=artifacts,
             iteration=iteration,
             target_island=target_island,
+            token_usage=token_usage,
         )
 
     except Exception as e:
@@ -564,6 +578,10 @@ async def run_evolution(
 
         next_iteration = current_iteration
         completed_iterations = 0
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+        total_tokens = 0
+        total_llm_calls = 0
 
         # Early stopping tracking
         early_stopping_enabled = self.config.early_stopping_patience is not None
@@ -609,8 +627,21 @@ async def run_evolution(
                 timeout_seconds = self.config.evaluator.timeout + 30
                 result = future.result(timeout=timeout_seconds)
 
+                token_str = ""
+                if result.token_usage:
+                    pt = result.token_usage.get("prompt_tokens", 0)
+                    ct = result.token_usage.get("completion_tokens", 0)
+                    tt = result.token_usage.get("total_tokens", 0)
+                    total_prompt_tokens += pt
+                    total_completion_tokens += ct
+                    total_tokens += tt
+                    total_llm_calls += 1
+                    token_str = f" | tokens: {tt} (prompt: {pt}, completion: {ct})"
+
                 if result.error:
-                    logger.warning(f"Iteration {completed_iteration} error: {result.error}")
+                    logger.warning(
+                        f"Iteration {completed_iteration} error: {result.error}{token_str}"
+                    )
                 elif result.child_program_dict:
                     # Reconstruct program from dict
                     child_program = Program(**result.child_program_dict)
@@ -651,6 +682,7 @@ async def run_evolution(
                                 metadata={
                                     "iteration_time": result.iteration_time,
                                     "changes": child_program.metadata.get("changes", ""),
+                                    "token_usage": result.token_usage,
                                 },
                             )
 
@@ -665,6 +697,7 @@ async def run_evolution(
                             program_id=child_program.id,
                             prompt=result.prompt,
                             responses=[result.llm_response] if result.llm_response else [],
+                            token_usage=result.token_usage,
                         )
 
                     # Island management
@@ -684,7 +717,7 @@ async def run_evolution(
                         f"Iteration {completed_iteration}: "
                         f"Program {child_program.id} "
                         f"(parent: {result.parent_id}) "
-                        f"completed in {result.iteration_time:.2f}s"
+                        f"completed in {result.iteration_time:.2f}s{token_str}"
                     )
 
                     if child_program.metrics:
@@ -847,6 +880,13 @@ async def run_evolution(
         else:
             logger.info("✅ Evolution completed - Maximum iterations reached")
 
+        if total_llm_calls > 0:
+            logger.info(
+                f"📊 Total LLM Token Usage ({total_llm_calls} calls): "
+                f"total_tokens={total_tokens} "
+                f"(prompt_tokens={total_prompt_tokens}, completion_tokens={total_completion_tokens})"
+            )
+
         return self.database.get_best_program()
 
     def _submit_iteration(

tests/test_llm_ensemble.py

@@ -62,5 +62,60 @@ def test_ensemble_initialization(self):
         self.assertEqual(ensemble.models[1].model, "b")
         self.assertEqual(ensemble.models[1].some_field, "value")
 
+
+class TestTokenUsageLogging(unittest.IsolatedAsyncioTestCase):
+    async def test_openai_llm_token_usage_logging_and_propagation(self):
+        from unittest.mock import MagicMock
+        from openevolve.llm.openai import OpenAILLM
+
+        cfg = LLMModelConfig(
+            name="gpt-4o-mini",
+            api_key="test",
+            api_base="http://test",
+            retries=0,
+            retry_delay=0.1,
+            timeout=10,
+        )
+        llm = OpenAILLM(cfg)
+
+        mock_response = MagicMock()
+        mock_response.choices[0].message.content = "evolved code response"
+        mock_response.usage.prompt_tokens = 1234
+        mock_response.usage.completion_tokens = 567
+        mock_response.usage.total_tokens = 1801
+
+        llm.client.chat.completions.create = MagicMock(return_value=mock_response)
+
+        with self.assertLogs("openevolve.llm.openai", level="INFO") as log_ctx:
+            result = await llm.generate_with_context(
+                system_message="sys",
+                messages=[{"role": "user", "content": "user"}],
+            )
+
+        self.assertEqual(result, "evolved code response")
+        self.assertIsNotNone(llm.last_usage)
+        self.assertEqual(llm.last_usage["prompt_tokens"], 1234)
+        self.assertEqual(llm.last_usage["completion_tokens"], 567)
+        self.assertEqual(llm.last_usage["total_tokens"], 1801)
+        self.assertEqual(llm.last_usage["model"], "gpt-4o-mini")
+        self.assertTrue(
+            any(
+                "LLM Token Usage (gpt-4o-mini): prompt_tokens=1234, completion_tokens=567, total_tokens=1801"
+                in msg
+                for msg in log_ctx.output
+            )
+        )
+
+        # Also verify LLMEnsemble propagates last_usage
+        ensemble = LLMEnsemble([cfg])
+        ensemble.models = [llm]
+        await ensemble.generate_with_context(
+            system_message="sys",
+            messages=[{"role": "user", "content": "user"}],
+        )
+        self.assertEqual(ensemble.last_usage, llm.last_usage)
+
+
 if __name__ == "__main__":
     unittest.main()
+