JetBrains detalha pipeline RAG de busca semântica de código em produção

O problema: agentes precisam de busca por significado

A JetBrains detalhou, em formato de diário de desenvolvimento, o pipeline RAG de busca semântica de código que originou o Air Context, produto hoje em prévia pública incluído nas licenças da empresa. A motivação: agentes de codificação gastam muito tempo localizando trechos relevantes em bases grandes usando ferramentas tradicionais como grep e busca por palavra-chave, que exigem saber de antemão o texto exato - um agente procurando onde tokens de sessão são renovados não pode contar que o código contenha a palavra "refresh". A proposta é indexar o código capturando sua semântica e permitir busca por linguagem natural.

Parsing e chunking com consciência de estrutura

O primeiro estágio do pipeline é dividir arquivos-fonte em unidades com escopo adequado. Chunking por tamanho fixo se mostrou semanticamente errado: agrupa trechos sem relação, como importações e conteúdo de funções. A JetBrains usou seus parsers maduros (fruto de 26 anos de desenvolvimento) no plataforma interna Code Engine para fazer chunking orientado à estrutura em nove linguagens: Kotlin, Java, Python, JavaScript, TypeScript, C#, PHP, Go e Rust; outras linguagens recebem fallback para divisão por linhas.

O algoritmo decide o escopo de cada chunk a partir do tipo e tamanho dos nós de sintaxe, mantém prefixos (documentação, anotações, modificadores) junto às declarações e preserva sufixos de fechamento com seus construtos. Anotações semanticamente vazias como @NotNull e @Override são removidas. A abordagem se assemelha ao cAST (Zhang et al., 2025), mas incorpora mais semântica de linguagem, como decorators de Python junto às definições. Chunks passam por normalização (trim de espaços, remoção de linhas em branco, ajuste de indentação) e são avaliados com uma estratégia de LLM-as-a-judge, que verifica se as fronteiras fazem sentido; mudanças no pipeline também passam por avaliação end-to-end de recuperação.

Vetorização: dimensões valem mais que precisão

Na vetorização, um modelo de embedding transforma cada chunk em um vetor de milhares de dimensões em que textos de significado parecido ficam próximos. Como um repositório grande gera milhões de chunks, o armazenamento vira o limite de custo - um vetor em floats de 32 bits pesa cerca de 16 KB, somando dezenas de gigabytes.

A equipe comparou duas alavancas de compressão: reduzir dimensões ou reduzir precisão. Com o mesmo orçamento (512 bytes por vetor), 128 dimensões em 32 bits perde claramente para 4.096 dimensões em 1 bit cada. A justificativa: cada dimensão funciona como uma "pergunta" aprendida sobre o texto; manter todas as dimensões com um bit preserva uma resposta aproximada para todas, enquanto truncar descarta informações irrecuperáveis.

A quantização binária é simples: componentes não-negativos viram 1, negativos viram 0. A métrica muda de similaridade de cosseno para distância de Hamming, computável em cerca de uma centena de instruções de CPU por comparação de 4.096 bits, contra milhares de multiplicações no caso dos floats. A perda de alguns pontos de recall foi aceita porque o consumidor é um agente de raciocínio, que abre e lê os candidatos de qualquer forma - importa que os arquivos relevantes apareçam entre os primeiros resultados, não a ordem exata.

Há um custo sutil: a quantização comprime a faixa de pontuações de similaridade. O ranking sobrevive, mas thresholding não - um painel que sugere código proativamente precisa de um vão utilizável entre "relacionado" e "não relacionado", que vetores binários não deixam. Nesses casos de julgamento absoluto de relevância, a equipe mantém floats de 16 bits.

Escopo de embedding e privacidade

Indexação e busca usam o mesmo modelo, mas com otimizações opostas: throughput na indexação (batchs de ~32 chunks na GPU) e tempo até o primeiro resultado na busca. É usado um modelo instruction-following com assimetria deliberada: documentos são embutidos como estão; consultas recebem uma instrução descrevendo a tarefa de recuperação.

Cada chunk é embutido junto com o caminho do arquivo, que fornece metadados úteis - mas no monorepo do IntelliJ IDEA, com mais de um milhão de arquivos e caminhos de até 218 caracteres, o caminho pode ocupar mais espaço que o próprio código. A solução: o caminho é abreviado preservando as duas pontas (segmentos iniciais do módulo e as dois últimos níveis, pai e nome do arquivo), eliminando o meio. Buscas com escopo em subdiretório inserem o escopo no texto da consulta na mesma forma abreviada usada na indexação, alinhando os dois lados do espaço vetorial.

Sobre privacidade, o Air Context não armazena código: os chunks guardam apenas referências, caminhos, offsets e referência a vetores - o servidor sabe onde algo relevante está, não o que é, e o trecho exibido é montado na máquina do cliente. Os embeddings usam um modelo open-weights rodando em GPUs da própria JetBrains, sem requisições a fornecedores externos, eliminando o risco de uso dos dados em treinamento. Em benchmarks internos de recuperação de código, os candidatos open-weights superaram as APIs de embedding hospedadas das grandes provedoras; segundo a equipe, o que resta de engenharia interessante está no que se alimenta aos modelos e como se os serve.