debug_tokens.py

import asyncio
import os
import json
import traceback
import logging
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from mcp_use import MCPAgent, MCPClient
from mcp_use.token_counting import TokenUsageSnapshot
import tiktoken

load_dotenv()

def setup_detailed_logging():
    """Nastavenie detailného logovania pre sledovanie nástrojov"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    logger = logging.getLogger("mcp_use")
    logger.setLevel(logging.DEBUG)
    logging.getLogger("httpx").setLevel(logging.WARNING)
    return logger

class TokenAnalyzer:
    """Trieda pre analýzu tokenov z skutočného MCP volania"""
    
    def __init__(self, agent):
        self.agent = agent
        self.encoding = tiktoken.encoding_for_model("gpt-4")
        self.tool_outputs = []
        self.conversation_history = []
        
    def token_usage_callback(self, snapshot: TokenUsageSnapshot):
        """Callback pre zachytávanie token usage"""
        print(f"\n📊 [TOKEN] Token Usage:")
        print(f"    ├── Operácia: {snapshot.operation or 'agent_execution'}")
        print(f"    ├── Input tokeny: {snapshot.usage.input_tokens}")
        print(f"    ├── Output tokeny: {snapshot.usage.output_tokens}")
        print(f"    ├── Tool tokeny: {snapshot.usage.tool_tokens}")
        print(f"    ├── Celkom tokeny: {snapshot.usage.total_tokens}")
        print(f"    └── Kumulatívne: {snapshot.cumulative_usage.total_tokens}")
        
    async def run_analysis(self, query: str):
        """Spustenie analýzy s automatickým volaním MCP servera"""
        print("🔍 DEBUGGING REAL TOKEN SIZES FROM MCP CALLS")
        print("=" * 60)
        
        # 1. Analýza základných komponentov
        await self._analyze_base_components(query)
        
        # 2. Spustenie skutočného MCP volania
        print(f"\n🚀 Spúšťam skutočné MCP volanie s otázkou: '{query}'")
        print("-" * 50)
        
        try:
            # Zachytenie tool outputs cez callback
            if hasattr(self.agent, 'adapter') and self.agent.adapter:
                # Nastavenie callback pre zachytávanie tool outputs
                original_callback = getattr(self.agent.adapter, 'token_callback', None)
                self.agent.adapter.set_token_callback(self._capture_tool_output)
            
            # Spustenie agenta
            result = await self.agent.run(query, manage_connector=False)
            
            print(f"\n✅ Agent dokončil spracovanie")
            print(f"📝 Výsledok: {result}")
            
            # 3. Analýza skutočných tool outputs
            await self._analyze_real_outputs()
            
            # 4. Porovnanie s očakávanými hodnotami
            await self._compare_with_expected()
            
            return result
            
        except Exception as e:
            print(f"❌ Chyba pri spustení agenta: {e}")
            traceback.print_exc()
            return None
    
    def _capture_tool_output(self, tool_name: str, output: str):
        """Zachytenie tool output pre analýzu"""
        if output:
            self.tool_outputs.append({
                'tool_name': tool_name,
                'output': str(output),
                'length': len(str(output)),
                'tokens': len(self.encoding.encode(str(output)))
            })
            print(f"🔧 [CAPTURE] Zachytený output z {tool_name}: {len(str(output))} znakov, ~{len(self.encoding.encode(str(output)))} tokenov")
    
    async def _analyze_base_components(self, query: str):
        """Analýza základných komponentov (system prompt, tools, query)"""
        print("📝 Analýza základných komponentov:")
        
        # System prompt
        system_msg = self.agent.get_system_message()
        if system_msg:
            system_tokens = len(self.encoding.encode(system_msg.content))
            print(f"   • System prompt: {len(system_msg.content):,} znakov, {system_tokens:,} tokenov")
        
        # Tool definitions
        if hasattr(self.agent, '_tools') and self.agent._tools:
            total_tool_tokens = 0
            print(f"   • Tool definitions ({len(self.agent._tools)} nástrojov):")
            
            for tool in self.agent._tools:
                tool_def = {
                    "name": tool.name,
                    "description": tool.description,
                }
                if hasattr(tool, 'args_schema') and tool.args_schema:
                    try:
                        tool_def["parameters"] = tool.args_schema.schema()
                    except:
                        tool_def["parameters"] = {}
                
                tool_json = json.dumps(tool_def)
                tool_tokens = len(self.encoding.encode(tool_json))
                total_tool_tokens += tool_tokens
                print(f"     - {tool.name}: {tool_tokens} tokenov")
            
            print(f"   • Celkom tool definitions: {total_tool_tokens:,} tokenov")
        
        # Query
        query_tokens = len(self.encoding.encode(query))
        print(f"   • Query: '{query}' = {query_tokens} tokenov")
        
        # Conversation history
        if hasattr(self.agent, '_conversation_history'):
            history = getattr(self.agent, '_conversation_history', [])
            if history:
                history_tokens = 0
                for msg in history:
                    if hasattr(msg, 'content') and msg.content:
                        msg_tokens = len(self.encoding.encode(str(msg.content)))
                        history_tokens += msg_tokens
                print(f"   • Conversation history: {len(history)} správ, {history_tokens} tokenov")
            else:
                print(f"   • Conversation history: prázdna")
    
    async def _analyze_real_outputs(self):
        """Analýza skutočných tool outputs"""
        print(f"\n🔧 Analýza skutočných tool outputs:")
        
        # Kontrola MCP adapter callback
        if not self.tool_outputs:
            print("   ❌ Žiadne tool outputs neboli zachytené cez MCP adapter!")
            
            # Skús získať tool outputs z LangChain callback
            if hasattr(self.agent, 'adapter') and hasattr(self.agent.adapter, '_token_callback'):
                callback = self.agent.adapter._token_callback
                if hasattr(callback, 'recent_tool_outputs') and callback.recent_tool_outputs:
                    print(f"   ✅ Zachytené {len(callback.recent_tool_outputs)} tool outputs z LangChain callback:")
                    for i, output in enumerate(callback.recent_tool_outputs):
                        output_length = len(str(output))
                        output_tokens = len(self.encoding.encode(str(output))) if output else 0
                        output_preview = str(output)[:100].replace('\n', ' ')
                        print(f"     - Tool output {i+1}: {output_length} znakov, ~{output_tokens} tokenov - {output_preview}...")
                else:
                    print("   ❌ Ani LangChain callback nezachytil tool outputs!")
            else:
                print("   ❌ LangChain callback nie je dostupný!")
            return
        
        total_output_tokens = 0
        for i, output_info in enumerate(self.tool_outputs, 1):
            print(f"   • Tool {i} ({output_info['tool_name']}):")
            print(f"     - Dĺžka: {output_info['length']:,} znakov")
            print(f"     - Tokeny: {output_info['tokens']:,}")
            print(f"     - Preview: {output_info['output'][:100]}...")
            total_output_tokens += output_info['tokens']
        
        print(f"   • Celkom tool outputs: {total_output_tokens:,} tokenov")
        
        # Analýza conversation history po volaní
        if hasattr(self.agent, '_conversation_history'):
            history = getattr(self.agent, '_conversation_history', [])
            if history:
                print(f"\n💬 Conversation history po volaní:")
                total_history_tokens = 0
                for i, msg in enumerate(history):
                    if hasattr(msg, 'content') and msg.content:
                        msg_tokens = len(self.encoding.encode(str(msg.content)))
                        total_history_tokens += msg_tokens
                        msg_type = type(msg).__name__
                        content_preview = str(msg.content)[:100].replace('\n', ' ')
                        print(f"   • Správa {i+1} ({msg_type}): {msg_tokens} tokenov - {content_preview}...")
                
                print(f"   • Celkom conversation history: {total_history_tokens:,} tokenov")
    
    async def _compare_with_expected(self):
        """Porovnanie s očakávanými hodnotami z feedback"""
        print(f"\n📈 Porovnanie s očakávanými hodnotami:")
        
        # Očakávané hodnoty z feedback
        expected_total_tokens = 51923  # Celkové tokeny za 4 kroky
        expected_step4_input = 29464   # Input tokeny pre 4. krok
        expected_step4_output = 141    # Output tokeny pre 4. krok
        
        # Získanie skutočných hodnot
        if hasattr(self.agent, 'token_tracker') and self.agent.token_tracker:
            actual_stats = self.agent.get_token_usage()
            if actual_stats:
                actual_total = actual_stats.get('cumulative_usage', {}).get('total_tokens', 0)
                actual_input = actual_stats.get('cumulative_usage', {}).get('input_tokens', 0)
                actual_output = actual_stats.get('cumulative_usage', {}).get('output_tokens', 0)
                
                print(f"   • Očakávané celkové tokeny: {expected_total_tokens:,}")
                print(f"   • Skutočné celkové tokeny: {actual_total:,}")
                print(f"   • Rozdiel: {abs(expected_total_tokens - actual_total):,}")
                print(f"   • Presnosť: {(min(actual_total, expected_total_tokens) / max(actual_total, expected_total_tokens) * 100):.1f}%")
                
                print(f"\n   • Očakávané input tokeny (posledný krok): {expected_step4_input:,}")
                print(f"   • Skutočné input tokeny: {actual_input:,}")
                print(f"   • Rozdiel: {abs(expected_step4_input - actual_input):,}")
                
                print(f"\n   • Očakávané output tokeny: {expected_step4_output}")
                print(f"   • Skutočné output tokeny: {actual_output}")
                print(f"   • Rozdiel: {abs(expected_step4_output - actual_output)}")
                
                # Detailný breakdown
                print(f"\n📊 Detailný breakdown:")
                operations = actual_stats.get('operations', {})
                for op_name, op_stats in operations.items():
                    print(f"   • {op_name}: {op_stats['count']}x, {op_stats['total_tokens']} tokenov")
            else:
                print("   ❌ Token tracking nie je dostupný")
        else:
            print("   ❌ Token tracker nie je inicializovaný")

async def main():
    """Hlavná funkcia pre debug analýzu"""
    logger = setup_detailed_logging()
    
    # Konfigurácia MCP
    config = {
        "mcpServers": {
            "wordpress_server": {
                "url": os.getenv("MCP_BASE_URL"),
                "headers": {
                    "Authorization": f"Bearer {os.getenv('JWT_TOKEN')}",
                    "Content-Type": "application/json"
                }
            }
        }
    }
    
    try:
        # Vytvorenie klienta a agenta
        print("🔍 Vytváram MCP client a agent...")
        client = MCPClient.from_dict(config)
        sessions = await client.create_all_sessions()
        print(f"✅ Vytvorené sessions: {list(sessions.keys())}")
        
        llm = ChatOpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"))
        agent = MCPAgent(
            llm=llm,
            client=client,
            max_steps=15,
            memory_enabled=True,
            auto_initialize=False,
            enable_token_counting=True
        )
        
        await agent.initialize()
        print(f"✅ Agent inicializovaný s {len(agent._tools)} nástrojmi")
        
        # Vytvorenie analyzátora
        analyzer = TokenAnalyzer(agent)
        
        # Pridanie callback pre token tracking
        if hasattr(agent, 'token_tracker') and agent.token_tracker:
            agent.add_token_callback(analyzer.token_usage_callback)
        
        # Spustenie analýzy s automatickou otázkou
        query = "Is the Buttoned Cotton Shirt in grey in stock?"
        result = await analyzer.run_analysis(query)
        
        print(f"\n🎉 Analýza dokončená!")
        
    except Exception as e:
        print(f"❌ Chyba v main: {e}")
        traceback.print_exc()
    
    finally:
        try:
            if 'client' in locals():
                await client.close_all_sessions()
            print("✅ Spojenia zatvorené")
        except Exception as e:
            print(f"⚠️ Chyba pri zatváraní: {e}")

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