@@ -21,6 +21,7 @@ class TurnTrackingRunHooks(RunHooks):
2121 """Records turn numbers seen by on_turn_start and on_turn_end."""
2222
2323 def __init__ (self ) -> None :
24+ """Initialise empty tracking lists and event counters."""
2425 self .turn_starts : list [int ] = []
2526 self .turn_ends : list [int ] = []
2627 self .events : dict [str , int ] = defaultdict (int )
@@ -31,6 +32,7 @@ async def on_turn_start(
3132 agent : Any ,
3233 turn_number : int ,
3334 ) -> None :
35+ """Record the turn number when a turn starts."""
3436 self .turn_starts .append (turn_number )
3537 self .events ["on_turn_start" ] += 1
3638
@@ -40,6 +42,7 @@ async def on_turn_end(
4042 agent : Any ,
4143 turn_number : int ,
4244 ) -> None :
45+ """Record the turn number when a turn ends."""
4346 self .turn_ends .append (turn_number )
4447 self .events ["on_turn_end" ] += 1
4548
@@ -48,6 +51,7 @@ class TurnTrackingAgentHooks(AgentHooks):
4851 """Records turn numbers seen on agent-level hooks."""
4952
5053 def __init__ (self ) -> None :
54+ """Initialise empty tracking lists."""
5155 self .turn_starts : list [int ] = []
5256 self .turn_ends : list [int ] = []
5357
@@ -57,6 +61,7 @@ async def on_turn_start(
5761 agent : Any ,
5862 turn_number : int ,
5963 ) -> None :
64+ """Record the turn number when a turn starts."""
6065 self .turn_starts .append (turn_number )
6166
6267 async def on_turn_end (
@@ -65,6 +70,7 @@ async def on_turn_end(
6570 agent : Any ,
6671 turn_number : int ,
6772 ) -> None :
73+ """Record the turn number when a turn ends."""
6874 self .turn_ends .append (turn_number )
6975
7076
@@ -109,16 +115,22 @@ async def test_on_turn_start_fires_before_llm() -> None:
109115 call_order : list [str ] = []
110116
111117 class OrderTrackingHooks (RunHooks ):
118+ """Tracks the order that turn/LLM lifecycle hooks are called."""
119+
112120 async def on_turn_start (self , context : Any , agent : Any , turn_number : int ) -> None :
121+ """Append a turn_start marker with the turn number."""
113122 call_order .append (f"turn_start:{ turn_number } " )
114123
115124 async def on_llm_start (self , context : Any , agent : Any , system_prompt : Any , input_items : Any ) -> None :
116- call_order .append (f"llm_start" )
125+ """Append an llm_start marker."""
126+ call_order .append ("llm_start" )
117127
118128 async def on_llm_end (self , context : Any , agent : Any , response : Any ) -> None :
119- call_order .append (f"llm_end" )
129+ """Append an llm_end marker."""
130+ call_order .append ("llm_end" )
120131
121132 async def on_turn_end (self , context : Any , agent : Any , turn_number : int ) -> None :
133+ """Append a turn_end marker with the turn number."""
122134 call_order .append (f"turn_end:{ turn_number } " )
123135
124136 model = FakeModel ()
0 commit comments