From 476f4f954d0f7fa9eb7d59202d456743b47b617c Mon Sep 17 00:00:00 2001 From: CM Consultor Date: Thu, 6 Aug 2026 22:03:29 +0000 Subject: [PATCH] feat: tool router para LiteLLM (callback + guardrail) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .dockerignore | 16 ++ .env.example | 20 ++ .gitignore | 36 +++ Dockerfile | 42 +++ README.md | 292 ++++++++++++++++++++ config.yaml | 76 ++++++ docker-compose.yml | 37 +++ requirements.txt | 13 + test_tool_router.py | 251 ++++++++++++++++++ tool_router_logger.py | 603 ++++++++++++++++++++++++++++++++++++++++++ 10 files changed, 1386 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 config.yaml create mode 100644 docker-compose.yml create mode 100644 requirements.txt create mode 100644 test_tool_router.py create mode 100644 tool_router_logger.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..604b7e1 --- /dev/null +++ b/.dockerignore @@ -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 diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..a144e03 --- /dev/null +++ b/.env.example @@ -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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..012dbd6 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c1cb428 --- /dev/null +++ b/Dockerfile @@ -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"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..4cef612 --- /dev/null +++ b/README.md @@ -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. diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..f875adb --- /dev/null +++ b/config.yaml @@ -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 \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..c6c1196 --- /dev/null +++ b/docker-compose.yml @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..f09b311 --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/test_tool_router.py b/test_tool_router.py new file mode 100644 index 0000000..d10efcb --- /dev/null +++ b/test_tool_router.py @@ -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()) \ No newline at end of file diff --git a/tool_router_logger.py b/tool_router_logger.py new file mode 100644 index 0000000..5da44c4 --- /dev/null +++ b/tool_router_logger.py @@ -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 \ No newline at end of file