litellm-tools-logger/test_tool_router.py

"""
Teste rápido do ToolRouter - simula chamadas pra provar a redução de tokens.
"""

import asyncio
import json
import sys
import os

# Adiciona o diretório atual ao path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from tool_router_logger import CustomToolRouter


# Simula tools como se fossem do hermes-agent
TOOLS_MOCK = [
    {
        "type": "function",
        "function": {
            "name": "get_time",
            "description": "Obtém a hora atual em um fuso horário específico",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {
                        "type": "string",
                        "description": "Fuso horário (ex: America/Sao_Paulo)"
                    }
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Pesquisa na web usando um mecanismo de busca. Retorna resultados de páginas web relevantes para a consulta.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Consulta de busca"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Lê o conteúdo de um arquivo no sistema de arquivos. Pode ler arquivos de texto e retornar seu conteúdo.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Caminho do arquivo"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Escreve conteúdo em um arquivo no sistema de arquivos.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Caminho do arquivo"},
                    "content": {"type": "string", "description": "Conteúdo"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_directory",
            "description": "Lista arquivos e diretórios em um determinado caminho.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Caminho"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "execute_bash",
            "description": "Executa um comando bash no terminal.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Comando bash"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Obtém a previsão do tempo para uma cidade.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "Nome da cidade"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "send_email",
            "description": "Envia um email para um destinatário.",
            "parameters": {
                "type": "object",
                "properties": {
                    "to": {"type": "string", "description": "Destinatário"},
                    "subject": {"type": "string", "description": "Assunto"},
                    "body": {"type": "string", "description": "Corpo do email"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Realiza cálculos matemáticos complexos.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "Expressão"}
                }
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_tools",
            "description": "Lista todas as ferramentas disponíveis com suas descrições.",
            "parameters": {"type": "object", "properties": {}}
        }
    },
]


async def test_tool_router():
    """Testa o router com diferentes mensagens."""
    
    router = CustomToolRouter(
        embedding_model="text-embedding-3-small",  # modelo real via API
        top_k=3,
        min_similarity=0.15,
        always_include_tools={"get_time", "help", "list_tools"},
        similarity_mode="jaccard",  # fallback local (sem precisar de API key)
    )
    
    test_cases = [
        ("👋 Oi, tudo bem?", "saudação simples - deve selecionar ~0 tools"),
        ("Que horas são em São Paulo?", "pergunta sobre hora"),
        ("Pesquise sobre IA generativa na web", "pesquisa web"),
        ("Leia o arquivo /tmp/test.txt pra mim", "leitura de arquivo"),
        ("Qual a previsão do tempo em SP?", "clima/weather"),
        ("Execute ls -la no terminal", "comando bash"),
        ("Me mande um email para teste@test.com", "envio de email"),
        ("Quanto é 2 + 2?", "cálculo"),
    ]
    
    print("=" * 80)
    print("🧪 TESTE DO TOOL ROUTER PARA REDUÇÃO DE TOKENS")
    print("=" * 80)
    print(f"\nTools disponíveis: {len(TOOLS_MOCK)}")
    print(f"Always include: {router.always_include_tools}")
    print(f"Modo: {router.similarity_mode}")
    print(f"Top-K: {router.top_k}, Min similarity: {router.min_similarity}")
    print()
    
    total_original = 0
    total_filtrado = 0
    
    for msg, desc in test_cases:
        print(f"\n{'─' * 70}")
        print(f"📝 Teste: '{msg}'")
        print(f"📌 {desc}")
        
        # Simula o data dict que o proxy passa
        data = {
            "messages": [
                {"role": "system", "content": "Você é um assistente útil."},
                {"role": "user", "content": msg},
            ],
            "tools": TOOLS_MOCK,
            "model": "gpt-4o-mini",
        }
        
        tokens_original = len(json.dumps(data["tools"]))
        
        # Executa o hook
        result = await router.async_pre_call_hook(
            user_api_key_dict=None,  # mock
            cache=None,               # mock
            data=data,
            call_type="completion",
        )
        
        if isinstance(result, dict):
            new_tools = result.get("tools", [])
        else:
            new_tools = TOOLS_MOCK
        
        tokens_filtrado = len(json.dumps(new_tools))
        reducao = (1 - tokens_filtrado / tokens_original) * 100 if tokens_original > 0 else 0
        
        selected = [t.get("function", t).get("name", "") for t in new_tools]
        
        total_original += tokens_original
        total_filtrado += tokens_filtrado
        
        print(f"   Tools: {len(TOOLS_MOCK)} → {len(new_tools)}")
        print(f"   Selecionadas: {selected}")
        print(f"   Tokens (chars): {tokens_original} → {tokens_filtrado} (🔥 -{reducao:.1f}%)")
    
    print(f"\n{'=' * 70}")
    reducao_total = (1 - total_filtrado / total_original) * 100
    print(f"📊 TOTAL GERAL: {total_original} → {total_filtrado} chars (🔥 -{reducao_total:.1f}%)")
    print(f"   Economia média por chamada: ~{total_original - total_filtrado} chars")
    print(f"{'=' * 70}")
    print()

    # --- Cenário real: estimativa de tokens LLM ---
    print("\n📈 ESTIMATIVA DE TOKENS REAIS (considerando modelo)")
    print("~4 chars = 1 token em média para tools JSON")
    
    tok_original = total_original // 4
    tok_filtrado = total_filtrado // 4
    economia_tokens = tok_original - tok_filtrado
    
    print(f"  Tokens originais (tools): ~{tok_original}")
    print(f"  Tokens filtrados (tools): ~{tok_filtrado}")
    print(f"  Economia por chamada: ~{economia_tokens} tokens")
    print(f"  Preço @ GPT-4o-mini: ${economia_tokens * 0.000000150:.6f}/chamada (só tools)")
    print(f"  Preço @ GPT-4o: ${economia_tokens * 0.0000025:.6f}/chamada (só tools)")


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