backend/docs/ai/rag-pipeline.md
Rodrigo Verdiani 90ee74b3b9 docs: add AI module documentation and AGENTS.md
Comprehensive documentation for the AI module covering architecture,
API design, data model, RAG pipeline, configuration, integration points,
and frontend integration guide. Includes IDA5 extraction points reference
and AGENTS.md with project conventions and coding standards.
2026-06-08 20:13:21 -03:00

16 KiB
Raw Blame History

Pipeline RAG (Retrieval-Augmented Generation)

Visão Geral

O pipeline RAG do Decompile-AI permite que o usuário faça perguntas em linguagem natural sobre um binário analisado. O sistema recupera automaticamente as funções e informações mais relevantes do binário e as injeta no prompt da LLM, que então gera uma resposta contextualizada.

O pipeline opera em duas fases distintas:

  1. Fase de Indexação — Processa os resultados da análise estática, gera embeddings e os armazena no pgvector
  2. Fase de Query — A cada pergunta do usuário, recupera contexto relevante e gera a resposta via LLM

Fase 1: Indexação

1.1 Gatilho

A indexação é disparada automaticamente quando uma análise estática é concluída com sucesso:

StaticAnalysisHandler.process() conclui
    │
    └── publica StaticAnalysisCompletedEvent(binaryId, projectId, workspaceId, analysisId, engine)
            │
            └── EmbeddingEventSubscriber.on(StaticAnalysisCompletedEvent)
                    │
                    └── cria Job(type=GENERATE_EMBEDDINGS, status=ENQUEUED)
                            │
                            └── RabbitMQ → EmbeddingGenerationHandler.process(job)

Também é possível reindexar manualmente (endpoint futuro).

1.2 Coleta de Dados

O EmbeddingGenerationHandler chama EmbeddingService.indexBinary(binaryId). O serviço:

  1. Deleta embeddings existentes para o binário (deleteByBinaryId) — hard delete, sem soft delete
  2. Busca a análise mais recente via AnalysisService.getAnalysis(analysisId) ou listAnalyses(binaryId)
  3. Itera sobre todas as StaticFunction da análise, carregando:
    • name, address, assembly (sempre presentes)
    • decompiledCode, signature (quando disponíveis — Ghidra/IDA 7)
    • StaticLabel associadas à função
    • StaticXref (callers e callees) — nomes das funções via join
  4. Busca metadados do binário: formato, arquitetura, compilador (tabela binaries)

1.3 Chunking

Cada função gera um chunk de texto que será embeddado. O template do chunk:

[FUNCTION] {name} at {address}
Binary: {filename} | {architecture} | {compiler} | {format}
Callers: {caller1}, {caller2}, ...
Callees: {callee1}, {callee2}, ...
Labels:
  {label1_address} ({label1_name})
  {label2_address} ({label2_name})

[ASSEMBLY]
{assembly_code}

Regras de chunking:

Regra Descrição
1 chunk = 1 função Unidade semântica natural da análise de binários
Truncagem Se o texto montado exceder chunk-max-chars (default 3000), o assembly é truncado
Labels inline Labels são incluídos como comentários contextuais para enriquecer o embedding
Xrefs como metadados Nomes de callers/callees são incluídos como "dicas semânticas" no texto
Chunk de metadados Um chunk extra (tipo METADATA) com informações gerais do binário é sempre gerado

Exemplo de chunk gerado (para malware.exe x86 com MSVC):

[FUNCTION] sub_401000 at 0x401000
Binary: malware.exe | x86 | MSVC 19.0 | PE32
Callers: WinMain, sub_402300
Callees: CreateFileA, WriteFile, CloseHandle, sub_401200
Labels:
  0x40100A (check_password)
  0x401050 (error_msg)

[ASSEMBLY]
push    ebp
mov     ebp, esp
sub     esp, 0x40
push    ecx
push    edx
mov     eax, [ebp+8]
push    eax
call    ds:CreateFileA
mov     [ebp-4], eax
cmp     eax, -1
jz      short loc_401050
...

Chunk de metadados (tipo METADATA):

[METADATA] Binary Overview
Filename: malware.exe
Format: PE32 (Portable Executable 32-bit)
Architecture: x86 (Intel 32-bit)
Compiler: Microsoft Visual C++ 19.0
Entry Point: 0x401200
Function count: 342
Analysis engine: IDA5

Esse chunk de metadados responde perguntas como "qual arquitetura?", "com o que foi compilado?", "quantas funções tem?".

1.4 Geração de Embeddings

Para cada chunk de texto:

  1. O texto é enviado ao Ollama via Spring AI EmbeddingModel.embed(document)
  2. O modelo configurado é qwen3-embedding:latest (8B parâmetros, dimensão configurável)
  3. A dimensão padrão dos vetores é 1024 (configurável em app.ai.embedding.dimension)
  4. Os embeddings são gerados em batch de até batch-size (default 20) chunks por chamada
  5. Cada embedding é um float[1024] (ou dimensão configurada)
  6. O vetor é persistido na tabela embedding_chunk junto com o texto original e metadados

Fluxo detalhado do batch:

funções: [f1, f2, f3, ..., f342]
    │
    │  particiona em batches de 20
    ▼
batch1: [f1..f20]  → Ollama embed() → 20 × float[1024] → INSERT batch
batch2: [f21..f40] → Ollama embed() → 20 × float[1024] → INSERT batch
...
batch18: [f341..f342] → Ollama embed() → 2 × float[1024] → INSERT batch
    │
    ▼
chunk METADATA → Ollama embed() → 1 × float[1024] → INSERT

Tratamento de erros:

  • Se a API Ollama falhar para um batch, o job é marcado como FAILED
  • O retry_count do job é incrementado (política de retry via RabbitMQ dead-letter queue no futuro)
  • Nenhum embedding parcial é salvo — a transação faz rollback

1.5 Re-indexação

Quando um binário é re-analisado (ex: trocar de engine IDA5 → Ghidra):

  1. O EmbeddingEventSubscriber cria novo job GENERATE_EMBEDDINGS
  2. EmbeddingService.indexBinary() chama deleteByBinaryId() primeiro → hard delete de todos os chunks antigos
  3. Novos chunks são gerados a partir da análise mais recente

Isso garante que não há embeddings órfãos ou inconsistentes.


Fase 2: Query (Chat)

2.1 Visão Geral do Fluxo

Usuário: "Quais funções nesse binário lidam com rede?"
    │
    ▼
ChatController → ChatService.chat(sessionId, message)
    │
    ├── 1. Busca sessão + histórico recente (últimas N mensagens)
    │
    ├── 2. ChatService.buildContext(binaryId, query, binary)
    │       │
    │       ├── 2a. SemanticSearch: embed query → pgvector <=>
    │       │        └── top-K chunks por similaridade de cosseno
    │       │
    │       ├── 2b. TextualFallback (se semantic < topK * 0.4):
    │       │        └── Extrai tokens da query → LIKE search no content
    │       │
    │       └── 2c. CallGraphExpansion:
    │                ├── Detecta nomes de função/endereço na query
    │                ├── Busca funções via static_function (ILIKE)
    │                └── Expande 1-hop (callers + callees) + assembly inline
    │
    ├── 3. Merge + dedup: requested functions first, then semantic/textual chunks
    │
    ├── 4. Build system prompt com contexto RAG + histórico
    │
    ├── 5. LLM (Gemma4:12b) → streaming
    │       │
    │       └── Flux<String> via StreamingChatModel.stream()
    │
    ├── 6. Salva UserMessage + AssistantMessage no banco
    │
    └── 7. Retorna SSE stream para o frontend

2.2 Busca Semântica (pgvector)

A busca vetorial é a espinha dorsal do retrieval. Utiliza o operador de distância de cosseno do pgvector:

SELECT id, content, metadata, token_count,
       1 - (embedding <=> CAST(:queryEmbedding AS vector)) AS similarity
FROM embedding_chunk
WHERE binary_id = :binaryId
  AND 1 - (embedding <=> CAST(:queryEmbedding AS vector)) > :threshold
ORDER BY embedding <=> CAST(:queryEmbedding AS vector)
LIMIT :limit
  • <=> é o operador de cosine distance do pgvector (0 = idêntico, 2 = oposto)
  • 1 - (<=>) converte para similaridade (1 = idêntico, -1 = oposto)
  • threshold (default 0.6) filtra resultados pouco relevantes
  • limit (default 10) retorna os top-K chunks
  • binary_id restringe a busca ao binário específico (escopo da sessão)

O índice HNSW (idx_embedding_chunk_embedding) acelera a busca para complexidade O(log N).

Modelo de embedding da query:

A pergunta do usuário é embeddada usando o mesmo modelo da indexação (qwen3-embedding:latest), garantindo que query e documentos estejam no mesmo espaço vetorial.

2.2.1 Fallback por Busca Textual

Quando a busca semântica retorna poucos resultados (menos de fallback-min-ratio * topK, default 0.4 * 10 = 4), o sistema complementa com busca textual direta via LIKE no conteúdo dos chunks:

SELECT * FROM embedding_chunk
WHERE binary_id = :binaryId
  AND LOWER(content) LIKE LOWER(CONCAT('%', :term, '%'))
LIMIT :limit

Estratégia:

  1. A query é tokenizada (split por espaços e pontuação)
  2. Tokens com ≥ 3 caracteres são usados como termos de busca
  3. Para cada termo, busca textual preenche os slots restantes até topK
  4. Resultados são deduplicados (mesmo chunk pode aparecer em múltiplas buscas)

Exemplo: Query "Quais funções lidam com rede?"

  • Tokens: Quais, funções, lidam, rede
  • Cada token vira LIKE '%rede%' no conteúdo dos chunks
  • Funções que mencionam WS2_32.send, socket, connect são recuperadas mesmo que o embedding semântico não as tenha rankeado bem

2.3 Expansão por Call Graph (Recuperação Estruturada)

O sistema detecta automaticamente menções a funções ou endereços na query do usuário e recupera o contexto via call graph:

Detecção:

  1. Tokenização da query + regex para padrões comuns: sub_XXXX, 0xXXXXXXXX, nomes em PascalCase
  2. Para cada token candidato, busca em static_function via:
    List<StaticFunction> findByNameIgnoreCaseContainingAndAnalysisId(name, analysisId);
    
  3. Até 5 funções são identificadas (limite para não poluir o contexto)

Expansão 1-hop:

  • Callers: quem chama a função detectada (xrefRepository.findByCalleeId)
  • Callees: quem a função detectada chama (xrefRepository.findByCallerId)
  • Para funções detectadas (que podem não ter embedding), o assembly é montado inline no contexto, truncado a 5000 caracteres

Prioridade no contexto:

  1. Funções explicitamente mencionadas na query + seu assembly (sempre incluídas, mesmo sem embedding)
  2. Callers/callees imediatos (1-hop, assembly truncado a 3000 caracteres)
  3. Chunks semânticos (pgvector top-K)
  4. Chunks textuais (LIKE fallback)

2.4 Construção do System Prompt

O ChatService.buildSystemPrompt() monta o prompt do sistema que instrui a LLM:

You are an expert reverse engineer and malware analyst.
You are analyzing a binary file with the following characteristics.
Use the provided function disassembly and context to answer the user's questions.
When referencing code, mention function names, addresses, and the binary context.
If the provided context does not contain enough information, acknowledge the limitation
and suggest what additional analysis would help.

[BINARY OVERVIEW]
Filename: malware.exe
Format: PE32 (Portable Executable 32-bit)
Architecture: x86
Compiler: Microsoft Visual C++ 19.0

[REQUESTED FUNCTION ANALYSIS]    ← call graph expansion (se houver)

[FUNCTION] sub_401000 at 0x401000
Callers: WinMain, sub_402300
Callees: CreateFileA, WriteFile, CloseHandle
[ASSEMBLY]
push    ebp
mov     ebp, esp
...

[RELEVANT CODE SECTIONS]        ← semantic + textual results

Section 1:
[FUNCTION] sub_401200 at 0x401200
Callers: sub_401000, sub_402500
Callees: WS2_32.send, WS2_32.recv
[ASSEMBLY]
...

[CONVERSATION GUIDELINES]
- Cite function addresses when relevant (e.g., "the function at 0x401000")
- Explain assembly in natural language when useful
- If you identify patterns (loops, conditionals, API calls), highlight them
- If you detect malicious intent, point it out objectively
- Be concise but thorough in your technical analysis
- If the context includes a [REQUESTED FUNCTION ANALYSIS] section with full assembly,
  analyze that function in depth rather than saying you don't have access to it
  • Explique o assembly em linguagem natural quando útil
  • Se identificar padrões (ex: loop, condicional, chamada de API), destaque-os
  • Se perceber intenção maliciosa, aponte objetivamente

### 2.6 Truncagem de Contexto

O contexto total (system prompt + chunks + histórico) deve caber na janela de contexto da LLM:

| Modelo | Janela de Contexto | Max Context Tokens Config |
|---|---|---|
| Gemma4:12b | 256K tokens | 8000 (padrão) |
| DeepSeek-chat | 64K tokens | 8000 (padrão) |

**Algoritmo de truncagem**:

1. **Reserva para system prompt**: ~500 tokens (fixo)
2. **Reserva para histórico**: últimas 20 mensagens, truncadas se necessário (~2000 tokens)
3. **Reserva para resposta**: ~1000 tokens (output)
4. **Disponível para chunks**: `max-context-tokens - 500 - 2000 - 1000 = ~4500 tokens`
5. **Preenche chunks em ordem de relevância** até atingir o limite
6. Chunks que não cabem são descartados (os menos relevantes primeiro)

### 2.7 Histórico de Conversa

O histórico da sessão (`chat_message`) é incluído no prompt para manter contexto multi-turn:

Messages enviados à LLM: [ { role: SYSTEM, content: "" }, { role: USER, content: "Quantas funções esse binário tem?" }, { role: ASSISTANT, content: "342 funções." }, { role: USER, content: "E dessas, quantas importam APIs de rede?" }, // ↑ A LLM usa o contexto do system prompt + histórico para responder ]


- Apenas as últimas `max-history-messages` (default 20) são incluídas
- Mensagens muito antigas são truncadas para caber no orçamento de tokens

### 2.8 Resposta Streaming (SSE)

Tokens são enviados ao frontend assim que gerados pela LLM:

data: {"type":"chunk","content":"O binário"}

data: {"type":"chunk","content":" importa "}

data: {"type":"chunk","content":"funções de"}

data: {"type":"chunk","content":" rede via"}

data: {"type":"chunk","content":" WS2_32"}

...

data: {"type":"done","messageId":"550e8400-e29b-...","sessionId":"660e8400-..."}


### 2.9 Tipos de Query e Estratégias

| Categoria | Exemplo de Query | Estratégia de Retrieval |
|---|---|---|
| **Análise de função** | "O que a função sub_401000 faz?" | Call graph expansion + assembly inline (1-hop) |
| **Busca por nome** | "Explique sub_f5" | Textual fallback (LIKE) + call graph se detectado |
| **Busca semântica** | "Quais funções lidam com criptografia?" | pgvector similarity search (top-10) |
| **Navegação do grafo** | "Quem chama WinMain?" | Call graph (callers diretos via xrefs) |
| **Metadados** | "Qual arquitetura/com pilador?" | Chunk METADATA (sempre incluído) |
| **Strings/Imports** | "Esse binário usa rede?" | pgvector (chunks que mencionam imports) |
| **Classificação** | "Isso parece malware?" | pgvector + METADATA (visão holística) |
| **Controle de fluxo** | "Explique o fluxo a partir de main" | Call graph (multi-hop callees) |

> **Nota**: O fallback textual e a expansão por call graph garantem que queries com nomes específicos de funções (ex: `sub_f5`, `0x401000`) sempre retornem o assembly relevante — mesmo que a busca semântica retorne 0 resultados.

---

## Modelo de Embedding: Qwen3 Embedding

| Propriedade | Valor |
|---|---|
| Modelo | `qwen3-embedding:latest` (8B) |
| Tamanho | 4.7 GB |
| Contexto máximo | 32K tokens |
| Dimensão configurável | 324096 (Spring AI 1.0.1 não expõe o parâmetro; usamos default do modelo: **4096**) |
| Ranking MTEB | #1 (score 70.58) |
| Idiomas | 100+ (incluindo código) |
| Execução | Ollama local (localhost:11434) |

**Escolha da dimensão 4096**: O modelo `qwen3-embedding:latest` (8B) retorna 4096 dimensões por padrão. O Spring AI 1.0.1 não expõe o parâmetro `dimensions` da API do Ollama (é um campo top-level, não dentro de `options`). Para 330 funções a 4096 floats (4 bytes cada), o storage total de vetores é ~5.4 MB. A coluna é `vector(4096)` e a busca usa sequential scan (sem índice — HNSW e IVFFlat do pgvector limitados a 2000d).