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
16
.dockerignore
Normal file
16
.dockerignore
Normal 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
20
.env.example
Normal 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
36
.gitignore
vendored
Normal 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
42
Dockerfile
Normal 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
292
README.md
Normal 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
76
config.yaml
Normal 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
37
docker-compose.yml
Normal 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
13
requirements.txt
Normal 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
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())
|
||||
603
tool_router_logger.py
Normal file
603
tool_router_logger.py
Normal 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
|
||||
Reference in New Issue
Block a user