diff --git a/openevolve/database.py b/openevolve/database.py index 8abe2bdc0..822852e66 100644 --- a/openevolve/database.py +++ b/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 = {} diff --git a/openevolve/llm/ensemble.py b/openevolve/llm/ensemble.py index b9161382a..c92928d6d 100644 --- a/openevolve/llm/ensemble.py +++ b/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""" diff --git a/openevolve/llm/openai.py b/openevolve/llm/openai.py index f7d0648de..75e7834ff 100644 --- a/openevolve/llm/openai.py +++ b/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( diff --git a/openevolve/process_parallel.py b/openevolve/process_parallel.py index b2cfeab78..20bc3fff7 100644 --- a/openevolve/process_parallel.py +++ b/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( diff --git a/tests/test_llm_ensemble.py b/tests/test_llm_ensemble.py index 7c11baea0..f6dfcaecb 100644 --- a/tests/test_llm_ensemble.py +++ b/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() +