cmconsultor/litellm-tools-logger
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: tool router para LiteLLM (callback + guardrail)

Browse Source 
- 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
2026-08-06 22:03:29 +00:00
commit 476f4f954d
10 changed files with 1386 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
.dockerignore Normal file
View File

@@ -0,0 +1,16 @@
.git
.gitignore
__pycache__
*.pyc
*.pyo
.venv
venv
.env
.env.local
*.log
README.md
.DS_Store
.idea
.vscode
docker-compose.yml
*.md

20
.env.example Normal file
View File

@@ -0,0 +1,20 @@
# Copie para .env e preencha:
# cp .env.example .env
# Obrigatórias para o proxy falar com o provider
OPENAI_API_KEY=sk-...
LITELLM_MASTER_KEY=sk-change-me
# Porta publicada (docker-compose)
PORT=4000
# ── Tool Router ──────────────────────────────────────────────────────────────
TOOL_ROUTER_SIMILARITY_MODE=jaccard
TOOL_ROUTER_TOP_K=3
TOOL_ROUTER_MIN_SIMILARITY=0.15
TOOL_ROUTER_ALWAYS_INCLUDE=get_time,help,list_tools
TOOL_ROUTER_EMBEDDING_MODEL=text-embedding-3-small
# ── Cache Valkey (opcional — só se habilitar no config.yaml) ─────────────────
# VALKEY_HOST=valkey
# VALKEY_PORT=6379

36
.gitignore vendored Normal file
View File

@@ -0,0 +1,36 @@
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
*.egg-info/
.eggs/
dist/
build/
.venv/
venv/
env/
# Env / secrets
.env
.env.local
*.pem
# IDE / OS
.idea/
.vscode/
.DS_Store
Thumbs.db
# Logs / tmp
*.log
/tmp/
.pytest_cache/
.mypy_cache/
.ruff_cache/
.coverage
htmlcov/
# Docker
*.tar

42
Dockerfile Normal file
View File

@@ -0,0 +1,42 @@
# =============================================================================
# Tool Router + LiteLLM Proxy
# =============================================================================
# Build: docker build -t tool-router-litellm .
# Run: docker run --rm -p 4000:4000 --env-file .env tool-router-litellm
# =============================================================================
FROM python:3.12-slim
ENV PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1 \
PYTHONPATH=/app \
PORT=4000
WORKDIR /app
# Dependências de sistema mínimas (certs para HTTPS das APIs LLM)
RUN apt-get update && apt-get install -y --no-install-recommends \
ca-certificates \
curl \
&& rm -rf /var/lib/apt/lists/*
# requirements.txt instala litellm + numpy + scikit-learn
# (litellm NÃO traz numpy/sklearn por padrão)
COPY requirements.txt .
RUN pip install --upgrade pip && pip install -r requirements.txt
# Código do router + configs
COPY tool_router_logger.py .
COPY config.yaml .
COPY test_tool_router.py .
# Porta do proxy
EXPOSE 4000
# Healthcheck simples no endpoint do proxy
HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \
CMD curl -fsS "http://127.0.0.1:${PORT}/health" || exit 1
# Variáveis TOOL_ROUTER_* e OPENAI_API_KEY / LITELLM_MASTER_KEY vêm do runtime
CMD ["sh", "-c", "litellm --config /app/config.yaml --port ${PORT} --host 0.0.0.0"]

292
README.md Normal file
View File

@@ -0,0 +1,292 @@
# Tool Router para LiteLLM
Custom callback / guardrail para o **LiteLLM Proxy** que filtra tools irrelevantes **antes** da chamada ao LLM, reduzindo tokens de input.
Em testes locais com 10 tools: **~74% de redução** de tokens só no payload de tools.
---
## Sumário
1. [O problema](#o-problema)
2. [Como funciona](#como-funciona)
3. [Arquivos do projeto](#arquivos-do-projeto)
4. [Jaccard vs Embedding](#jaccard-vs-embedding)
5. [Instalação](#instalação)
6. [Configuração no proxy](#configuração-no-proxy)
7. [Parâmetros e env vars](#parâmetros-e-env-vars)
8. [Como testar](#como-testar)
9. [Docker](#docker)
10. [FAQ](#faq)
---
## O problema
Agentes com muitas tools enviam **todas** as definições (JSON schemas) em **toda** request. Um simples `"Oi"` pode gastar milhares de tokens só com schemas irrelevantes.
## Como funciona
```
Usuário: "Qual o clima em SP?" + 10 tools
│
▼
┌───────────────────────────────────┐
│ async_pre_call_hook() │
│ 1. Extrai a query do usuário │
│ 2. Similaridade com cada tool │
│ 3. Top-K + always_include │
└──────────┬────────────────────────┘
│
▼
Só 2–3 tools relevantes → LLM
```
Dois modos de uso:
| Modo | Quando roda | Como ativar |
|------|-------------|-------------|
| **Callback** | Todas as chamadas | `callbacks: [tool_router_logger.tool_router_instance]` |
| **Guardrail** | Por request (ou `default_on: true`) | `guardrails: [...]` + body `"guardrails": ["tool_router"]` |
---
## Arquivos do projeto
```
tool-router-litellm/
├── tool_router_logger.py # CustomToolRouter + ToolRouterGuardrail
├── config.yaml # Exemplo de config do proxy
├── test_tool_router.py # Teste offline (jaccard)
├── requirements.txt # Dependências
├── Dockerfile # Imagem do proxy + router
├── docker-compose.yml # Subir com 1 comando
├── .env.example # Variáveis de ambiente
├── .gitignore
└── README.md
```
---
## Jaccard vs Embedding
| | **Jaccard** (default) | **Embedding** |
|---|---|---|
| Dependência de API | Nenhuma | Precisa de modelo de embedding |
| Custo extra | Zero | Tokens de embedding |
| Latência | ~ms | Depende da API |
| Precisão | Boa (sinônimos PT/EN) | Melhor semanticamente |
| Offline | Sim | Não |
Recomendação: comece com **jaccard**. Mude para **embedding** se o filtro errar em queries muito abstratas.
```bash
# Jaccard (default)
export TOOL_ROUTER_SIMILARITY_MODE=jaccard
# Embedding
export TOOL_ROUTER_SIMILARITY_MODE=embedding
export TOOL_ROUTER_EMBEDDING_MODEL=text-embedding-3-small
```
---
## Instalação
### O LiteLLM já traz numpy/scikit-learn?
**Não.** O `litellm` **não** inclui `numpy` nem `scikit-learn` por padrão.
Esses pacotes são **obrigatórios** para o Tool Router (jaccard e cosine similarity).
```bash
# 1. Clone / entre no repo
cd tool-router-litellm
# 2. Instale as deps (litellm + numpy + scikit-learn)
pip install -r requirements.txt
# 3. Configure env
cp .env.example .env
# edite OPENAI_API_KEY, LITELLM_MASTER_KEY, etc.
# 4. Suba o proxy
litellm --config config.yaml --port 4000
```
> O `PYTHONPATH` precisa incluir o diretório deste repo para o proxy importar `tool_router_logger`.
> No Docker isso já é resolvido. Localmente:
```bash
export PYTHONPATH="$(pwd):$PYTHONPATH"
litellm --config config.yaml --port 4000
```
---
## Configuração no proxy
### Opção A — Callback (recomendado: roda sempre)
```yaml
# config.yaml
litellm_settings:
callbacks:
- tool_router_logger.tool_router_instance
```
Parâmetros vêm de **env vars** `TOOL_ROUTER_*` (ver abaixo).
### Opção B — Guardrail (controle por request)
```yaml
litellm_settings:
guardrails:
- guardrail_name: tool_router
litellm_params:
guardrail: tool_router_logger.ToolRouterGuardrail
mode: pre_call
default_on: false # true = roda em todas as calls
top_k: 3
min_similarity: 0.15
similarity_mode: jaccard # ou embedding
always_include_tools:
- get_time
- help
- list_tools
```
Ativar no request:
```bash
curl http://localhost:4000/chat/completions \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
-H "Content-Type: application/json" \
-d '{
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "Qual o clima em SP?"}],
"tools": [ ... ],
"guardrails": ["tool_router"]
}'
```
### Callback vs Guardrail
| | Callback | Guardrail |
|---|---|---|
| Roda sempre | Sim | Só se `default_on: true` ou pedido no body |
| Controle por request | Não | Sim (`"guardrails": ["tool_router"]`) |
| Setup | Mais simples | Mais flexível |
| Ideal para | Reduzir tokens em tudo | Filtrar só em alguns fluxos |
---
## Parâmetros e env vars
| Env var | Default | Descrição |
|---------|---------|-----------|
| `TOOL_ROUTER_SIMILARITY_MODE` | `jaccard` | `jaccard` ou `embedding` |
| `TOOL_ROUTER_TOP_K` | `3` | Quantas tools (além de always_include) |
| `TOOL_ROUTER_MIN_SIMILARITY` | `0.15` | Threshold mínimo |
| `TOOL_ROUTER_ALWAYS_INCLUDE` | `get_time,help,list_tools` | Tools sempre enviadas (CSV) |
| `TOOL_ROUTER_EMBEDDING_MODEL` | `text-embedding-3-small` | Modelo no modo embedding |
No modo **guardrail**, os mesmos campos podem ir em `litellm_params` do `config.yaml`.
---
## Como testar
### 1. Teste offline (sem API key, sem proxy)
```bash
python test_tool_router.py
```
Saída esperada: redução de ~70–75% nos chars das tools.
### 2. Proxy local
```bash
export PYTHONPATH="$(pwd):$PYTHONPATH"
export OPENAI_API_KEY=sk-...
export LITELLM_MASTER_KEY=sk-test-master
export TOOL_ROUTER_SIMILARITY_MODE=jaccard
litellm --config config.yaml --port 4000 --detailed_debug
```
```bash
curl http://localhost:4000/chat/completions \
-H "Authorization: Bearer sk-test-master" \
-H "Content-Type: application/json" \
-d '{
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "Qual a previsão do tempo em São Paulo?"}],
"tools": [
{"type":"function","function":{"name":"get_time","description":"Obtém a hora atual","parameters":{"type":"object","properties":{}}}},
{"type":"function","function":{"name":"get_weather","description":"Obtém a previsão do tempo","parameters":{"type":"object","properties":{}}}},
{"type":"function","function":{"name":"send_email","description":"Envia um email","parameters":{"type":"object","properties":{}}}},
{"type":"function","function":{"name":"execute_bash","description":"Executa comando bash","parameters":{"type":"object","properties":{}}}},
{"type":"function","function":{"name":"search_web","description":"Pesquisa na web","parameters":{"type":"object","properties":{}}}}
]
}'
```
No log com `--detailed_debug`, o `completion()` deve receber só as tools filtradas (ex.: `get_time` + `get_weather`).
---
## Docker
### Build e run
```bash
cp .env.example .env
# edite OPENAI_API_KEY e LITELLM_MASTER_KEY
docker compose up --build -d
```
Proxy em `http://localhost:4000`.
### Só Dockerfile
```bash
docker build -t tool-router-litellm .
docker run --rm -p 4000:4000 \
--env-file .env \
-e TOOL_ROUTER_SIMILARITY_MODE=jaccard \
tool-router-litellm
```
A imagem:
1. Instala `requirements.txt` (**litellm + numpy + scikit-learn**)
2. Copia o código para `/app`
3. Sobe `litellm --config config.yaml --port 4000`
4. Já coloca `/app` no `PYTHONPATH` para importar `tool_router_logger`
---
## FAQ
**Preciso instalar o `requirements.txt` manualmente?**
Sim, se for local. O LiteLLM sozinho **não** traz `numpy`/`scikit-learn`. No Docker o `Dockerfile` já instala tudo.
**Posso usar só como guardrail?**
Sim. Comente `callbacks` e descomente o bloco `guardrails` no `config.yaml`.
**O cache Valkey é obrigatório?**
Não. No `config.yaml` de exemplo ele está configurado; se não tiver Valkey, remova `cache` / `cache_params` ou use o `config` embutido no Docker (sem Valkey por default no compose simples).
**Funciona com function calling / tools do OpenAI?**
Sim. O hook altera `data["tools"]` no pre-call; o resto do fluxo do proxy permanece igual.
**E se o embedding falhar?**
O router loga warning e faz fallback (não derruba a request). Em produção, prefira jaccard se a API de embedding for instável.
---
## Licença
Use como quiser no seu proxy / agente. Sem warranty.

76
config.yaml Normal file
View File

@@ -0,0 +1,76 @@
# ============================================================================
# config.yaml - LiteLLM Proxy com Tool Router + Valkey Semantic Cache
# ============================================================================
# Para rodar:
# litellm --config config.yaml --detailed_debug
# ============================================================================
model_list:
- model_name: gpt-4o-mini
litellm_params:
model: openai/gpt-4o-mini
api_key: os.environ/OPENAI_API_KEY
# Temperatura baixa para ferramentas
temperature: 0.1
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
# Cache semântico com Valkey (reduz tokens em resposta)
cache: true
cache_params:
type: valkey-semantic
host: os.environ/VALKEY_HOST
port: os.environ/VALKEY_PORT
valkey_semantic_cache_embedding_model: "text-embedding-3-small"
similarity_threshold: 0.92 # só retorna cache se >= 92% similar
ttl: 3600 # 1 hora
# ── Tool Router (roda em TODAS as chamadas) ────────────────────────────────
# Formato válido no LiteLLM 1.95+: `callbacks` aponta para uma INSTÂNCIA
# de CustomLogger (`tool_router_instance` é criada no import do módulo).
# Parâmetros configuráveis via env vars: TOOL_ROUTER_* (ver tool_router_logger.py).
callbacks:
- tool_router_logger.tool_router_instance
# ── Opção GUARDRAIL (ativa por request: `guardrails: ["tool_router"]`) ─────
# Descomente se preferir controlar por request em vez de global.
# guardrails:
# - guardrail_name: tool_router
# litellm_params:
# guardrail: tool_router_logger.ToolRouterGuardrail
# mode: pre_call
# default_on: false # true = roda em todas as chamadas
# top_k: 3
# min_similarity: 0.15
# similarity_mode: jaccard # ou "embedding" (exige API de embedding)
# Callbacks para monitoramento
success_callback: ["langfuse"] # opcional: rastreamento
# Desliga cache de requisições (deixa só o semântico) se quiser
# request_model_request_count: false
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
# database_url: os.environ/DATABASE_URL # se quiser persitir usage
# Proxy custom headers
proxy_batch_write_at: 1 # escreve logs em batch
# Segurança
allowed_routes:
- chat/completions
- models
# Rate limiting (opcional)
# alerting: ["slack"]
# alerting_threshold: 300 # 5 minutos sem resposta
# Desabilitar rotas indesejadas
# disable_end_user_cost_tracking: true
# Para ambientes com muitos tools, configurar:
max_request_size_mb: 15

37
docker-compose.yml Normal file
View File

@@ -0,0 +1,37 @@
services:
litellm:
build: .
image: tool-router-litellm:latest
container_name: tool-router-litellm
ports:
- "${PORT:-4000}:4000"
env_file:
- .env
environment:
# Garante import do tool_router_logger
PYTHONPATH: /app
# Defaults do router (sobrescreva no .env se quiser)
TOOL_ROUTER_SIMILARITY_MODE: ${TOOL_ROUTER_SIMILARITY_MODE:-jaccard}
TOOL_ROUTER_TOP_K: ${TOOL_ROUTER_TOP_K:-3}
TOOL_ROUTER_MIN_SIMILARITY: ${TOOL_ROUTER_MIN_SIMILARITY:-0.15}
TOOL_ROUTER_ALWAYS_INCLUDE: ${TOOL_ROUTER_ALWAYS_INCLUDE:-get_time,help,list_tools}
TOOL_ROUTER_EMBEDDING_MODEL: ${TOOL_ROUTER_EMBEDDING_MODEL:-text-embedding-3-small}
volumes:
# Permite editar config sem rebuild
- ./config.yaml:/app/config.yaml:ro
- ./tool_router_logger.py:/app/tool_router_logger.py:ro
restart: unless-stopped
healthcheck:
test: ["CMD", "curl", "-fsS", "http://127.0.0.1:4000/health"]
interval: 30s
timeout: 5s
retries: 3
start_period: 20s
# Opcional: descomente se for usar cache semântico Valkey no config.yaml
# valkey:
# image: valkey/valkey:8-alpine
# container_name: tool-router-valkey
# ports:
# - "6379:6379"
# restart: unless-stopped

13
requirements.txt Normal file
View File

@@ -0,0 +1,13 @@
# Dependências do Tool Router + LiteLLM Proxy
# Instalação: pip install -r requirements.txt
#
# NOTA: litellm NÃO inclui numpy/scikit-learn por padrão.
# Esses pacotes são necessários para o modo jaccard (cosine_similarity)
# e para o modo embedding (vetores).
litellm>=1.95.0
numpy>=1.26.0
scikit-learn>=1.3.0
# Opcional: cache semântico com Valkey/Redis
# redis>=5.0.0

251
test_tool_router.py Normal file
View 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())

603
tool_router_logger.py Normal file
View File

@@ -0,0 +1,603 @@
"""
Tool Router - LiteLLM Custom Logger / Guardrail
===============================================
Intercepta as chamadas do proxy do LiteLLM e faz retrieval das tools,
reduzindo drasticamente o consumo de tokens.
Como funciona:
1. Indexa cada tool (name + description)
2. No pre-call hook, calcula similaridade com a query do usuário
3. Seleciona top-K tools (+ always_include)
4. Só envia essas tools para o LLM
Uso como CALLBACK (roda em todas as chamadas):
litellm_settings:
callbacks:
- tool_router_logger.tool_router_instance
Uso como GUARDRAIL (ativa por request):
litellm_settings:
guardrails:
- guardrail_name: tool_router
litellm_params:
guardrail: tool_router_logger.ToolRouterGuardrail
mode: pre_call
default_on: false
Ou programaticamente:
litellm.callbacks = [CustomToolRouter(...)]
"""
import asyncio
import hashlib
import json
import logging
import os
import re
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching import DualCache
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Embedding local leve (sem dependência externa pesada)
# ---------------------------------------------------------------------------
class FastEmbedder:
"""
Usa o modelo de embedding do LiteLLM (qualquer provider via API)
pra evitar depender de sentence-transformers local.
Alternativa: Cache as embeddings no Valkey para reuse entre chamadas.
"""
def __init__(self, embedding_model: str = "text-embedding-ada-002", cache: Optional[DualCache] = None):
self.embedding_model = embedding_model
self.cache = cache
self._local_cache: Dict[str, List[float]] = {}
async def embed(self, texts: List[str], litellm_kwargs: Optional[Dict] = None) -> List[List[float]]:
"""Gera embeddings em lote usando o LiteLLM."""
import litellm
to_embed = []
indices = []
results = [None] * len(texts)
# Verifica cache primeiro (se disponível)
for i, text in enumerate(texts):
text_hash = hashlib.md5(text.encode()).hexdigest()
cache_key = f"tool_embed:{text_hash}"
if cache_key in self._local_cache:
results[i] = self._local_cache[cache_key]
elif self.cache:
cached = await self.cache.async_get_cache(cache_key)
if cached is not None:
val = cached
if isinstance(val, str):
val = json.loads(val)
self._local_cache[cache_key] = val
results[i] = val
else:
to_embed.append(text)
indices.append(i)
else:
to_embed.append(text)
indices.append(i)
if not to_embed:
return results # type: ignore
# Faz chamada de embedding
kwargs = {"model": self.embedding_model, "input": to_embed}
if litellm_kwargs:
kwargs.update(litellm_kwargs)
try:
response = await litellm.aembedding(**kwargs)
for idx, emb_data in zip(indices, response.data):
embedding = emb_data["embedding"] if isinstance(emb_data, dict) else emb_data.embedding
results[idx] = embedding
# Salva no cache
text = to_embed[indices.index(idx)]
text_hash = hashlib.md5(text.encode()).hexdigest()
cache_key = f"tool_embed:{text_hash}"
self._local_cache[cache_key] = embedding
if self.cache:
await self.cache.async_set_cache(cache_key, embedding)
except Exception as e:
logger.warning(f"Embedding failed, falling back to empty: {e}")
for idx in indices:
results[idx] = [0.0] * 1536 # fallback
return results # type: ignore
# ---------------------------------------------------------------------------
# Similaridade textual inteligente (fallback sem embedding externo)
# ---------------------------------------------------------------------------
import string
from collections import Counter
# Mapa de sinônimos e lematização para português + inglês (tools geralmente
# têm descrições em inglês, queries em português)
_SYNONYM_MAP = {
# verbos para substantivos (português)
"pesquise": "pesquisa", "pesquisar": "pesquisa", "buscar": "pesquisa", "busca": "pesquisa",
"procure": "pesquisa", "procurar": "pesquisa", "encontre": "pesquisa",
"leia": "leitura", "ler": "leitura", "abrir": "leitura", "abra": "leitura",
"escreva": "escrita", "escrever": "escrita", "salvar": "escrita",
"execute": "execução", "executar": "execução", "rodar": "execução", "rode": "execução",
"liste": "listagem", "listar": "listagem", "mostre": "listagem",
"envie": "envio", "enviar": "envio", "mande": "envio",
"calcule": "cálculo", "calcular": "cálculo", "calcular": "cálculo",
# português → inglês (tools descriptions são em inglês)
"hora": "time", "horas": "time", "horário": "time", "horario": "time",
"relógio": "time", "relogio": "time",
"tempo": "weather time", # ambos
"clima": "weather", "previsão": "weather", "previsao": "weather",
"chuva": "weather", "temperatura": "weather", "temperatures": "weather",
"cidade": "city",
"arquivo": "file", "arquivos": "file", "documento": "file",
"pasta": "directory", "diretório": "directory", "diretorio": "directory",
"email": "email", "e-mail": "email", "mensagem": "email",
"web": "web", "internet": "web", "site": "web", "pesquisar": "search",
"bash": "bash", "terminal": "bash", "comando": "bash command", "shell": "bash",
"linha": "line",
"navegador": "browser",
"sistema": "system", "arquivos": "files",
}
# Stopwords comuns em português e inglês (palavras muito curtas/frequentes)
_STOPWORDS = {
"a", "e", "o", "em", "de", "da", "do", "na", "no", "um", "uma",
"é", "ser", "tem", "com", "para", "por", "pra", "os", "as",
"ao", "aos", "das", "dos", "nas", "nos", "que", "se", "mas",
"mais", "muito", "como", "vai", "vou", "me", "te", "se",
"the", "a", "an", "in", "on", "at", "to", "for", "of",
"is", "it", "be", "with", "and", "or", "from", "by",
"um", "uns", "uma", "umas",
}
def _tokenize(text: str) -> list:
"""Tokeniza e normaliza, retornando lista (não set) pra preservar frequência."""
text = text.lower().translate(str.maketrans('', '', string.punctuation))
tokens = re.findall(r'[a-záéíóúãõâêîôûç]+', text)
# Remove stopwords
tokens = [t for t in tokens if t not in _STOPWORDS and len(t) > 1]
# Aplica sinônimos
expanded = []
for t in tokens:
expanded.append(t)
if t in _SYNONYM_MAP:
expanded.extend(_SYNONYM_MAP[t].split())
return expanded
def _jaccard_similarity(a: set, b: set) -> float:
if not a or not b:
return 0.0
return len(a & b) / len(a | b)
def _tf_keyword_score(query_tokens: list, tool_text: str) -> float:
"""
Score baseado em quantas palavras-chave da query aparecem no texto da tool.
Penaliza tools que não contêm NENHUMA palavra relevante.
"""
tool_lower = tool_text.lower()
matches = sum(1 for q in query_tokens if q in tool_lower)
if matches == 0:
return 0.0
# Normaliza pelo tamanho da query pra favorecer matches parciais
return matches / max(len(set(query_tokens)), 1)
def _combined_text_similarity(query: str, tool_text: str) -> float:
"""
Combina Jaccard (set overlap) + TF-keyword score pra dar
similaridade mais robusta sem precisar de embeddings externos.
"""
query_tokens = _tokenize(query)
tool_tokens = _tokenize(tool_text)
query_set = set(query_tokens)
tool_set = set(tool_tokens)
# Jaccard entre sets
jac = _jaccard_similarity(query_set, tool_set)
# TF-score: quantas palavras da query estão no texto da tool
tf = _tf_keyword_score(query_tokens, tool_text)
# Identifica match exato no nome da tool (peso máximo)
query_words = set(w for w in query_tokens if len(w) > 2)
tool_name = tool_text.split()[0] if tool_text else ""
name_match = 1.0 if tool_name in query_words else 0.0
# Se o nome da tool apareceu na query, é quase certeza
if name_match > 0:
return 1.0
# Peso maior pro TF quando Jaccard é baixo (captura relações indiretas)
if jac == 0 and tf > 0:
return tf * 0.4 # meio caminho, mais conservador
# Só retorna algo relevante se passou de um mínimo
raw = max(jac, tf * 0.6)
return raw
# ---------------------------------------------------------------------------
# CustomLogger principal
# ---------------------------------------------------------------------------
DEFAULT_TOOL_NAMES = {"get_time", "help", "ping", "status"}
class CustomToolRouter(CustomLogger):
"""
Router semântico de tools para LiteLLM Proxy.
No pre-call hook, faz retrieval das tools mais relevantes
para a query do usuário, reduzindo o prompt de tools enviado ao LLM.
Args:
embedding_model: Modelo de embedding (ex: 'text-embedding-ada-002')
top_k: Número de tools a incluir (default: 3)
min_similarity: Similaridade mínima para incluir tool (default: 0.15)
always_include_tools: Lista de nomes de tools sempre enviadas
similarity_mode: "embedding" (preferido) ou "jaccard" (fallback leve)
cache_host: Host do Valkey/Redis (opcional)
cache_port: Porta do Valkey/Redis (opcional)
cache_password: Senha (opcional)
"""
def __init__(
self,
embedding_model: str = "text-embedding-ada-002",
top_k: int = 3,
min_similarity: float = 0.15,
always_include_tools: Optional[list] = None,
similarity_mode: str = "embedding",
cache_host: Optional[str] = None,
cache_port: Optional[str] = None,
cache_password: Optional[str] = None,
):
super().__init__()
self.top_k = top_k
self.min_similarity = min_similarity
self.always_include_tools = set(always_include_tools or DEFAULT_TOOL_NAMES)
self.similarity_mode = similarity_mode
self.embedding_model = embedding_model
# Cache compartilhado (inicializado no async_pre_call_hook)
self._dual_cache: Optional[DualCache] = None
self._embedder: Optional[FastEmbedder] = None
# Cache local de embeddings de tools
self._tool_embeddings: Dict[str, List[float]] = {}
self._tool_texts: Dict[str, str] = {}
logger.info(
f"CustomToolRouter initialized: "
f"top_k={top_k}, min_similarity={min_similarity}, "
f"mode={similarity_mode}, model={embedding_model}, "
f"always_include={self.always_include_tools}"
)
def _extract_user_message(self, messages: List[Dict]) -> str:
"""Extrai a última mensagem do usuário."""
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content", "")
if role == "user":
if isinstance(content, list):
# Pode ser multimodal - extrai texto
texts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
texts.append(part.get("text", ""))
return " ".join(texts)
return str(content)
return ""
def _build_tool_text(self, tool: Dict) -> str:
"""Cria um texto combinado para a tool (nome + descrição + parâmetros)."""
func = tool.get("function", tool)
name = func.get("name", "")
description = func.get("description", "")
parameters = func.get("parameters", {})
# Concatena informações relevantes
parts = [name, description]
# Adiciona nomes dos parâmetros
if isinstance(parameters, dict):
props = parameters.get("properties", {})
if isinstance(props, dict):
param_names = " ".join(props.keys())
if param_names:
parts.append(param_names)
return " ".join(parts).lower()
async def _build_tool_index(self, tools: List[Dict]):
"""Indexa as tools: gera/cacheia embeddings."""
texts = {}
for tool in tools:
func = tool.get("function", tool)
name = func.get("name", "")
text = self._build_tool_text(tool)
texts[name] = text
self._tool_texts[name] = text
# Gera embeddings para tools ainda não indexadas
if self.similarity_mode == "embedding" and self._embedder:
new_tools = [name for name in texts if name not in self._tool_embeddings]
if new_tools:
tool_texts_list = [texts[name] for name in new_tools]
try:
embeddings = await self._embedder.embed(tool_texts_list)
for name, emb in zip(new_tools, embeddings):
self._tool_embeddings[name] = emb
except Exception as e:
logger.warning(f"Tool indexing failed, using jaccard: {e}")
self.similarity_mode = "jaccard"
async def _compute_similarities(
self, query: str, tool_names: List[str]
) -> List[Tuple[str, float]]:
"""Calcula similaridade entre query e cada tool."""
if not query.strip():
return [(name, 0.0) for name in tool_names]
if self.similarity_mode == "embedding" and self._embedder:
try:
query_emb = (await self._embedder.embed([query]))[0]
similarities = []
for name in tool_names:
tool_emb = self._tool_embeddings.get(name)
if tool_emb:
sim = cosine_similarity([query_emb], [tool_emb])[0][0]
similarities.append((name, float(sim)))
else:
similarities.append((name, 0.0))
return similarities
except Exception as e:
logger.warning(f"Embedding similarity failed: {e}")
# fallback pra jaccard
# Similaridade combinada (fallback inteligente)
similarities = []
for name in tool_names:
tool_text = self._tool_texts.get(name, name)
sim = _combined_text_similarity(query, tool_text)
similarities.append((name, sim))
return similarities
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
) -> Optional[Union[Exception, str, dict]]:
"""
Hook principal: intercepta a requisição e filtra tools.
Retorna:
- dict modificado → segue com as tools filtradas
- None → continua sem modificação
- Exception → rejeita a request
"""
# Só processa completion
if call_type not in ("completion", "acompletion"):
return None
messages = data.get("messages", [])
tools = data.get("tools", [])
# Se não tem tools ou só 1, não precisa fazer nada
if not tools or len(tools) <= 1:
return None
# Inicializa cache e embedder
self._dual_cache = cache
if self._embedder is None:
self._embedder = FastEmbedder(
embedding_model=self.embedding_model,
cache=cache,
)
# Extrai query do usuário
user_query = self._extract_user_message(messages)
# Indexa as tools (se necessário)
await self._build_tool_index(tools)
# Nomes de todas as tools
all_tool_names = []
for tool in tools:
func = tool.get("function", tool)
all_tool_names.append(func.get("name", ""))
# Calcula similaridades
similarities = await self._compute_similarities(user_query, all_tool_names)
# Filtra: sempre inclui + top-k por similaridade
selected_names = set()
# 1. Tools always-include
for name in self.always_include_tools:
if name in all_tool_names:
selected_names.add(name)
# 2. Top-k por similaridade (acima do threshold)
sorted_tools = sorted(similarities, key=lambda x: x[1], reverse=True)
for name, sim in sorted_tools:
if len(selected_names) >= self.top_k:
break
if sim >= self.min_similarity:
selected_names.add(name)
# Se não selecionou nada, pelo menos a top-1
if not selected_names and sorted_tools:
selected_names.add(sorted_tools[0][0])
# Monta novo tools array
new_tools = [t for t in tools if (
t.get("function", t).get("name", "") in selected_names
)]
# Log do que aconteceu
logger.info(
f"ToolRouter: tools {len(tools)}→{len(new_tools)} "
f"(query: '{user_query[:60]}...' ) "
f"selected: {sorted(selected_names)} "
f"similarities: {dict(sorted_tools)}"
)
# Retorna dados modificados
if len(new_tools) < len(tools):
data["tools"] = new_tools
# Se reduziu drasticamente, também podemos avisar o modelo
if len(new_tools) <= 1 and messages:
system_msg = (
"Você tem acesso apenas a tools selecionadas para esta conversa. "
"Se precisar de uma tool específica que não está disponível, "
"explique qual precisa e o usuário poderá liberá-la."
)
# Adiciona system message sutil, sem poluir
existing_system = False
for msg in messages:
if msg.get("role") == "system":
existing_system = True
if system_msg not in msg.get("content", ""):
msg["content"] = msg.get("content", "") + "\n\n" + system_msg
break
if not existing_system:
messages.insert(0, {"role": "system", "content": system_msg})
data["messages"] = messages
return data
async def async_log_success_event(self, *args, **kwargs):
"""Apenas para logging/monitoramento."""
pass
# ---------------------------------------------------------------------------
# Factory para criação do logger via string path
# ---------------------------------------------------------------------------
# =============================================================================
# Instância pronta + factory + wrapper de guardrail
# =============================================================================
#
# Uso como CALLBACK (recomendado — roda em todas as chamadas):
# litellm_settings:
# callbacks: ["tool_router_logger.tool_router_instance"]
#
# Uso como GUARDRAIL (ativa por request com `guardrails: ["tool_router"]`):
# litellm_settings:
# guardrails:
# - guardrail_name: tool_router
# litellm_params:
# guardrail: tool_router_logger.ToolRouterGuardrail
# mode: pre_call
# default_on: false
# top_k: 3
# min_similarity: 0.15
# similarity_mode: jaccard
#
# Parâmetros configuráveis via env vars:
# TOOL_ROUTER_EMBEDDING_MODEL, TOOL_ROUTER_TOP_K, TOOL_ROUTER_MIN_SIMILARITY,
# TOOL_ROUTER_ALWAYS_INCLUDE (csv), TOOL_ROUTER_SIMILARITY_MODE
# =============================================================================
def _router_params_from_env() -> dict:
"""Lê os parâmetros do router a partir de env vars (com defaults sãos)."""
return {
"embedding_model": os.environ.get("TOOL_ROUTER_EMBEDDING_MODEL", "text-embedding-3-small"),
"top_k": int(os.environ.get("TOOL_ROUTER_TOP_K", "3")),
"min_similarity": float(os.environ.get("TOOL_ROUTER_MIN_SIMILARITY", "0.15")),
"always_include_tools": set(
filter(None, os.environ.get("TOOL_ROUTER_ALWAYS_INCLUDE", "get_time,help,list_tools").split(","))
),
"similarity_mode": os.environ.get("TOOL_ROUTER_SIMILARITY_MODE", "jaccard"),
}
# Instância global reutilizável — é o que o proxy carrega via:
# callbacks: ["tool_router_logger.tool_router_instance"]
tool_router_instance: CustomToolRouter = CustomToolRouter(**_router_params_from_env())
def get_custom_tool_router(**kwargs) -> CustomToolRouter:
"""
Factory function para uso programático.
Se nenhum kwarg for passado, retorna a instância global
(equivalente a `tool_router_instance`). Se kwargs forem passados,
cria uma instância dedicada com aqueles parâmetros.
"""
global tool_router_instance
if kwargs:
return CustomToolRouter(**kwargs)
return tool_router_instance
class ToolRouterGuardrail(CustomGuardrail):
"""
Wrapper para usar o ToolRouter como guardrail do LiteLLM Proxy.
Guardrails permitem ativar o roteamento de tools por request:
curl ... -d '{"model": "gpt-4o-mini", "guardrails": ["tool_router"], ...}'
Ou globalmente com `default_on: true` no config.yaml.
"""
_ROUTER_PARAMS = (
"embedding_model", "top_k", "min_similarity",
"always_include_tools", "similarity_mode",
"cache_host", "cache_port", "cache_password",
)
def __init__(self, guardrail_name: str = "tool_router", **kwargs):
# `guardrail_name`, `event_hook`, `default_on` e demais campos do
# guardrail são repassados à base; os parâmetros do router são
# extraídos para criar a instância interna.
super().__init__(guardrail_name=guardrail_name, **kwargs)
router_params = {k: v for k, v in kwargs.items() if k in self._ROUTER_PARAMS}
self._router = CustomToolRouter(**router_params)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
) -> Optional[Union[Exception, str, dict]]:
"""Delega o filtro de tools ao CustomToolRouter."""
return await self._router.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type=call_type,
)
async def async_log_success_event(self, *args, **kwargs):
pass