diff --git a/python/packages/autogen-agentchat/docs/optimization.md b/python/packages/autogen-agentchat/docs/optimization.md new file mode 100644 index 000000000000..beffabbb1faa --- /dev/null +++ b/python/packages/autogen-agentchat/docs/optimization.md @@ -0,0 +1,92 @@ +# AutoGen Agent Optimizer + +The AutoGen Agent Optimizer provides a unified interface for optimizing AutoGen agents using various optimization backends. This allows you to improve agent performance by automatically tuning system messages and tool descriptions based on training data. + +## Installation + +The base optimization interface is included with `autogen-agentchat`. To use the DSPy backend, you'll also need to install DSPy: + +```bash +pip install autogen-ext[dspy] +# or directly +pip install dspy +``` + +## Basic Usage + +```python +from autogen_agentchat.optimize import compile, list_backends + +# Check available backends +print("Available backends:", list_backends()) + +# Optimize an agent +optimized_agent, report = compile( + agent=my_agent, + trainset=training_examples, + metric=evaluation_function, + backend="dspy", + optimizer_name="MIPROv2", + optimizer_kwargs={"max_steps": 16} +) +``` + +## Interface + +### `compile(agent, trainset, metric, *, backend="dspy", **kwargs)` + +Optimizes an AutoGen agent by tuning its system message and tool descriptions. + +**Parameters:** +- `agent`: Any AutoGen agent (e.g., AssistantAgent) +- `trainset`: Iterable of training examples (DSPy Examples or backend-specific format) +- `metric`: Evaluation function `(gold, pred) → float | bool` +- `backend`: Name of optimization backend (default: "dspy") +- `**kwargs`: Additional parameters passed to the backend + +**Returns:** +- `(optimized_agent, report)`: Tuple of the optimized agent and optimization report + +### `list_backends()` + +Returns a list of available optimization backends. + +## Backends + +### DSPy Backend + +The DSPy backend uses the DSPy optimization framework to improve agent prompts. + +**Supported optimizers:** +- SIMBA (default) +- MIPROv2 +- And any other DSPy optimizer + +**Backend-specific parameters:** +- `lm_client`: Language model client (defaults to agent's model client) +- `optimizer_name`: Name of DSPy optimizer (default: "SIMBA") +- `optimizer_kwargs`: Additional optimizer parameters + +## Example + +See `examples/optimization_demo.py` for a complete example demonstrating the interface. + +## Adding New Backends + +To add a new optimization backend: + +1. Create a class inheriting from `BaseBackend` +2. Set the `name` class attribute +3. Implement the `compile()` method +4. The backend will be automatically registered when imported + +```python +from autogen_agentchat.optimize._backend import BaseBackend + +class MyBackend(BaseBackend): + name = "my_backend" + + def compile(self, agent, trainset, metric, **kwargs): + # Your optimization logic here + return optimized_agent, report +``` \ No newline at end of file diff --git a/python/packages/autogen-agentchat/examples/optimization_demo.py b/python/packages/autogen-agentchat/examples/optimization_demo.py new file mode 100644 index 000000000000..b90eab590d83 --- /dev/null +++ b/python/packages/autogen-agentchat/examples/optimization_demo.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +""" +Example demonstrating the AutoGen Agent Optimizer interface. + +This example shows how to use the optimization interface, including: +1. Creating an agent with tools +2. Preparing a training dataset +3. Using the compile() function to optimize the agent +4. Checking available backends + +Note: This example requires DSPy to be installed for actual optimization. +Run: pip install dspy +""" + +import asyncio +from unittest.mock import Mock + +def main(): + """Demonstrate the AutoGen Agent Optimizer interface.""" + print("=== AutoGen Agent Optimizer Demo ===\n") + + # Import the optimization interface + from autogen_agentchat.optimize import compile, list_backends + + # Import DSPy backend to register it (safe import) + try: + from autogen_ext.optimize import dspy # This will register the backend + except ImportError: + pass # DSPy backend not available + + print("1. Available optimization backends:") + backends = list_backends() + print(f" {backends}") + print() + + # ➊ Build a toy agent --------------------------------------------------- + print("2. Creating a simple agent with tools...") + + class SimpleAgent: + """Mock agent for demonstration.""" + def __init__(self, name: str, system_message: str): + self.name = name + self.system_message = system_message + self._system_messages = [] + if system_message: + # Mock SystemMessage class + class SystemMessage: + def __init__(self, content): + self.content = content + self._system_messages = [SystemMessage(system_message)] + + # Mock tools + self._tools = [] + self.model_client = Mock() # Mock model client + + def add_tool(self, name: str, description: str): + """Add a mock tool to the agent.""" + class MockTool: + def __init__(self, name, description): + self.name = name + self.description = description + + self._tools.append(MockTool(name, description)) + + # Create the agent + agent = SimpleAgent( + name="calc", + system_message="You are a helpful calculator assistant." + ) + + # Add a tool + agent.add_tool("add", "Add two numbers together") + + print(f" Agent: {agent.name}") + print(f" System message: {agent.system_message}") + print(f" Tools: {[(t.name, t.description) for t in agent._tools]}") + print() + + # ➋ Minimal trainset ---------------------------------------------------- + print("3. Creating training dataset...") + + # Mock DSPy Example format + class MockExample: + def __init__(self, user_request: str, answer: str): + self.user_request = user_request + self.answer = answer + + def with_inputs(self, *inputs): + return self + + train = [ + MockExample(user_request="2+2", answer="4").with_inputs("user_request"), + MockExample(user_request="Add 3 and 5", answer="8").with_inputs("user_request"), + ] + + print(f" Training examples: {len(train)}") + for i, ex in enumerate(train): + print(f" Example {i+1}: '{ex.user_request}' -> '{ex.answer}'") + print() + + # ➌ Define metric -------------------------------------------------------- + print("4. Defining evaluation metric...") + + def metric(gold, pred, **kwargs): + """Simple exact match metric.""" + return getattr(gold, 'answer', gold) == getattr(pred, 'answer', pred) + + print(" Using exact match metric") + print() + + # ➍ Optimize using the unified API -------------------------------------- + print("5. Attempting optimization...") + + try: + # This is the main interface as specified in the issue + opt_agent, report = compile( + agent=agent, + trainset=train, + metric=metric, + backend="dspy", # default anyway + optimizer_name="MIPROv2", + optimizer_kwargs=dict(max_steps=8), + ) + + print("✓ Optimization completed successfully!") + print("\nOptimization Report:") + for key, value in report.items(): + print(f" {key}: {value}") + + print(f"\nOptimized agent system message:") + print(f" {opt_agent.system_message}") + + except ImportError as e: + print(f"⚠ DSPy not available: {e}") + print("\nTo run actual optimization, install DSPy:") + print(" pip install dspy") + print("\nThe interface is ready to use once DSPy is installed!") + + except Exception as e: + print(f"❌ Optimization failed: {e}") + + print("\n=== Demo Complete ===") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/__init__.py index c5bdfc2b51cb..6dfedeabc394 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/__init__.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/__init__.py @@ -11,4 +11,7 @@ EVENT_LOGGER_NAME = "autogen_agentchat.events" """Logger name for event logs.""" -__version__ = importlib.metadata.version("autogen_agentchat") +try: + __version__ = importlib.metadata.version("autogen_agentchat") +except importlib.metadata.PackageNotFoundError: + __version__ = "0.6.1-dev" diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/__init__.py new file mode 100644 index 000000000000..4f9ad03eae17 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/__init__.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import Any, Callable, Dict, Iterable, List, Tuple + +from ._backend import _BACKENDS, get_backend + + +def compile( + agent: Any, + trainset: Iterable[Any], + metric: Callable[[Any, Any], float | bool], + *, + backend: str = "dspy", + **kwargs: Any, +) -> Tuple[Any, Dict[str, Any]]: + """ + Optimise the `system_message` and tool descriptions of an AutoGen agent. + + Parameters + ---------- + agent + Any subclass of autogen_core.agents.base.Agent. + trainset + Iterable of supervision examples (DSPy Examples or anything the + back-end accepts). + metric + Callable(gold, pred) → float | bool used by the optimiser. + backend + Name of the registered optimisation backend (default: "dspy"). + kwargs + Extra parameters forwarded verbatim to the backend. + + Returns + ------- + (optimised_agent, report) + """ + backend_impl = get_backend(backend) + return backend_impl.compile(agent, trainset, metric, **kwargs) + + +def list_backends() -> List[str]: + """Return the names of all available optimisation back-ends.""" + return sorted(_BACKENDS.keys()) \ No newline at end of file diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/_backend.py b/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/_backend.py new file mode 100644 index 000000000000..2c98e72b8b88 --- /dev/null +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/optimize/_backend.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, Iterable, Tuple + +# Simple registry so new back-ends can self-register +_BACKENDS: Dict[str, type["BaseBackend"]] = {} + + +class BaseBackend(ABC): + """Contract every optimiser back-end must fulfil.""" + + #: name used in compile(... backend="") + name: str = "" + + def __init_subclass__(cls, **kw: Any) -> None: + super().__init_subclass__(**kw) + if cls.name: + _BACKENDS[cls.name] = cls + + # ---- required API -------------------------------------------------- + @abstractmethod + def compile( + self, + agent: Any, + trainset: Iterable[Any], + metric: Callable[[Any, Any], float | bool], + **kwargs: Any, + ) -> Tuple[Any, Dict[str, Any]]: + """Return (optimised_agent, diagnostics/report).""" + ... + + +def get_backend(name: str) -> BaseBackend: + """Get a backend instance by name.""" + try: + backend_cls = _BACKENDS[name] + return backend_cls() + except KeyError: + raise ValueError( + f"Unknown backend '{name}'. Available: {', '.join(_BACKENDS)}" + ) from None \ No newline at end of file diff --git a/python/packages/autogen-agentchat/tests/test_optimize.py b/python/packages/autogen-agentchat/tests/test_optimize.py new file mode 100644 index 000000000000..67349c3c3bad --- /dev/null +++ b/python/packages/autogen-agentchat/tests/test_optimize.py @@ -0,0 +1,149 @@ +import pytest +from unittest.mock import Mock +from autogen_agentchat.optimize import list_backends, compile +from autogen_agentchat.optimize._backend import BaseBackend, get_backend + + +def test_backend_registry(): + """Test that the backend registry works.""" + # Get initial backends (might include dspy if loaded) + backends_before = set(list_backends()) + + # Create a dummy backend for testing + class DummyBackend(BaseBackend): + name = "dummy" + + def compile(self, agent, trainset, metric, **kwargs): + return agent, {"optimizer": "dummy", "status": "test"} + + # Should now have the dummy backend + backends_after = set(list_backends()) + assert "dummy" in backends_after + assert backends_after == backends_before.union({"dummy"}) + + +def test_backend_not_found(): + """Test error handling for unknown backend.""" + with pytest.raises(ValueError, match="Unknown backend 'nonexistent'"): + get_backend("nonexistent") + + +def test_compile_with_dummy_backend(): + """Test compile function with dummy backend.""" + # Create a dummy backend + class TestBackend(BaseBackend): + name = "test" + + def compile(self, agent, trainset, metric, **kwargs): + return agent, {"optimizer": "test", "result": "success"} + + # Create a dummy agent + class DummyAgent: + def __init__(self): + self.system_message = "You are a helpful assistant." + + agent = DummyAgent() + trainset = [] + metric = lambda x, y: True + + optimized_agent, report = compile(agent, trainset, metric, backend="test") + + assert optimized_agent is agent + assert report["optimizer"] == "test" + assert report["result"] == "success" + + +def test_dspy_backend_registration(): + """Test that DSPy backend is properly registered when module is imported.""" + # Import the DSPy backend to register it + from autogen_ext.optimize.dspy import DSPyBackend + + backends = list_backends() + assert "dspy" in backends + + +def test_dspy_backend_unavailable(): + """Test that DSPy backend gracefully handles missing DSPy dependency.""" + from autogen_ext.optimize.dspy import DSPyBackend + + class DummyAgent: + def __init__(self): + self.system_message = "You are helpful" + self._model_client = Mock() + + agent = DummyAgent() + trainset = [] + metric = lambda x, y: True + + backend = DSPyBackend() + + # Should raise ImportError about missing DSPy + with pytest.raises(ImportError, match="DSPy is required for optimization"): + backend.compile(agent, trainset, metric) + + +def test_dspy_backend_missing_model_client(): + """Test DSPy backend error handling for missing model client.""" + from autogen_ext.optimize.dspy import DSPyBackend + + class DummyAgent: + def __init__(self): + self.system_message = "You are helpful" + # No model_client attribute + + agent = DummyAgent() + trainset = [] + metric = lambda x, y: True + + backend = DSPyBackend() + + # Should raise ValueError about missing model client + with pytest.raises(ValueError, match="Could not find model_client"): + backend.compile(agent, trainset, metric) + + +def test_compile_function_with_kwargs(): + """Test that compile function forwards kwargs to backend.""" + class KwargsTestBackend(BaseBackend): + name = "kwargs_test" + + def compile(self, agent, trainset, metric, **kwargs): + return agent, {"received_kwargs": kwargs} + + class DummyAgent: + pass + + agent = DummyAgent() + trainset = [] + metric = lambda x, y: True + + # Pass some kwargs + test_kwargs = {"optimizer_name": "MIPROv2", "max_steps": 10} + optimized_agent, report = compile( + agent, trainset, metric, + backend="kwargs_test", + **test_kwargs + ) + + assert report["received_kwargs"] == test_kwargs + + +def test_list_backends_returns_sorted(): + """Test that list_backends returns a sorted list.""" + # Create multiple backends + class BackendZ(BaseBackend): + name = "z_backend" + def compile(self, agent, trainset, metric, **kwargs): + return agent, {} + + class BackendA(BaseBackend): + name = "a_backend" + def compile(self, agent, trainset, metric, **kwargs): + return agent, {} + + backends = list_backends() + + # Should be sorted alphabetically + assert backends == sorted(backends) + assert "a_backend" in backends + assert "z_backend" in backends \ No newline at end of file diff --git a/python/packages/autogen-ext/pyproject.toml b/python/packages/autogen-ext/pyproject.toml index 6e529921b59f..30a2aac59805 100644 --- a/python/packages/autogen-ext/pyproject.toml +++ b/python/packages/autogen-ext/pyproject.toml @@ -155,6 +155,8 @@ canvas = [ "unidiff>=0.7.5", ] +dspy = ["dspy>=2.4.0"] + [tool.hatch.build.targets.wheel] packages = ["src/autogen_ext"] diff --git a/python/packages/autogen-ext/src/autogen_ext/__init__.py b/python/packages/autogen-ext/src/autogen_ext/__init__.py index bd2c9ca453aa..da7d8e5cd4ee 100644 --- a/python/packages/autogen-ext/src/autogen_ext/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/__init__.py @@ -1,3 +1,6 @@ import importlib.metadata -__version__ = importlib.metadata.version("autogen_ext") +try: + __version__ = importlib.metadata.version("autogen_ext") +except importlib.metadata.PackageNotFoundError: + __version__ = "0.6.1-dev" diff --git a/python/packages/autogen-ext/src/autogen_ext/optimize/__init__.py b/python/packages/autogen-ext/src/autogen_ext/optimize/__init__.py new file mode 100644 index 000000000000..3da9a75aea64 --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/optimize/__init__.py @@ -0,0 +1,8 @@ +# Optimization backends for AutoGen agents + +# Try to import DSPy backend to register it +try: + from .dspy import DSPyBackend # noqa: F401 +except ImportError: + # DSPy not available, skip registration + pass \ No newline at end of file diff --git a/python/packages/autogen-ext/src/autogen_ext/optimize/_utils.py b/python/packages/autogen-ext/src/autogen_ext/optimize/_utils.py new file mode 100644 index 000000000000..8cde026d89cf --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/optimize/_utils.py @@ -0,0 +1,163 @@ +"""Utility glue for turning AutoGen agents into DSPy-friendly modules.""" +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, TYPE_CHECKING + +if TYPE_CHECKING: + try: + import dspy + except ImportError: + dspy = None + + from autogen_core.models import AssistantMessage, SystemMessage, UserMessage + from autogen_core.models.base import ChatCompletionClient + + +def _check_dspy_available() -> None: + """Check if DSPy is available and raise helpful error if not.""" + try: + import dspy # noqa: F401 + except ImportError as e: + raise ImportError( + "DSPy is required for optimization but not installed. " + "Please install it with: pip install dspy" + ) from e + + +# --------------------------------------------------------------------- # +# 1. LM adaptor –– AutoGen client ➜ DSPy.LM +# --------------------------------------------------------------------- # +class AutoGenLM: + """Adapts an AutoGen ChatCompletionClient to DSPy.LM interface.""" + + def __init__(self, client: "ChatCompletionClient") -> None: + _check_dspy_available() + import dspy + + self.client = client + self.model = getattr(client, "model", "unknown") + # Initialize basic attributes expected by DSPy + self.history = [] + + async def _acall(self, messages: List[Dict[str, str]], **kw: Any) -> str: + """Convert DSPy messages to AutoGen format and call the client.""" + from autogen_core.models import AssistantMessage, SystemMessage, UserMessage + + autogen_msgs = [] + for m in messages: + role, content = m["role"], m["content"] + if role == "user": + autogen_msgs.append(UserMessage(content=content)) + elif role == "assistant": + autogen_msgs.append(AssistantMessage(content=content)) + else: + autogen_msgs.append(SystemMessage(content=content)) + + resp = await self.client.create(autogen_msgs, **kw) + return resp.content + + def __call__(self, messages: List[Dict[str, str]], **kw: Any) -> str: + """Synchronous interface for DSPy compatibility.""" + return asyncio.run(self._acall(messages, **kw)) + + def basic_request(self, messages: List[Dict[str, str]], **kw: Any) -> str: + """DSPy interface method.""" + return self(messages, **kw) + + +# --------------------------------------------------------------------- # +# 2. DSPy module wrapper around an existing Agent +# --------------------------------------------------------------------- # +class DSPyAgentWrapper: + """ + Exposes `agent.system_message` and each tool description as learnable prompts. + This is a DSPy Module that wraps an AutoGen agent. + """ + + def __init__(self, agent: Any) -> None: + _check_dspy_available() + import dspy + + self._agent = agent + + # Turn system prompt & each tool description into learnable strings + system_message = self._get_system_message(agent) + self._system_prompt = dspy.Prompt(system_message or "You are a helpful assistant.") + + self._tool_prompts = [] + tools = self._get_tools(agent) + for tool in tools: + description = getattr(tool, "description", "") or "" + self._tool_prompts.append(dspy.Prompt(description)) + + # Make this a proper DSPy Module by adding the predict component + class AgentSignature(dspy.Signature): + """Agent signature for processing user requests.""" + user_request: str = dspy.InputField() + answer: str = dspy.OutputField() + + self._predict = dspy.Predict(AgentSignature) + + def _get_system_message(self, agent: Any) -> str | None: + """Extract system message from agent.""" + # Try different ways agents might store system messages + if hasattr(agent, "system_message"): + return agent.system_message + elif hasattr(agent, "_system_messages") and agent._system_messages: + return agent._system_messages[0].content + return None + + def _get_tools(self, agent: Any) -> List[Any]: + """Extract tools from agent.""" + return getattr(agent, "_tools", []) or getattr(agent, "tools", []) + + def forward(self, user_request: str) -> Any: + """Forward pass through the agent. + + In a full implementation, this would: + 1. Update the agent with optimized prompts + 2. Call the agent's run method + 3. Return the result + + For now, we use a simple predict as a placeholder. + """ + _check_dspy_available() + import dspy + + # Patch live values into the wrapped agent + self._update_agent_prompts() + + # In an ideal implementation, we'd call: + # result = await self._agent.run(task=user_request) + # But this requires proper async handling and depends on the agent interface + + # For now, use DSPy predict as a fallback + prediction = self._predict(user_request=user_request) + return dspy.Prediction(answer=prediction.answer) + + def _update_agent_prompts(self) -> None: + """Update the agent with current prompt values.""" + # Update system message + if hasattr(self._agent, "system_message"): + self._agent.system_message = self._system_prompt.value + elif hasattr(self._agent, "_system_messages") and self._agent._system_messages: + from autogen_core.models import SystemMessage + self._agent._system_messages[0] = SystemMessage(content=self._system_prompt.value) + + # Update tool descriptions + tools = self._get_tools(self._agent) + for prompt, tool in zip(self._tool_prompts, tools): + if hasattr(tool, "description"): + tool.description = prompt.value + + # Convenient handles for back-end to read tuned texts later + @property + def learnable_system_prompt(self) -> Any: + """Get the learnable system prompt.""" + return self._system_prompt + + @property + def learnable_tool_prompts(self) -> List[Any]: + """Get the learnable tool prompts.""" + return self._tool_prompts \ No newline at end of file diff --git a/python/packages/autogen-ext/src/autogen_ext/optimize/dspy.py b/python/packages/autogen-ext/src/autogen_ext/optimize/dspy.py new file mode 100644 index 000000000000..cf617216892f --- /dev/null +++ b/python/packages/autogen-ext/src/autogen_ext/optimize/dspy.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import importlib +from typing import Any, Callable, Dict, Iterable, Tuple, TYPE_CHECKING + +if TYPE_CHECKING: + try: + import dspy + except ImportError: + dspy = None + +from autogen_agentchat.optimize._backend import BaseBackend +from ._utils import AutoGenLM, DSPyAgentWrapper + + +def _check_dspy_available() -> None: + """Check if DSPy is available and raise helpful error if not.""" + try: + import dspy # noqa: F401 + except ImportError as e: + raise ImportError( + "DSPy is required for optimization but not installed. " + "Please install it with: pip install dspy" + ) from e + + +class DSPyBackend(BaseBackend): + """Optimise AutoGen agents with any DSPy optimiser.""" + + name = "dspy" + + def compile( + self, + agent: Any, + trainset: Iterable[Any], + metric: Callable[[Any, Any], float | bool], + *, + lm_client: Any | None = None, + optimizer_name: str = "SIMBA", + optimizer_kwargs: Dict[str, Any] | None = None, + **extra: Any, + ) -> Tuple[Any, Dict[str, Any]]: + """ + Compile/optimize an AutoGen agent using DSPy. + + Parameters + ---------- + agent + AutoGen agent to optimize + trainset + Training examples for optimization + metric + Evaluation metric function + lm_client + Language model client (if None, uses agent.model_client) + optimizer_name + Name of DSPy optimizer to use + optimizer_kwargs + Additional arguments for the optimizer + extra + Additional keyword arguments (ignored) + + Returns + ------- + Tuple of (optimized_agent, report) + """ + _check_dspy_available() + import dspy + + if not optimizer_kwargs: + optimizer_kwargs = {} + + # 1. Configure DSPy with the supplied AutoGen LM (or grab from agent) + lm_client = lm_client or getattr(agent, "model_client", None) or getattr(agent, "_model_client", None) + if lm_client is None: + raise ValueError("Could not find model_client in agent and none provided") + + dspy_lm = AutoGenLM(lm_client) + dspy.configure(lm=dspy_lm) + + # 2. Wrap agent + wrapper = DSPyAgentWrapper(agent) + + # 3. Create optimiser instance + try: + opt_mod = importlib.import_module("dspy.optimizers") + OptimCls = getattr(opt_mod, optimizer_name) + except (ImportError, AttributeError) as e: + raise ValueError(f"Could not find optimizer '{optimizer_name}' in dspy.optimizers") from e + + optimiser = OptimCls(metric=metric, **optimizer_kwargs) + + # 4. Compile + compiled = optimiser.compile(wrapper, trainset=trainset) + + # 5. Write back tuned texts into the *original* live agent + if hasattr(agent, "system_message"): + agent.system_message = compiled.learnable_system_prompt.value + elif hasattr(agent, "_system_messages") and agent._system_messages: + from autogen_core.models import SystemMessage + agent._system_messages[0] = SystemMessage(content=compiled.learnable_system_prompt.value) + + # Update tool descriptions + tools = getattr(agent, "_tools", []) or getattr(agent, "tools", []) + for new_desc, tool in zip(compiled.learnable_tool_prompts, tools): + if hasattr(tool, "description"): + tool.description = new_desc.value + + # Prepare report + best_metric = getattr(optimiser, "best_metric", None) + + # Get current system message for report + current_system_message = None + if hasattr(agent, "system_message"): + current_system_message = agent.system_message + elif hasattr(agent, "_system_messages") and agent._system_messages: + current_system_message = agent._system_messages[0].content + + report = { + "optimizer": optimizer_name, + "best_metric": best_metric, + "tuned_system_prompt": current_system_message, + "tuned_tool_descriptions": [ + getattr(t, "description", "") for t in tools + ], + } + + return agent, report \ No newline at end of file