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