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