@@ -48,29 +48,29 @@ class FakeEvaluator:
4848
4949 def __init__ (self , block_on : str | None = None ) -> None :
5050 self .block_on = block_on
51- self .calls : List [tuple [str , dict ]] = []
51+ self .calls : List [tuple [str , dict [ str , Any ] ]] = []
5252
5353 def _record (self , hook : str , ** kwargs : Any ) -> None :
5454 self .calls .append ((hook , kwargs ))
5555 if self .block_on == hook :
5656 raise GovernanceBlockException ("blocked" ) # type: ignore[call-arg]
5757
58- def evaluate_before_agent (self , ** kwargs : Any ) -> None :
58+ def evaluate_before_agent (self , * args : Any , ** kwargs : Any ) -> Any :
5959 self ._record ("before_agent" , ** kwargs )
6060
61- def evaluate_after_agent (self , ** kwargs : Any ) -> None :
61+ def evaluate_after_agent (self , * args : Any , ** kwargs : Any ) -> Any :
6262 self ._record ("after_agent" , ** kwargs )
6363
64- def evaluate_before_model (self , ** kwargs : Any ) -> None :
64+ def evaluate_before_model (self , * args : Any , ** kwargs : Any ) -> Any :
6565 self ._record ("before_model" , ** kwargs )
6666
67- def evaluate_after_model (self , ** kwargs : Any ) -> None :
67+ def evaluate_after_model (self , * args : Any , ** kwargs : Any ) -> Any :
6868 self ._record ("after_model" , ** kwargs )
6969
70- def evaluate_tool_call (self , ** kwargs : Any ) -> None :
70+ def evaluate_tool_call (self , * args : Any , ** kwargs : Any ) -> Any :
7171 self ._record ("tool_call" , ** kwargs )
7272
73- def evaluate_after_tool (self , ** kwargs : Any ) -> None :
73+ def evaluate_after_tool (self , * args : Any , ** kwargs : Any ) -> Any :
7474 self ._record ("after_tool" , ** kwargs )
7575
7676
@@ -279,10 +279,10 @@ async def request(self, messages, settings, params):
279279 return ModelResponse (parts = [TextPart (content = "Your balance is 1000." )])
280280
281281 gm = GovernanceModel .__new__ (GovernanceModel ) # bypass WrapperModel init
282- gm .wrapped = FakeWrapped () # type: ignore[attr-defined]
282+ gm .wrapped = FakeWrapped () # type: ignore[assignment] # test double, not a real Model
283283 gm ._callbacks = cb
284284 messages = [ModelRequest (parts = [UserPromptPart (content = "What is my balance?" )])]
285- await gm .request (messages , None , None )
285+ await gm .request (messages , None , None ) # type: ignore[arg-type]
286286
287287 assert order == ["MODEL_CALL" ]
288288 assert _hooks (ev ) == ["before_model" , "after_model" ]
@@ -307,11 +307,11 @@ async def request_stream(self, *_a, **_k):
307307 yield SimpleNamespace (get = lambda : denied )
308308
309309 gm = GovernanceModel .__new__ (GovernanceModel ) # bypass WrapperModel init
310- gm .wrapped = FakeWrapped () # type: ignore[attr-defined]
310+ gm .wrapped = FakeWrapped () # type: ignore[assignment] # test double, not a real Model
311311 gm ._callbacks = cb
312312 messages = [ModelRequest (parts = [UserPromptPart (content = "hi" )])]
313313 with pytest .raises (GovernanceBlockException ):
314- async with gm .request_stream (messages , None , None ) as stream :
314+ async with gm .request_stream (messages , None , None ) as stream : # type: ignore[arg-type]
315315 assert stream is not None
316316
317317
@@ -331,11 +331,11 @@ async def request_stream(self, *_a, **_k):
331331 yield SimpleNamespace (get = lambda : final )
332332
333333 gm = GovernanceModel .__new__ (GovernanceModel )
334- gm .wrapped = FakeWrapped () # type: ignore[attr-defined]
334+ gm .wrapped = FakeWrapped () # type: ignore[assignment] # test double, not a real Model
335335 gm ._callbacks = cb
336336 messages = [ModelRequest (parts = [UserPromptPart (content = "the question" )])]
337337
338- async with gm .request_stream (messages , None , None ) as stream :
338+ async with gm .request_stream (messages , None , None ) as stream : # type: ignore[arg-type]
339339 # BEFORE_MODEL already fired; AFTER_MODEL deferred until the stream
340340 # context exits (final response is assembled).
341341 assert _hooks (ev ) == ["before_model" ]
@@ -361,12 +361,12 @@ async def request_stream(self, *_a, **_k):
361361 yield SimpleNamespace (get = lambda : final )
362362
363363 gm = GovernanceModel .__new__ (GovernanceModel )
364- gm .wrapped = FakeWrapped () # type: ignore[attr-defined]
364+ gm .wrapped = FakeWrapped () # type: ignore[assignment] # test double, not a real Model
365365 gm ._callbacks = cb
366366 messages = [ModelRequest (parts = [UserPromptPart (content = "hi" )])]
367367
368368 with pytest .raises (RuntimeError , match = "consumer blew up" ):
369- async with gm .request_stream (messages , None , None ):
369+ async with gm .request_stream (messages , None , None ): # type: ignore[arg-type]
370370 raise RuntimeError ("consumer blew up" )
371371
372372 assert "after_model" in _hooks (ev ) # ran despite the consumer error
0 commit comments