476f4f954d
- 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
293 lines
8.5 KiB
Markdown
293 lines
8.5 KiB
Markdown
# 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.
|