Implement Agent Optimizer interface and DSPy backend

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
+```

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

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"

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

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>")
+    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