Shortwave: assistente de e-mail com RAG

01 · CONTEXTO

Um assistente de e-mail com RAG

A Shortwave, um aplicativo de e-mail, construiu um assistente de IA com geração aumentada por recuperação (RAG). A caixa de entrada de cada usuário reúne correspondência, convites de calendário, recibos, notificações e newsletters. A empresa queria uma base de conhecimento a partir dela, mas as tentativas anteriores com busca tradicional tinham sido decepcionantes, e os LLMs podiam mudar isso.

Para os usuários, a empresa descreve a IA como uma pessoa sentada ao lado deles, com conhecimento perfeito do histórico de e-mails. Esse modelo mental simples ajuda a prever o que o assistente consegue fazer e como usá-lo melhor.

02 · ARQUITETURA

Um prompt grande, quatro etapas

O princípio central é que como responder a uma pergunta deve ser tratado pelo próprio LLM; o trabalho do sistema é encontrar os dados certos. Segundo a equipe, cadeias longas de chamadas de LLM introduziram perda de dados e erros em cada estágio, com respostas de baixa qualidade. Em vez delas, a resposta sai de um único LLM que recebe todo o contexto num prompt grande, o que depende de um modelo muito capaz e de uma janela de contexto grande.

Para organizar isso, a Shortwave criou a abstração de ferramenta (Tool). Cada tipo de dado, como calendário, histórico de e-mail ou rascunho, vem de uma ferramenta diferente, e cada ferramenta decide internamente quais dados específicos entram no prompt final.

Ao responder, o assistente segue quatro etapas: seleção de ferramentas, recuperação dos dados das ferramentas, resposta à pergunta e pós-processamento.

1. A pergunta do usuário vai imediatamente a um LLM. Essa chamada não serve para responder, e sim para determinar os tipos de dados que precisam ser reunidos.

2. A seleção pode devolver nenhuma, uma ou várias ferramentas, e o prompt dessa etapa leva detalhes como o estado atual do app.

3. As ferramentas escolhidas recuperam seus dados em paralelo. A mais complexa é a EmailHistory, que dá ao assistente conhecimento de todo o histórico de e-mail.

4. Com os dados reunidos, o sistema monta um prompt com a pergunta original e o contexto trazido pelas ferramentas e o entrega a um LLM.

5. Aqui é preciso decidir como alocar tokens entre as ferramentas, porque a EmailHistory costuma produzir bem mais conteúdo do que cabe no prompt.

6. Por fim, o pós-processamento converte a saída em texto rico, adiciona citações das fontes e sugere ações ao usuário.

03 · ORÇAMENTO DE CONTEXTO

Mais dados do que cabem no prompt

O AI Search procura e-mails relevantes em todo o histórico do usuário e os ordena por utilidade, para que trocas inteligentes possam ser feitas na hora de responder. Muitas vezes os resultados chegam a mil ou mais, e é preciso reduzi-los a algumas dezenas, e rápido.

Na etapa de resposta isso vira uma questão de orçamento de contexto: é preciso decidir como alocar tokens entre as ferramentas, porque a EmailHistory costuma produzir bem mais conteúdo do que cabe no prompt. O post fala em heurísticas para fazer essas trocas.

04 · AI SEARCH

Da pergunta ao ranking de e-mails

A primeira etapa é a reformulação da consulta. Ela reescreve com um LLM uma consulta que não traz o contexto necessário, para que faça sentido sozinha; o post dá o exemplo de uma pergunta de continuação numa conversa, que não faz sentido como pergunta isolada. Para isso, considera o que está visível na tela a que o usuário pode estar se referindo, incluindo o histórico do chat, a conversa visível e o rascunho atual.

Depois, features extraídas da consulta alimentam uma busca tradicional e uma busca vetorial. Na ingestão, cada e-mail vira um embedding por um modelo aberto, calculado em servidores próprios com aceleração por GPU, e fica num banco vetorial com espaço separado por usuário.

1. A primeira etapa do AI Search é reformular a consulta com um LLM, para que ela faça sentido sozinha.

2. Um conjunto de LLMs rápidos extrai features da consulta reformulada, cada uma com escore de confiança. Com elas, a busca tradicional por texto completo ou metadados pode ser o melhor caminho.

3. A consulta reformulada é usada na busca por e-mails semanticamente similares no banco vetorial, com espaço separado por usuário. Features como o intervalo de datas restringem a busca a um subconjunto menor.

4. A primeira fase de re-ranking usa heurísticas rápidas para cortar a lista grande, ordenando pela similaridade semântica e aplicando regras sobre as features, como um filtro gaussiano nas datas.

5. Na segunda fase, um cross-encoder aberto em GPUs próprias recebe a consulta reformulada e os fragmentos pré-selecionados. Ele é mais esperto que as heurísticas, porém lento, e por isso só processa os fragmentos mais bem ranqueados.

6. As heurísticas são aplicadas de novo aos escores do cross-encoder, e a lista final ordenada de fragmentos volta para o assistente.

05 · TRADE-OFFS

Duas fases de re-ranking e o custo total

A equipe otimizou o sistema para responder quase todas as perguntas em 3-5 segundos, e essa latência ponta a ponta inclui várias chamadas de LLM, consultas ao ElasticSearch e chamadas a modelos abertos. O sistema usa concorrência, streaming e pipelining e muito poder de computação: servidores robustos e clusters de GPUs de ponta.