Python: [BREAKING] Update Agent Framework Lab Lightning to use Agent-lightning v0.2.0 API (#1644)

* Merge changes from AGL release

* Merge changes from AGL release

* fix mypy

* fix tool call with pydantic

* Apply suggestion from @ekzhu

* fix lint

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>

python/packages/lab/lightning/README.md

@@ -26,7 +26,7 @@ pip install -e ".[lightning,math]"
 pip install -e ".[lightning,tau2]"
 ```
 
-To prepare for RL training, you'll also need to install dependencies like PyTorch, Ray, and vLLM. See the [Agent-lightning setup instructions](https://github.com/microsoft/agent-lightning) for more details.
+To prepare for RL training, you'll also need to install dependencies like PyTorch, Ray, and vLLM. See the [Agent-lightning setup instructions](https://microsoft.github.io/agent-lightning/stable/tutorials/installation/) for more details.
 
 ## Usage Patterns
 
@@ -40,7 +40,7 @@ The basic usage pattern follows these steps:
 ### Example Implementation
 
 ```python
-from agent_framework.lab.lightning import init
+from agent_framework.lab.lightning import AgentFrameworkTracer
 from agentlightning import rollout, Trainer, LLM, Dataset
 from agentlightning.algorithm.verl import VERL
 
@@ -73,10 +73,10 @@ config = {
     # ... additional config
 }
 
-# Initialize agent-framework to send telemetry data to agent-lightning's observability backend
-init()
+# Initialize agent-framework tracer to send telemetry data to agent-lightning's observability backend
+tracer = AgentFrameworkTracer()
 
-trainer = Trainer(algorithm=VERL(config), n_workers=2)
+trainer = Trainer(algorithm=VERL(config), tracer=tracer, n_workers=2)
 # Both train_dataset and val_dataset are lists of TaskType
 trainer.fit(math_agent, train_dataset, val_data=val_dataset)
 ```

python/packages/lab/lightning/agent_framework_lab_lightning/__init__.py

@@ -2,12 +2,10 @@
 
 """RL Module for Microsoft Agent Framework."""
 
-# ruff: noqa: F403
-
 import importlib.metadata
 
 from agent_framework.observability import OBSERVABILITY_SETTINGS
-from agentlightning import *  # type: ignore
+from agentlightning import AgentOpsTracer  # type: ignore
 
 try:
     __version__ = importlib.metadata.version(__name__)
@@ -15,9 +13,22 @@ except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"  # Fallback for development mode
 
 
-def init() -> None:
-    """Initialize the agent-framework-lab-lightning for training."""
-    OBSERVABILITY_SETTINGS.enable_otel = True
+class AgentFrameworkTracer(AgentOpsTracer):  # type: ignore
+    """Tracer for Agent-framework.
+
+    Tracer that enables OpenTelemetry observability for the Agent-framework,
+    so that the traces are visible to Agent-lightning.
+    """
+
+    def init(self) -> None:
+        """Initialize the agent-framework-lab-lightning for training."""
+        OBSERVABILITY_SETTINGS.enable_otel = True
+        super().init()
+
+    def teardown(self) -> None:
+        """Teardown the agent-framework-lab-lightning for training."""
+        super().teardown()
+        OBSERVABILITY_SETTINGS.enable_otel = False
 
 
-__all__: list[str] = ["init"]
+__all__: list[str] = ["AgentFrameworkTracer"]

python/packages/lab/lightning/samples/train_math_agent.py

@@ -18,11 +18,9 @@ import string
 from typing import TypedDict, cast
 
 import sympy  # type: ignore[import-untyped,reportMissingImports]
-from agent_framework._agents import ChatAgent
-from agent_framework._mcp import MCPStdioTool
-from agent_framework._types import AgentRunResponse
-from agent_framework.openai._chat_client import OpenAIChatClient
-from agent_framework_lab_lightning import init as lightning_init
+from agent_framework import AgentRunResponse, ChatAgent, MCPStdioTool
+from agent_framework.lab.lightning import AgentFrameworkTracer
+from agent_framework.openai import OpenAIChatClient
 from agentlightning import LLM, Dataset, Trainer, rollout
 from agentlightning.algorithm.verl import VERL
 
@@ -192,10 +190,6 @@ def main():
     # This configuration controls all aspects of the RL training process.
     # Key sections: algorithm, data, rollout, actor, trainer
     rl_training_config = {
-        "agentlightning": {
-            # The port to communicate between the rollout workers and the RL training process
-            "port": 9999,
-        },
         "algorithm": {
             # Advantage estimator type: "gae", "grpo", "reinforce_plus_plus", etc.
             "adv_estimator": "grpo"
@@ -280,10 +274,6 @@ def main():
         },
     }
 
-    # Initialize and run training
-    # lightning_init() enables observability integration with agent-framework
-    lightning_init()
-
     # Load your datasets
     train_dataset = _load_jsonl("data/math/train.jsonl")
     val_dataset = _load_jsonl("data/math/test.jsonl")
@@ -298,13 +288,13 @@ def main():
 
     # Create trainer with VERL algorithm and start training
     # n_workers: Number of rollout workers (processes) for parallel data collection
-    trainer = Trainer(algorithm=VERL(rl_training_config), n_workers=2)
+    trainer = Trainer(algorithm=VERL(rl_training_config), tracer=AgentFrameworkTracer(), n_workers=2)
 
     # This starts the actual RL training loop:
     # 1. Collect rollouts using current model
     # 2. Compute advantages and train the model
     # 3. Repeat for specified number of epochs
-    trainer.fit(math_agent, train_dataset, val_data=val_dataset)
+    trainer.fit(math_agent, train_dataset, val_dataset=val_dataset)
 
 
 def debug():

python/packages/lab/lightning/samples/train_tau2_agent.py

@@ -17,14 +17,15 @@ import asyncio
 import json
 import os
 import random
+import time
 import traceback
 from pathlib import Path
 from typing import TypedDict, cast
 
+from agent_framework.lab.lightning import AgentFrameworkTracer
 from agent_framework.lab.tau2 import ASSISTANT_AGENT_ID, patch_env_set_state  # type: ignore
 from agent_framework.lab.tau2 import TaskRunner as Tau2TaskRunner  # type: ignore
 from agent_framework.openai import OpenAIChatClient
-from agent_framework_lab_lightning import init as lightning_init
 from agentlightning import LLM, Dataset, LitAgent, NamedResources, Rollout, Trainer
 from agentlightning.algorithm.verl import VERL
 from tau2.data_model.tasks import Task as Tau2Task  # type: ignore[import-untyped]
@@ -133,9 +134,6 @@ def main():
     """Main entrypoint."""
     # RL config with higher resource requirements and W&B logging
     rl_training_config = {
-        "agentlightning": {
-            "port": 9999,
-        },
         "algorithm": {"adv_estimator": "grpo"},
         "data": {
             "train_batch_size": 8,
@@ -187,7 +185,6 @@ def main():
         },
     }
 
-    lightning_init()
     patch_env_set_state()  # Tau2-specific environment setup
 
     train_dataset, val_dataset = _load_dataset()
@@ -196,14 +193,13 @@ def main():
     # Only the assistant agent is trained; user simulator remains fixed
     tau2_agent = Tau2Agent(trained_agents=ASSISTANT_AGENT_ID)
 
-    trainer = Trainer(algorithm=VERL(rl_training_config), n_workers=4)
-    trainer.fit(tau2_agent, train_dataset, val_data=val_dataset)
+    tracer = AgentFrameworkTracer()
+    trainer = Trainer(algorithm=VERL(rl_training_config), tracer=tracer, n_workers=4)
+    trainer.fit(tau2_agent, train_dataset, val_dataset=val_dataset)
 
 
 def debug():
     """Debug mode for testing multi-agent setup and Tau2 integration."""
-    lightning_init()
-
     train_dataset, _ = _load_dataset()
     tau2_agent = Tau2Agent(trained_agents=ASSISTANT_AGENT_ID)
 
@@ -218,7 +214,7 @@ def debug():
         tau2_agent.rollout_async(
             train_dataset[0],
             resources={"main_llm": LLM(model="gpt-4.1", endpoint=openai_base_url)},
-            rollout=Rollout(rollout_id="dummy"),
+            rollout=Rollout(rollout_id="dummy", input="dummy_input", start_time=time.time()),
         )
     )
 

python/packages/lab/lightning/tests/test_lightning.py

@@ -7,14 +7,13 @@ from unittest.mock import AsyncMock, patch
 import pytest
 from agent_framework import (
     AgentExecutor,
+    AgentRunEvent,
     ChatAgent,
     WorkflowBuilder,
 )
-from agent_framework._workflows._events import AgentRunEvent
+from agent_framework.lab.lightning import AgentFrameworkTracer
 from agent_framework.openai import OpenAIChatClient
-from agent_framework_lab_lightning import init
-from agentlightning.adapter import TraceTripletAdapter
-from agentlightning.tracer import AgentOpsTracer
+from agentlightning import TracerTraceToTriplet
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
 
@@ -134,30 +133,30 @@ async def test_observability(workflow_two_agents):
             |                    |
         [chat gpt-4o]        [chat gpt-4o]
     """
-    init()
-
-    tracer = AgentOpsTracer()
+    tracer = AgentFrameworkTracer()
     try:
         tracer.init()
         tracer.init_worker(0)
 
-        with tracer.trace_context():
+        async with tracer.trace_context():
             await workflow_two_agents.run("Please analyze the quarterly sales data")
 
-        triplets = TraceTripletAdapter(agent_match=None, llm_call_match="chat").adapt(tracer.get_last_trace())
+        triplets = TracerTraceToTriplet(agent_match=None, llm_call_match="chat").adapt(tracer.get_last_trace())
         assert len(triplets) == 2
 
-        triplets = TraceTripletAdapter(agent_match="analyzer", llm_call_match="chat").adapt(tracer.get_last_trace())
+        triplets = TracerTraceToTriplet(agent_match="analyzer", llm_call_match="chat").adapt(tracer.get_last_trace())
         assert len(triplets) == 1
 
-        triplets = TraceTripletAdapter(agent_match="advisor", llm_call_match="chat").adapt(tracer.get_last_trace())
+        triplets = TracerTraceToTriplet(agent_match="advisor", llm_call_match="chat").adapt(tracer.get_last_trace())
         assert len(triplets) == 1
 
         # Parent agent is not matched
-        triplets = TraceTripletAdapter(agent_match="DataAnalyzer", llm_call_match="chat").adapt(tracer.get_last_trace())
+        triplets = TracerTraceToTriplet(agent_match="DataAnalyzer", llm_call_match="chat").adapt(
+            tracer.get_last_trace()
+        )
         assert len(triplets) == 0
 
-        triplets = TraceTripletAdapter(agent_match="InvestmentAdvisor|advisor", llm_call_match="chat").adapt(
+        triplets = TracerTraceToTriplet(agent_match="InvestmentAdvisor|advisor", llm_call_match="chat").adapt(
             tracer.get_last_trace()
         )
         assert len(triplets) == 1

python/packages/lab/pyproject.toml

@@ -36,8 +36,7 @@ gaia = [
 
 # Lightning RL training module dependencies
 lightning = [
-  # Till 0.2.0 is released
-  "agentlightning @ git+https://github.com/microsoft/agent-lightning@138ad0e48780b65ccefa628ee7a5fb9fb27aca01",
+  "agentlightning>=0.2.0,<0.3.0",
 ]
 
 # TAU2 benchmark module dependencies

python/pyproject.toml

@@ -68,6 +68,13 @@ environments = [
     "sys_platform == 'linux'",
     "sys_platform == 'win32'"
 ]
+override-dependencies = [
+  # A conflict between the dependency of litellm[proxy] < 0.30.0, which is a dependency of agent-lightning
+  # and uvicorn >= 0.34.0, which is a dependency of tau2
+  "uvicorn==0.38.0",
+  # Similar problem with websockets, which is a dependency conflict between litellm[proxy] and mcp
+  "websockets==15.0.1",
+]
 
 [tool.uv.workspace]
 members = [ "packages/*" ]