visible logs

* Update RAI transparency documentation

* Update RAI_README.md

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

RAI_README.md

@@ -0,0 +1,77 @@
+# Responsible AI Transparency Documentation - Agent Lightning
+
+## OVERVIEW
+
+Agent Lightning is a flexible and extensible framework that enables seamless agent optimization for any existing agent framework. Agent optimization includes various data-driven techniques to customize the agent for better performance, including but not limited to model fine-tuning, prompt tuning, and model selection. And the agent frameworks refer to popular and easy-to-use agent developing frameworks such as OpenAI Agents SDK, Microsoft AutoGen, and LangChain.
+
+### WHAT CAN AGENT LIGHTNING DO
+Agent lightning was developed to bridge the gap between agent workflow development and agent optimization, empowering developers to go beyond static, pre-trained models and unlock the full potential of adaptive, learning-based agents. Agent Lightning is a training framework which can be used for any LLMs.
+
+### INTENDED USES
+Agent Lightning is best suited for agent researchers and developers. They can easily fine-tune models in existing agent frameworks with Agent Lightning. This can improve model performance on the targeted scenarios.
+
+### OUT-OF-SCOPE USES
+Agent Lightning is not well-suited for users who are not familiar with agent development and machine learning concepts.
+
+We do not recommend using Agent Lightning in commercial or real-world applications without further testing and development. It is being released for research purposes.
+
+Agent Lightning was not designed or evaluated for all possible downstream purposes. Developers should consider its inherent limitations as they select use cases, and evaluate and mitigate for accuracy, safety, and fairness concerns specific to each intended downstream use.
+
+Agent Lightning should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.
+
+We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare). 
+
+## HOW TO GET STARTED
+To begin using Agent Lightning, here are some instructions.
+1.	Install dependencies, including Python, uv, PyTorch, FlashAttention, vLLM, verl.
+2.	Clone and install Agent Lightning.
+3.	Convert the dataset (provided by the user) into parquet file, which contains multiple columns. Each column contains a data id, an input and an expected output.
+4.	Run agent, which is developed by the user.
+5.	Run the training process via “bash train.sh”
+
+## EVALUATION
+Agent Lightning was evaluated on its ability to correctly complete 3 example tasks: (1) Math. The model needs to answer some math questions, and when answering one question, the model can use the calculator as its tool to help answer. (2) Text2SQL. The model is given a question related to the database, and it is required to generate a SQL which can query the database, find the information to answer the question. (3) Retrieval-Augmented Generation (RAG). The model is given a question which needs some information from Wikipedia to answer. The model is required to generate some queries to find the related information in Wikipedia, and answer the question according to retrieved documents.
+
+### EVALUATION METHODS AND RESULTS
+For detailed evaluation methods and results, please refer to the latest version of our  [technical report](https://arxiv.org/abs/2508.03680). 
+
+
+## LIMITATIONS
+Agent Lightning was developed for research and experimental purposes. Further testing and validation are needed before considering its application in commercial or real-world scenarios.
+
+Agent Lightning was designed and tested using the English language. Performance in other languages may vary and should be assessed by someone who is both an expert in the expected outputs and a native speaker of that language.
+
+Outputs generated by AI may include factual errors, fabrication, or speculation. Users are responsible for assessing the accuracy of generated content. All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
+Agent Lightning inherits any biases, errors, or omissions produced by its base model. Developers are advised to choose an appropriate base LLM/MLLM carefully, depending on the intended use case.
+We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model. 
+
+## BEST PRACTICES
+Better performance can be achieved by following the instructions in how to get started section.
+
+We strongly encourage users to use LLMs/MLLMs that support robust Responsible AI mitigations, such as Azure Open AI (AOAI) services. Such services continually update their safety and RAI mitigations with the latest industry standards for responsible use. For more on AOAI’s best practices when employing foundations models for scripts and applications:
+- [Blog post on responsible AI features in AOAI that were presented at Ignite 2023](https://techcommunity.microsoft.com/t5/ai-azure-ai-services-blog/announcing-new-ai-safety-amp-responsible-ai-features-in-azure/ba-p/3983686)
+- [Overview of Responsible AI practices for Azure OpenAI models](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/overview)
+- [Azure OpenAI Transparency Note](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/transparency-note)
+- [OpenAI’s Usage policies](https://openai.com/policies/usage-policies)
+- [Azure OpenAI’s Code of Conduct](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/code-of-conduct)
+  
+Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.   
+
+Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning. 
+
+It is the user’s responsibility to ensure that the use of Agent Lightning complies with relevant data protection regulations and organizational guidelines.
+
+## LICENSE
+We use the MIT license.
+
+## CONTACT
+We welcome feedback and collaboration from our audience. If you have suggestions, questions, or observe unexpected/offensive behavior in our technology, please contact us at agent-lightning@microsoft.com.
+
+If the team receives reports of undesired behavior or identifies issues independently, we will update this repository with appropriate mitigations.
+
+
+
+---
+
+*Last updated: September 6, 2025*
+*Document version: 1.0*

agentlightning/litagent.py

@@ -94,32 +94,6 @@ class LitAgent:
             raise ValueError("Runner reference is no longer valid (object has been garbage collected).")
         return runner
 
-    def on_rollout_start(self, task: Task, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called immediately before a rollout begins.
-
-        Args:
-            task: The :class:`Task` object that will be processed.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method to implement custom logic such as
-        logging, metric collection, or resource setup. By default, this is a
-        no-op.
-        """
-
-    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer) -> None:
-        """Hook called after a rollout completes.
-
-        Args:
-            task: The :class:`Task` object that was processed.
-            rollout: The resulting :class:`Rollout` object.
-            runner: The :class:`AgentRunner` managing the rollout.
-            tracer: The tracer instance associated with the runner.
-
-        Subclasses can override this method for cleanup or additional
-        logging. By default, this is a no-op.
-        """
-
     def training_rollout(self, task: TaskInput, rollout_id: str, resources: NamedResources) -> RolloutRawResult:
         """Defines the agent's behavior for a single training task.
 

agentlightning/logging.py

@@ -1,4 +1,16 @@
+from __future__ import annotations
+
+import json
 import logging
+from typing import Optional, TYPE_CHECKING
+import numpy as np
+
+from .types import Hook
+
+if TYPE_CHECKING:
+    from .types import Task, Rollout
+    from .tracer import BaseTracer
+    from .runner import AgentRunner
 
 
 def configure_logger(level: int = logging.INFO, name: str = "agentlightning") -> logging.Logger:
@@ -14,3 +26,192 @@ def configure_logger(level: int = logging.INFO, name: str = "agentlightning") ->
     logger.setLevel(level)
     logger.propagate = False  # prevent double logging
     return logger
+
+
+logger = logging.getLogger(__name__)
+
+
+class LightningLogger(Hook):
+    """Agent-lightning logger that supports tracing events and metrics throughout the training."""
+
+    def log_event(self, event: str, data: dict):
+        """
+        Log an event with its associated data when something happens.
+        """
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        """
+        Log a metric with its value and an optional step.
+        """
+
+    def log_message(self, level: int, message: str):
+        """
+        Log a message at a specific logging level.
+        """
+
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer):
+        """
+        By default, each logger automatically logs the rollout event at the end of each rollout.
+        """
+        self.log_event("rollout", {"task": task.model_dump(), "rollout": rollout.model_dump()})
+
+
+class ConsoleLogger(LightningLogger):
+    """A simple logger that logs messages to the console using Python's logging module."""
+
+    def __init__(self, level: int = logging.INFO):
+        self.logger = configure_logger(level, name="agentlightning.ConsoleLogger")
+        self.default_level = level
+        self.worker_id: Optional[int] = None
+
+    def init_worker(self, worker_id: int):
+        super().init_worker(worker_id)
+        self.worker_id = worker_id
+
+    def teardown_worker(self, worker_id: int):
+        super().teardown_worker(worker_id)
+        self.worker_id = None
+
+    def log_event(self, event: str, data: dict):
+        data_str = str(data)
+        if len(data_str) > 512:
+            data_str = f"{data_str[:512]}... (truncated)"
+        message = f"Event: {event}, Data: {data_str}"
+        self.log_message(self.default_level, message)
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        step_str = f" at step {step}" if step is not None else ""
+        message = f"Metric: {metric} = {value}{step_str}"
+        self.log_message(self.default_level, message)
+
+    def log_message(self, level: int, message: str):
+        if level >= self.default_level:
+            if self.worker_id is not None:
+                message = f"(Worker-{self.worker_id}) {message}"
+            else:
+                message = f"(Main) {message}"
+            self.logger.log(level, message)
+        # else skip logging if below default level
+
+
+class WandbLogger(LightningLogger):
+
+    def __init__(
+        self,
+        project: str,
+        entity: Optional[str] = None,
+        name: Optional[str] = None,
+        config: Optional[dict] = None,
+        *,
+        flush_every_n_events: int = 128,
+        aggregate_every_n_metrics: int = 128,
+    ):
+        import wandb
+        from wandb.sdk.wandb_run import Run
+
+        self.wandb_run: Optional[Run] = None
+        self.wandb_run_id: Optional[str] = None
+
+        self.project = project
+        self.entity = entity
+        self.name = name or wandb.util.generate_id()
+        self.config = config or {}
+
+        self.event_table: Optional[wandb.Table] = None
+        self.flush_every_n_events = flush_every_n_events
+        self.aggregate_every_n_metrics = aggregate_every_n_metrics
+
+        self.metrics_buffer: dict[str, list[float]] = {}
+
+    def init_worker(self, worker_id: int):
+        import wandb
+
+        super().init_worker(worker_id)
+        self.wandb_run = wandb.init(
+            project=self.project,
+            entity=self.entity,
+            group=self.name,
+            job_type=f"worker_{worker_id}",
+            config=self.config,
+        )
+        logger.info(f"Wandb run initialized: {self.name} (Worker {worker_id})")
+        if self.wandb_run is None:
+            raise RuntimeError("Failed to initialize Wandb run.")
+        self.wandb_run_id = self.wandb_run.id
+
+    def teardown_worker(self, worker_id: int):
+        import wandb
+
+        super().teardown_worker(worker_id)
+
+        for metric in self.metrics_buffer:
+            if len(self.metrics_buffer[metric]) > 0:
+                self._log_aggregated_metrics(metric)
+
+        if len(self.event_table.data) > 0:
+            logger.info(f"Flushing {len(self.event_table.data)} events to Wandb before finishing...")
+            wandb.log({"client/events": self.event_table})
+            self.event_table = None
+
+        wandb.finish(exit_code=0)
+
+    def teardown(self):
+        import wandb
+
+        super().teardown()
+        if self.wandb_run is not None:
+            wandb.finish(exit_code=0)
+            self.wandb_run = None
+            self.wandb_run_id = None
+
+    def log_event(self, event: str, data: dict):
+        import wandb
+
+        if self.event_table is None:
+            self.event_table = wandb.Table(columns=["event", "data"])
+
+        try:
+            data_str = json.dumps(data)  # Ensure data is JSON serializable
+        except (TypeError, ValueError):
+            data_str = str(data)
+        self.event_table.add_data(event, data_str)
+
+        if len(self.event_table.data) % self.flush_every_n_events == 0:
+            logger.info(f"Flushing {len(self.event_table.data)} events to Wandb...")
+            wandb.log({"client/events": self.event_table})
+
+    def log_metric(self, metric: str, value: float, step: Optional[int] = None):
+        import wandb
+
+        if step is not None:
+            wandb.log({"client_metric/" + metric: value}, step=step)
+        else:
+            wandb.log({"client_metric/" + metric: value})
+
+        if metric not in self.metrics_buffer:
+            self.metrics_buffer[metric] = []
+        self.metrics_buffer[metric].append(value)
+        if len(self.metrics_buffer[metric]) >= self.aggregate_every_n_metrics:
+            self._log_aggregated_metrics(metric, step)
+            self.metrics_buffer[metric] = []
+
+    def log_message(self, level: int, message: str):
+        pass  # Wandb handles logging internally, so we don't need to implement this
+
+    def _log_aggregated_metrics(self, metric, step: Optional[int] = None):
+        import wandb
+
+        arr = np.array(self.metrics_buffer[metric])
+        aggregated_value = {
+            "mean": float(np.mean(arr)),
+            "max": float(np.max(arr)),
+            "min": float(np.min(arr)),
+            "std": float(np.std(arr)),
+            "count": int((~np.isnan(arr)).sum()),
+        }
+        for key, value in aggregated_value.items():
+            if value is not None:
+                if step is not None:
+                    wandb.log({"client_agg/" + metric + "/" + key: value}, step=step)
+                else:
+                    wandb.log({"client_agg/" + metric + "/" + key: value})

agentlightning/runner.py

@@ -7,12 +7,11 @@ from contextlib import nullcontext
 from typing import List, Optional, Union, Dict, Any
 
 import agentops
-
 from opentelemetry.sdk.trace import ReadableSpan
+
 from .client import AgentLightningClient
 from .litagent import LitAgent
-from .types import Rollout, Task, Triplet, RolloutRawResult
-from .types import ParallelWorkerBase
+from .types import Rollout, Task, Triplet, RolloutRawResult, ParallelWorkerBase, Hook
 from .tracer.base import BaseTracer
 from .tracer import TripletExporter
 
@@ -43,12 +42,14 @@ class AgentRunner(ParallelWorkerBase):
         triplet_exporter: TripletExporter,
         worker_id: Optional[int] = None,
         max_tasks: Optional[int] = None,
+        hooks: Optional[List[Hook]] = None,
     ):
         super().__init__()
         self.agent = agent
         self.client = client
         self.tracer = tracer
         self.triplet_exporter = triplet_exporter
+        self.hooks = hooks or []
 
         # Worker-specific attributes
         self.worker_id = worker_id
@@ -158,10 +159,11 @@ class AgentRunner(ParallelWorkerBase):
         rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
 
         try:
-            try:
-                self.agent.on_rollout_start(task, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_start(task, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook: {hook}.")
 
             with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                 start_time = time.time()
@@ -180,10 +182,11 @@ class AgentRunner(ParallelWorkerBase):
         except Exception:
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
-            try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_end(task, rollout_obj, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook: {hook}.")
             self.client.post_rollout(rollout_obj)
 
         return True
@@ -227,10 +230,11 @@ class AgentRunner(ParallelWorkerBase):
         rollout_obj = Rollout(rollout_id=task.rollout_id)  # Default empty rollout
 
         try:
-            try:
-                self.agent.on_rollout_start(task, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_start(task, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook: {hook}.")
 
             with self.tracer.trace_context(name=f"rollout_{rollout_id}"):
                 start_time = time.time()
@@ -248,10 +252,11 @@ class AgentRunner(ParallelWorkerBase):
         except Exception:
             logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
         finally:
-            try:
-                self.agent.on_rollout_end(task, rollout_obj, self, self.tracer)
-            except Exception:
-                logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook.")
+            for hook in self.hooks:
+                try:
+                    hook.on_rollout_end(task, rollout_obj, self, self.tracer)
+                except Exception:
+                    logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_end hook: {hook}.")
             await self.client.post_rollout_async(rollout_obj)
 
         return True

agentlightning/trainer.py

@@ -1,16 +1,14 @@
 import asyncio
 import logging
 import multiprocessing
-import os
 import signal
 import time
 from typing import List, Optional, Union
 import importlib
 
-import agentops
-
 from .client import AgentLightningClient
 from .litagent import LitAgent
+from .logging import LightningLogger
 from .runner import AgentRunner
 from .types import ParallelWorkerBase
 from .tracer.base import BaseTracer
@@ -52,6 +50,7 @@ class Trainer(ParallelWorkerBase):
         daemon: bool = True,
         tracer: Union[BaseTracer, str, dict, None] = None,
         triplet_exporter: Union[TripletExporter, dict, None] = None,
+        loggers: Optional[List[LightningLogger]] = None,
     ):
         super().__init__()
         self.n_workers = n_workers
@@ -60,6 +59,8 @@ class Trainer(ParallelWorkerBase):
         self.dev = dev
         self._client: AgentLightningClient | None = None  # Will be initialized in fit method
 
+        self._loggers = loggers or []
+
         self.tracer = self._make_tracer(tracer)
         if isinstance(triplet_exporter, TripletExporter):
             self.triplet_exporter = triplet_exporter
@@ -109,12 +110,18 @@ class Trainer(ParallelWorkerBase):
 
         self.tracer.init()
 
+        for _logger in self.loggers:
+            _logger.init()
+
         logger.info(f"Trainer main initialization complete.")
 
     def teardown(self) -> None:
         logger.info(f"Cleaning up Trainer...")
         self.tracer.teardown()
 
+        for _logger in self.loggers:
+            _logger.teardown()
+
         self._client = None
         logger.info(f"Trainer main cleanup complete.")
 
@@ -187,6 +194,8 @@ class Trainer(ParallelWorkerBase):
                 num_processed = asyncio.run(loop.iter_async())
             else:
                 num_processed = loop.iter()
+        except KeyboardInterrupt:
+            logger.info(f"[Worker {worker_id}] KeyboardInterrupt received. Exiting worker loop...")
         except Exception:
             logger.exception(f"[Worker {worker_id}] Unhandled exception in worker loop.")
         finally:
@@ -197,11 +206,21 @@ class Trainer(ParallelWorkerBase):
     def _initialize_worker_env(self, worker_id: int):
         logger.info(f"[Worker {worker_id}] Setting up trainer environment...")  # worker_id included in process name
         self.tracer.init_worker(worker_id)
+        for _logger in self.loggers:
+            _logger.init_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Worker environment setup complete.")
 
     def _teardown_worker_env(self, worker_id: int):
         logger.info(f"[Worker {worker_id}] Cleaning up trainer environment...")
         self.tracer.teardown_worker(worker_id)
         logger.info(f"[Worker {worker_id}] Environment cleanup complete.")
+        for _logger in self.loggers:
+            _logger.teardown_worker(worker_id)
+        logger.info(f"[Worker {worker_id}] Worker environments torn down.")
+
+    @property
+    def loggers(self) -> List[LightningLogger]:
+        return self._loggers
 
     @staticmethod
     def kill_orphaned_processes() -> None:

agentlightning/types.py

@@ -1,8 +1,15 @@
-from typing import Any, Dict, List, Optional, Union, Literal, Annotated
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union, Literal, Annotated, TYPE_CHECKING
 
 from pydantic import BaseModel, Field, Discriminator
 from opentelemetry.sdk.trace import ReadableSpan
 
+if TYPE_CHECKING:
+    from .runner import AgentRunner
+    from .tracer import BaseTracer
+
+
 __all__ = [
     "Triplet",
     "Rollout",
@@ -18,6 +25,7 @@ __all__ = [
     "ResourcesUpdate",
     "GenericResponse",
     "ParallelWorkerBase",
+    "Hook",
 ]
 
 
@@ -202,3 +210,33 @@ class ParallelWorkerBase:
 
     def teardown(self, *args: Any, **kwargs: Any) -> None:
         pass
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    def on_rollout_start(self, task: Task, runner: AgentRunner, tracer: BaseTracer) -> None:
+        """Hook called immediately before a rollout begins.
+
+        Args:
+            task: The :class:`Task` object that will be processed.
+            runner: The :class:`AgentRunner` managing the rollout.
+            tracer: The tracer instance associated with the runner.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    def on_rollout_end(self, task: Task, rollout: Rollout, runner: AgentRunner, tracer: BaseTracer) -> None:
+        """Hook called after a rollout completes.
+
+        Args:
+            task: The :class:`Task` object that was processed.
+            rollout: The resulting :class:`Rollout` object.
+            runner: The :class:`AgentRunner` managing the rollout.
+            tracer: The tracer instance associated with the runner.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """

examples/apo/apo_client.py

@@ -35,5 +35,10 @@ if __name__ == "__main__":
     configure_logger()
     dotenv.load_dotenv()
     agent = SimpleAgent()
-    trainer = Trainer(n_workers=2)
+    from agentlightning.logging import ConsoleLogger, WandbLogger
+    loggers = [
+        ConsoleLogger(),
+        WandbLogger(project="agent-lightning-debug"),
+    ]
+    trainer = Trainer(n_workers=2, loggers=loggers)
     trainer.fit(agent, backend="http://127.0.0.1:9997")

pyproject.toml

@@ -14,6 +14,9 @@ dependencies = [
   "uvicorn",
   "fastapi",
   "aiohttp",
+  "opentelemetry-api",
+  "opentelemetry-sdk",
+  "numpy",
 ]
 
 [project.optional-dependencies]
@@ -49,6 +52,7 @@ agent = [
   "sqlparse",
   "nltk",
   "uv",
+  "wandb",
 ]
 
 [build-system]

scripts/validate_example_wandb.py

@@ -3,6 +3,7 @@ import sys
 
 if len(sys.argv) != 3:
     print("Usage: python validate_example_wandb.py <project> <run_name>")
+    sys.exit(1)
 
 project = sys.argv[1]
 run_name = sys.argv[2]

tests/test_litagent_hooks.py

@@ -1,8 +1,7 @@
 import pytest
 from contextlib import contextmanager
 
-
-from agentlightning import LitAgent, Task, ResourcesUpdate
+from agentlightning import Hook, LitAgent, Task, ResourcesUpdate
 from agentlightning.runner import AgentRunner
 from agentlightning.tracer import BaseTracer, TripletExporter
 
@@ -53,11 +52,6 @@ class DummyAsyncClient:
 
 
 class HookAgent(LitAgent):
-    def __init__(self):
-        super().__init__()
-        self.start_called = False
-        self.end_called = False
-        self.end_rollout = None
 
     def training_rollout(self, task, rollout_id, resources):
         return 0.5
@@ -65,6 +59,9 @@ class HookAgent(LitAgent):
     async def training_rollout_async(self, task, rollout_id, resources):
         return 0.5
 
+
+class TestHook(Hook):
+
     def on_rollout_start(self, task, runner, tracer):
         self.start_called = True
         self.start_task = task
@@ -78,12 +75,13 @@ def test_runner_calls_hooks():
     agent = HookAgent()
     client = DummyClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    hook = TestHook()
+    runner = AgentRunner(agent, client, tracer, TripletExporter(), hooks=[hook])
 
     assert runner.run() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert hook.start_called
+    assert hook.end_called
+    assert hook.end_rollout.final_reward == 0.5
 
 
 @pytest.mark.asyncio
@@ -91,9 +89,10 @@ async def test_runner_calls_hooks_async():
     agent = HookAgent()
     client = DummyAsyncClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    hook = TestHook()
+    runner = AgentRunner(agent, client, tracer, TripletExporter(), hooks=[hook])
 
     assert await runner.run_async() is True
-    assert agent.start_called
-    assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert hook.start_called
+    assert hook.end_called
+    assert hook.end_rollout.final_reward == 0.5