fix(llm): avoid assuming custom clients store a model field

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]:

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()
-