test_token_fix.py

#!/usr/bin/env python3
"""Test script to verify token counting fix."""

import asyncio
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from mcp_use import MCPAgent, MCPClient
from mcp_use.token_counting import TokenUsageSnapshot

# Load environment variables
load_dotenv()

def token_usage_callback(snapshot: TokenUsageSnapshot):
    """Test callback for token usage."""
    print(f"\n🔢 Token Usage Test:")
    print(f"   Operation: {snapshot.operation}")
    print(f"   Input tokens: {snapshot.usage.input_tokens}")
    print(f"   Output tokens: {snapshot.usage.output_tokens}")
    print(f"   Total tokens: {snapshot.usage.total_tokens}")
    print(f"   Cumulative: {snapshot.cumulative_usage.total_tokens}")

async def test_token_counting():
    """Test token counting functionality."""
    
    # Simple test config
    config = {
        "mcpServers": {
            "filesystem": {
                "command": "npx",
                "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
            }
        }
    }
    
    client = MCPClient.from_dict(config)
    
    try:
        # Create agent with token counting
        llm = ChatOpenAI(model="gpt-4o-mini", api_key=os.getenv("OPENAI_API_KEY"))
        
        agent = MCPAgent(
            llm=llm,
            client=client,
            enable_token_counting=True,
            token_tracking_callbacks=[token_usage_callback],
            max_steps=3
        )
        
        print("🚀 Testing token counting...")
        
        # Simple test query
        result = await agent.run("What is 2+2?")
        
        print(f"\n✅ Result: {result}")
        
        # Check token stats
        stats = agent.get_token_usage()
        if stats:
            print(f"\n📊 Final Stats:")
            print(f"   Total operations: {stats.get('total_operations', 0)}")
            print(f"   Total tokens: {stats.get('cumulative_usage', {}).get('total_tokens', 0)}")
        else:
            print("\n❌ No token stats available")
            
    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        await client.close_all_sessions()

if __name__ == "__main__":
    asyncio.run(test_token_counting())