feat: tool router para LiteLLM (callback + guardrail)
- CustomToolRouter como callback (roda em todas as chamadas) - ToolRouterGuardrail para ativação por request - Modos: jaccard (offline) e embedding (API) - Dockerfile + docker-compose prontos - README com docs completas - Teste offline validando ~74% redução de tokens
This commit is contained in:
CM Consultor
251
test_tool_router.py
Normal file
251
test_tool_router.py
Normal file
@@ -0,0 +1,251 @@
|
||||
"""
|
||||
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())
|
||||
Reference in New Issue
Block a user