usage_logger.py

"""Logger for tracking MCP tool usage and token estimates."""

import logging
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

class UsageLogger:
    """Track MCP tool calls and estimate token usage."""
    
    def __init__(self, log_file: str = "usage_log.jsonl"):
        self.log_file = Path(log_file)
        self.setup_logging()
    
    def setup_logging(self):
        """Configure logging."""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger("pdf-rag")
    
    def estimate_tokens(self, text: str) -> int:
        """
        Rough token estimate (1 token ≈ 4 characters for English).
        This is approximate - actual tokenization may vary.
        """
        return len(text) // 4
    
    def log_tool_call(self, tool_name: str, arguments: Dict[str, Any], 
                     result: Any, response_text: str):
        """Log a tool call with token estimates."""
        args_str = json.dumps(arguments)
        
        # Estimate tokens
        input_tokens = self.estimate_tokens(args_str)
        output_tokens = self.estimate_tokens(response_text)
        total_tokens = input_tokens + output_tokens
        
        # Create log entry
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            "tool": tool_name,
            "arguments": arguments,
            "input_tokens_estimate": input_tokens,
            "output_tokens_estimate": output_tokens,
            "total_tokens_estimate": total_tokens,
            "output_length_chars": len(response_text),
            "output_length_words": len(response_text.split())
        }
        
        # Write to JSONL file (one JSON object per line)
        with open(self.log_file, 'a') as f:
            f.write(json.dumps(log_entry) + '\n')
        
        # Also log to console
        self.logger.info(
            f"Tool: {tool_name} | "
            f"Tokens: ~{total_tokens} (in: {input_tokens}, out: {output_tokens}) | "
            f"Output: {len(response_text)} chars"
        )
        
        return log_entry
    
    def get_usage_summary(self) -> Dict[str, Any]:
        """Get summary of total usage from log file."""
        if not self.log_file.exists():
            return {"error": "No usage data yet"}
        
        total_tokens = 0
        total_calls = 0
        tool_usage = {}
        
        with open(self.log_file, 'r') as f:
            for line in f:
                if line.strip():
                    entry = json.loads(line)
                    total_tokens += entry.get("total_tokens_estimate", 0)
                    total_calls += 1
                    
                    tool = entry.get("tool", "unknown")
                    if tool not in tool_usage:
                        tool_usage[tool] = {"calls": 0, "tokens": 0}
                    tool_usage[tool]["calls"] += 1
                    tool_usage[tool]["tokens"] += entry.get("total_tokens_estimate", 0)
        
        return {
            "total_calls": total_calls,
            "total_tokens_estimate": total_tokens,
            "tool_breakdown": tool_usage,
            "log_file": str(self.log_file)
        }