33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
import time
|
|
from backend.logger import setup_logger
|
|
from agents import AgentHooks, RunContextWrapper, Agent
|
|
|
|
logger = setup_logger(__name__)
|
|
|
|
class LegalAssistantHooks(AgentHooks):
|
|
"""Tracks execution time, tool calls and token usage per agent run."""
|
|
|
|
def __init__(self):
|
|
self.start_time: float = 0
|
|
self.tool_calls: int = 0
|
|
|
|
async def on_start(self, ctx: RunContextWrapper, agent: Agent) -> None:
|
|
"""Resets counters and logs agent start."""
|
|
self.start_time = time.time()
|
|
self.tool_calls = 0
|
|
logger.info(f"[AGENT {agent.name} STARTED]")
|
|
|
|
async def on_tool_start(self, ctx: RunContextWrapper, agent: Agent, tool) -> None:
|
|
"""Logs each tool invocation."""
|
|
self.tool_calls += 1
|
|
logger.info(f"[TOOL {tool.name} CALLED]")
|
|
|
|
async def on_tool_end(self, ctx: RunContextWrapper, agent: Agent, tool, result) -> None:
|
|
"""Marker indicating the end of the tool invocation."""
|
|
logger.info(f"[TOOL {tool.name} ENDED]")
|
|
|
|
async def on_end(self, ctx: RunContextWrapper, agent: Agent, output) -> None:
|
|
"""Logs elapsed time and total tokens used."""
|
|
elapsed = time.time() - self.start_time
|
|
logger.info(f"[AGENT {agent.name} ENDED] | tools_called={self.tool_calls} | elapsed={elapsed:.2f}s")
|
|
|