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
2. Como funciona
3. Arquivos do projeto
4. Jaccard vs Embedding
5. Instalação
6. Configuração no proxy
7. Parâmetros e env vars
8. Como testar
9. Docker
10. 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.

# 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).

# 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:

export PYTHONPATH="$(pwd):$PYTHONPATH"
litellm --config config.yaml --port 4000

---

Configuração no proxy

Opção A — Callback (recomendado: roda sempre)

# 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)

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:

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)

python test_tool_router.py

Saída esperada: redução de ~70–75% nos chars das tools.

2. Proxy local

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

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

cp .env.example .env
# edite OPENAI_API_KEY e LITELLM_MASTER_KEY

docker compose up --build -d

Proxy em http://localhost:4000.

Só Dockerfile

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.