Callback example

gistfile1.txt

"""Custom callback handlers for sbagents CLI."""

import sys
from datetime import datetime
from typing import Any


# ANSI color codes
GRAY = "\033[38;5;245m"
ORANGE = "\033[38;5;214m"
RED = "\033[38;5;196m"
RESET = "\033[0m"

# Global cost tracker across all handlers
_total_input_cost = [0.0]
_total_output_cost = [0.0]


def get_total_input_cost() -> float:
    """Get total input cost accumulated across all handler invocations."""
    return _total_input_cost[0]


def get_total_output_cost() -> float:
    """Get total output cost accumulated across all handler invocations."""
    return _total_output_cost[0]


def get_total_cost() -> float:
    """Get total cost (input + output) accumulated across all handler invocations."""
    return _total_input_cost[0] + _total_output_cost[0]


def get_total_input_tokens() -> int:
    """Get total input tokens accumulated across all handler invocations."""
    return 0  # Not tracked separately for now


def get_total_output_tokens() -> int:
    """Get total output tokens accumulated across all handler invocations."""
    return 0  # Not tracked separately for now


def reset_total_cost():
    """Reset total cost and token counters."""
    _total_input_cost[0] = 0.0
    _total_output_cost[0] = 0.0


def create_cli_callback_handler(agent_id: str, stream_output: bool = True):
    """Create a callback handler for CLI output.
    
    Args:
        agent_id: Agent identifier for logging
        stream_output: Whether to stream output to stdout
        
    Returns:
        Callback handler function
    """
    tool_name = [None]
    model_id = [None]
    total_usage = [{"inputTokens": 0, "outputTokens": 0, "totalTokens": 0}]

    def get_timestamp() -> str:
        return datetime.now().strftime("%H:%M:%S")

    def callback_handler(**kwargs: Any):
        nonlocal tool_name, model_id, total_usage
        
        # Store model from first data event
        if "data" in kwargs and not model_id[0] and "agent" in kwargs:
            model = kwargs["agent"].model
            model_id[0] = model.config.get("model_id", str(model))

        if "current_tool_use" in kwargs:
            tool = kwargs["current_tool_use"]
            if tool.get("name"):
                tool_name[0] = tool["name"]
                timestamp = get_timestamp()
                print(f"\n{GRAY}{timestamp}{RESET} {ORANGE}[{agent_id} uses tool: {tool['name']}]{RESET}", file=sys.stderr, flush=True)

        if stream_output and "data" in kwargs:
            print(kwargs["data"], end="", flush=True)

        if "message" in kwargs:
            msg = kwargs["message"]
            if msg.get("role") == "user":
                for block in msg.get("content", []):
                    if "toolResult" in block:
                        result = block["toolResult"]
                        tool_id = result.get("toolUseId", "unknown")
                        content = result.get("content", [])
                        result_text = ""
                        for c in content:
                            if "text" in c:
                                result_text = c["text"]
                        timestamp = get_timestamp()
                        print(f"\n{GRAY}{timestamp}{RESET} {ORANGE}[tool result: {tool_name[0] or tool_id} = {result_text}]{RESET}", file=sys.stderr, flush=True)
                        tool_name[0] = None

        # Track model usage from result event
        if "result" in kwargs:
            result = kwargs["result"]
            metrics = result.metrics
            
            # Get model from agent
            if "agent" in kwargs:
                model = kwargs["agent"].model
                model_id[0] = model.config.get("model_id", str(model))
            
            # Get accumulated usage
            if hasattr(metrics, "accumulated_usage") and metrics.accumulated_usage:
                total_usage[0] = metrics.accumulated_usage
            
            # Print final stats
            if model_id[0] and total_usage[0]:
                timestamp = get_timestamp()
                input_tok = total_usage[0].get("inputTokens", 0)
                output_tok = total_usage[0].get("outputTokens", 0)
                print(f"\n{GRAY}{timestamp}{RESET} {ORANGE}[tokens: in={input_tok} out={output_tok}]{RESET}", file=sys.stderr, flush=True)
                

    return callback_handler