from ai_infra.callbacks import CallbacksABCBase class for callback handlers. Override methods to receive events. All methods have default no-op implementations, so you only need to override the ones you care about.
class MetricsCallbacks(Callbacks): def __init__(self): self.total_tokens = 0 self.call_count = 0 def on_llm_end(self, event: LLMEndEvent): self.call_count += 1 if event.total_tokens: self.total_tokens += event.total_tokens