diff --git a/openevolve/llm/ensemble.py b/openevolve/llm/ensemble.py index 1d912971f..be398ecdd 100644 --- a/openevolve/llm/ensemble.py +++ b/openevolve/llm/ensemble.py @@ -115,7 +115,7 @@ def _sample_model(self) -> LLMInterface: """Sample a model from the ensemble based on weights""" index = self.random_state.choices(range(len(self.models)), weights=self.weights, k=1)[0] sampled_model = self.models[index] - logger.info(f"Sampled model: {vars(sampled_model)['model']}") + logger.info(f"Sampled model: {self.models_cfg[index].name}") return sampled_model async def generate_multiple(self, prompt: str, n: int, **kwargs) -> List[str]: diff --git a/tests/test_llm_ensemble.py b/tests/test_llm_ensemble.py index 45f798b00..189a11713 100644 --- a/tests/test_llm_ensemble.py +++ b/tests/test_llm_ensemble.py @@ -8,6 +8,7 @@ from openevolve.config import LLMModelConfig from openevolve.llm.base import LLMInterface + class TestLLMEnsemble(unittest.TestCase): def test_weighted_sampling(self): models = [ @@ -35,7 +36,6 @@ def test_weighted_sampling(self): self.assertEqual(len(sampled_models), len(models)) - class TestEnsembleInit(unittest.TestCase): class MyCustomLLM(LLMInterface): def __init__(self, model, some_field): @@ -45,7 +45,9 @@ def __init__(self, model, some_field): async def generate(self, prompt: str, **kwargs) -> str: return "custom-generate" - async def generate_with_context(self, system_message: str, messages: List[Dict[str, str]], **kwargs) -> str: + async def generate_with_context( + self, system_message: str, messages: List[Dict[str, str]], **kwargs + ) -> str: return "custom-generate-with-context" def init_custom_llm(self, model_cfg): @@ -116,6 +118,47 @@ async def test_openai_llm_token_usage_logging_and_propagation(self): self.assertEqual(ensemble.last_usage, llm.last_usage) +class TestCustomClientSampling(unittest.IsolatedAsyncioTestCase): + class InterfaceOnlyLLM(LLMInterface): + def __init__(self): + self.calls = [] + + async def generate(self, prompt, **kwargs): + self.calls.append((prompt, kwargs)) + return "custom answer" + + async def generate_with_context(self, system_message, messages, **kwargs): + self.calls.append((system_message, messages, kwargs)) + return "custom context answer" + + class PropertyModelLLM(InterfaceOnlyLLM): + @property + def model(self): + return "custom model" + + async def check_custom_client(self, client): + ensemble = LLMEnsemble( + [LLMModelConfig(name="custom model", init_client=lambda _cfg: client)] + ) + messages = [{"role": "user", "content": "question"}] + with self.assertLogs("openevolve.llm.ensemble", level="INFO") as logs: + self.assertEqual(await ensemble.generate("question", temperature=0), "custom answer") + self.assertEqual( + await ensemble.generate_with_context("system", messages, max_tokens=8), + "custom context answer", + ) + self.assertEqual( + client.calls, + [("question", {"temperature": 0}), ("system", messages, {"max_tokens": 8})], + ) + self.assertTrue(all("custom model" in message for message in logs.output)) + + async def test_interface_only_client_can_generate(self): + await self.check_custom_client(self.InterfaceOnlyLLM()) + + async def test_property_backed_model_client_can_generate(self): + await self.check_custom_client(self.PropertyModelLLM()) + + if __name__ == "__main__": unittest.main() -