Rodrigo Verdiani
90ee74b3b9
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.
370 lines
13 KiB
Markdown
370 lines
13 KiB
Markdown
# Modelo de Dados — Módulo AI
|
|
|
|
## Diagrama Entidade-Relacionamento
|
|
|
|
```
|
|
┌─────────────────────────────────────────────────────────────┐
|
|
│ TABELAS EXISTENTES │
|
|
│ │
|
|
│ binaries ──────┬── static_analysis ──┬── static_function │
|
|
│ (workspace) │ ├── static_label │
|
|
│ │ └── static_xref │
|
|
│ │ │
|
|
│ └── jobs │
|
|
└─────────────────────────────────────────────────────────────┘
|
|
│
|
|
│ FK: binary_id
|
|
▼
|
|
┌─────────────────────────────────────────────────────────────┐
|
|
│ NOVAS TABELAS (V008) │
|
|
│ │
|
|
│ embedding_chunk ──── chat_session ──── chat_message │
|
|
│ │ │ │
|
|
│ │ │ FK: binary_id → binaries │
|
|
│ │ │ FK: project_id → projects │
|
|
│ │ │ │
|
|
│ │ FK: binary_id │ │
|
|
│ │ → binaries │ │
|
|
│ │ │ │
|
|
│ │ FK: source_id │ │
|
|
│ │ → opcional │ │
|
|
│ │ (static_func) │ │
|
|
└─────────────────────────────────────────────────────────────┘
|
|
```
|
|
|
|
## Migration V008 — DDL Completo
|
|
|
|
```sql
|
|
-- V008__ai_embedding_and_chat.sql
|
|
|
|
-- Garantir que a extensão pgvector está ativada
|
|
CREATE EXTENSION IF NOT EXISTS vector;
|
|
|
|
-- ============================================================
|
|
-- Tabela: embedding_chunk
|
|
-- Armazena chunks de texto embeddados para busca semântica
|
|
-- ============================================================
|
|
CREATE TABLE embedding_chunk (
|
|
id UUID PRIMARY KEY,
|
|
|
|
-- Referência ao binário (obrigatória)
|
|
binary_id UUID NOT NULL REFERENCES binaries(id) ON DELETE CASCADE,
|
|
|
|
-- Tipo de chunk: FUNCTION, METADATA
|
|
chunk_type VARCHAR(30) NOT NULL,
|
|
|
|
-- Tipo da fonte: STATIC_FUNCTION, IMPORT, STRING (extensível)
|
|
source_type VARCHAR(30),
|
|
|
|
-- FK opcional para a entidade fonte (ex: static_function.id)
|
|
source_id UUID,
|
|
|
|
-- Texto completo que foi embeddado
|
|
content TEXT NOT NULL,
|
|
|
|
-- Vetor de embedding (dimensão configurável, default 4096)
|
|
-- qwen3-embedding:latest (8B) retorna 4096 dimensões
|
|
embedding vector(4096),
|
|
|
|
-- Metadados flexíveis em JSON:
|
|
-- {
|
|
-- "function_name": "sub_401000",
|
|
-- "address": "0x401000",
|
|
-- "callers": ["WinMain", "sub_402300"],
|
|
-- "callees": ["CreateFileA", "WriteFile"],
|
|
-- "label_count": 2
|
|
-- }
|
|
metadata JSONB,
|
|
|
|
-- Contagem estimada de tokens (para controle de contexto)
|
|
token_count INTEGER,
|
|
|
|
-- Timestamp de criação
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
|
|
-- Índice para busca por binário (filtro de escopo)
|
|
CREATE INDEX idx_embedding_chunk_binary ON embedding_chunk(binary_id);
|
|
|
|
-- Índice para lookup por fonte (ex: achar chunk de uma função específica)
|
|
CREATE INDEX idx_embedding_chunk_source ON embedding_chunk(source_type, source_id);
|
|
|
|
-- Índice de busca textual (fallback para queries específicas)
|
|
CREATE INDEX idx_embedding_chunk_content ON embedding_chunk
|
|
USING gin (content gin_trgm_ops);
|
|
|
|
-- NOTA: Não há índice vetorial. HNSW e IVFFlat do pgvector limitados a 2000d.
|
|
-- Sequential scan é rápido para < 10K vetores de 4096d.
|
|
|
|
-- ============================================================
|
|
-- Tabela: chat_session
|
|
-- Representa uma sessão de conversa com o chat IA
|
|
-- ============================================================
|
|
CREATE TABLE chat_session (
|
|
id UUID PRIMARY KEY,
|
|
|
|
-- Binário em análise (escopo da conversa)
|
|
binary_id UUID NOT NULL REFERENCES binaries(id) ON DELETE CASCADE,
|
|
|
|
-- Projeto pai (para navegação/breadcrumb)
|
|
project_id UUID NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
|
|
|
|
-- Título da sessão (ex: "Analisando funções de rede")
|
|
title VARCHAR(255),
|
|
|
|
-- Modelo de chat usado nesta sessão
|
|
-- Ex: "gemma4:12b", "deepseek-chat", "gemma4:e4b"
|
|
chat_model VARCHAR(100) NOT NULL,
|
|
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
|
|
CREATE INDEX idx_chat_session_binary ON chat_session(binary_id);
|
|
|
|
-- ============================================================
|
|
-- Tabela: chat_message
|
|
-- Mensagens individuais dentro de uma sessão
|
|
-- ============================================================
|
|
CREATE TABLE chat_message (
|
|
id UUID PRIMARY KEY,
|
|
|
|
-- Sessão pai
|
|
session_id UUID NOT NULL REFERENCES chat_session(id) ON DELETE CASCADE,
|
|
|
|
-- Papel do emissor: USER, ASSISTANT, SYSTEM
|
|
role VARCHAR(20) NOT NULL,
|
|
|
|
-- Conteúdo textual da mensagem
|
|
content TEXT NOT NULL,
|
|
|
|
-- Estimativa de tokens (para tracking de uso)
|
|
token_count INTEGER,
|
|
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
|
|
-- Índice para buscar histórico em ordem cronológica
|
|
CREATE INDEX idx_chat_message_session ON chat_message(session_id, created_at);
|
|
```
|
|
|
|
## Entidades JPA
|
|
|
|
### `EmbeddingChunk`
|
|
|
|
```java
|
|
@Entity
|
|
@Table(name = "embedding_chunk")
|
|
@Getter @Setter @NoArgsConstructor @AllArgsConstructor @Builder
|
|
public class EmbeddingChunk {
|
|
|
|
@Id
|
|
@GeneratedValue(strategy = GenerationType.UUID)
|
|
private UUID id;
|
|
|
|
@Column(name = "binary_id", nullable = false)
|
|
private UUID binaryId;
|
|
|
|
@Column(name = "chunk_type", nullable = false, length = 30)
|
|
private String chunkType;
|
|
// Valores: "FUNCTION", "METADATA"
|
|
|
|
@Column(name = "source_type", length = 30)
|
|
private String sourceType;
|
|
// Valores: "STATIC_FUNCTION", null (para METADATA)
|
|
|
|
@Column(name = "source_id")
|
|
private UUID sourceId;
|
|
// FK lógica para static_function.id (sem constraint de FK no banco
|
|
// porque pode ser null ou apontar para tabelas diferentes no futuro)
|
|
|
|
@Column(nullable = false, columnDefinition = "TEXT")
|
|
private String content;
|
|
|
|
@JdbcTypeCode(SqlTypes.VECTOR)
|
|
@Column(name = "embedding", columnDefinition = "vector(1024)")
|
|
private float[] embedding;
|
|
|
|
@JdbcTypeCode(SqlTypes.JSON)
|
|
@Column(name = "metadata", columnDefinition = "jsonb")
|
|
private String metadata;
|
|
// JSON string: {"function_name":"...", "address":"...", "callers":[...], ...}
|
|
|
|
@Column(name = "token_count")
|
|
private Integer tokenCount;
|
|
|
|
@CreationTimestamp
|
|
@Column(name = "created_at", nullable = false, updatable = false)
|
|
private Instant createdAt;
|
|
}
|
|
```
|
|
|
|
**Notas sobre o mapeamento**:
|
|
- `embedding` usa `@JdbcTypeCode(SqlTypes.VECTOR)` — requer Hibernate 6.2+ (Spring Boot 4.0.6 inclui)
|
|
- `metadata` é `String` com `@JdbcTypeCode(SqlTypes.JSON)` — Hibernate converte de/para JSONB automaticamente
|
|
- `sourceId` é FK lógica (não constraint no banco) para permitir flexibilidade futura (apontar para outras tabelas além de `static_function`)
|
|
|
|
### `ChatSession`
|
|
|
|
```java
|
|
@Entity
|
|
@Table(name = "chat_session")
|
|
@Getter @Setter @NoArgsConstructor @AllArgsConstructor @Builder
|
|
public class ChatSession {
|
|
|
|
@Id
|
|
@GeneratedValue(strategy = GenerationType.UUID)
|
|
private UUID id;
|
|
|
|
@Column(name = "binary_id", nullable = false)
|
|
private UUID binaryId;
|
|
|
|
@ManyToOne(fetch = FetchType.LAZY)
|
|
@JoinColumn(name = "binary_id", insertable = false, updatable = false)
|
|
private Binary binary;
|
|
|
|
@Column(name = "project_id", nullable = false)
|
|
private UUID projectId;
|
|
|
|
@ManyToOne(fetch = FetchType.LAZY)
|
|
@JoinColumn(name = "project_id", insertable = false, updatable = false)
|
|
private Project project;
|
|
|
|
@Column(length = 255)
|
|
private String title;
|
|
|
|
@Column(name = "chat_model", nullable = false, length = 100)
|
|
private String chatModel;
|
|
|
|
@OneToMany(mappedBy = "session", cascade = CascadeType.ALL, orphanRemoval = true)
|
|
@Builder.Default
|
|
private List<ChatMessage> messages = new ArrayList<>();
|
|
|
|
@CreationTimestamp
|
|
@Column(name = "created_at", nullable = false, updatable = false)
|
|
private Instant createdAt;
|
|
|
|
@UpdateTimestamp
|
|
@Column(name = "updated_at", nullable = false)
|
|
private Instant updatedAt;
|
|
}
|
|
```
|
|
|
|
### `ChatMessage`
|
|
|
|
```java
|
|
@Entity
|
|
@Table(name = "chat_message")
|
|
@Getter @Setter @NoArgsConstructor @AllArgsConstructor @Builder
|
|
public class ChatMessage {
|
|
|
|
@Id
|
|
@GeneratedValue(strategy = GenerationType.UUID)
|
|
private UUID id;
|
|
|
|
@ManyToOne(fetch = FetchType.LAZY)
|
|
@JoinColumn(name = "session_id", nullable = false)
|
|
private ChatSession session;
|
|
|
|
@Column(nullable = false, length = 20)
|
|
private String role; // "USER", "ASSISTANT", "SYSTEM"
|
|
|
|
@Column(nullable = false, columnDefinition = "TEXT")
|
|
private String content;
|
|
|
|
@Column(name = "token_count")
|
|
private Integer tokenCount;
|
|
|
|
@CreationTimestamp
|
|
@Column(name = "created_at", nullable = false, updatable = false)
|
|
private Instant createdAt;
|
|
}
|
|
```
|
|
|
|
## Relacionamentos
|
|
|
|
```
|
|
chat_session (1) ────────── (N) chat_message
|
|
│ │
|
|
│ FK binary_id → binaries │ FK session_id → chat_session
|
|
│ FK project_id → projects │
|
|
│ │
|
|
└── cascade: ALL, orphanRemoval └── cascade: ALL, orphanRemoval
|
|
(deletar sessão deleta (deletar sessão deleta
|
|
todas as mensagens) todas as mensagens)
|
|
|
|
embedding_chunk (N) ────────── (1) binaries
|
|
│
|
|
│ FK binary_id → binaries
|
|
│ ON DELETE CASCADE (deletar binário deleta todos os chunks)
|
|
│
|
|
└── source_id: FK lógica para static_function.id
|
|
(não é constraint; usado para join manual)
|
|
```
|
|
|
|
## Índices e Performance
|
|
|
|
### Índices
|
|
|
|
| Índice | Coluna(s) | Uso |
|
|
|---|---|---|
|
|
| `idx_embedding_chunk_binary` | `binary_id` | Filtrar chunks por binário (todas as queries) |
|
|
| `idx_embedding_chunk_source` | `source_type, source_id` | Lookup por fonte (ex: achar chunk de uma função específica) |
|
|
| `idx_embedding_chunk_content` | `content` (GIN trgm) | Busca textual de fallback (`ILIKE '%term%'`) |
|
|
| `idx_chat_session_binary` | `binary_id` | Listar sessões de um binário |
|
|
| `idx_chat_message_session` | `session_id, created_at` | Buscar histórico em ordem cronológica |
|
|
|
|
> **Performance**: Sem índice vetorial (HNSW/IVFFlat limitados a 2000d). Para < 10K vetores de 4096d, sequential scan < 100ms.
|
|
|
|
## Coluna `metadata` (JSONB)
|
|
|
|
O campo `metadata` armazena informações estruturadas que não fazem parte do texto embeddado mas são úteis para filtragem e exibição:
|
|
|
|
```json
|
|
{
|
|
"function_name": "sub_401000",
|
|
"address": "0x401000",
|
|
"callers": ["WinMain", "sub_402300"],
|
|
"callees": ["CreateFileA", "WriteFile", "CloseHandle"],
|
|
"label_count": 2,
|
|
"assembly_length": 1247,
|
|
"has_decompiled_code": false
|
|
}
|
|
```
|
|
|
|
Para chunks do tipo `METADATA`:
|
|
|
|
```json
|
|
{
|
|
"filename": "malware.exe",
|
|
"format": "PE32",
|
|
"architecture": "x86",
|
|
"compiler": "MSVC 19.0",
|
|
"function_count": 342
|
|
}
|
|
```
|
|
|
|
Vantagens do JSONB sobre colunas separadas:
|
|
- Flexível: adicionar campos sem migration
|
|
- Indexável: `CREATE INDEX ... ON embedding_chunk USING gin (metadata jsonb_path_ops)` se necessário
|
|
- Permite queries como: `SELECT * FROM embedding_chunk WHERE metadata->>'function_name' = 'WinMain'`
|
|
|
|
## Políticas de Deleção
|
|
|
|
| Operação | Comportamento |
|
|
|---|---|
|
|
| Deletar binário | `ON DELETE CASCADE` → deleta todos os `embedding_chunk`, `chat_session` e `chat_message` |
|
|
| Deletar projeto | `ON DELETE CASCADE` via binário → cascata completa |
|
|
| Re-analisar binário | Hard delete de `embedding_chunk` via `deleteByBinaryId()` antes de reindexar |
|
|
| Deletar sessão | `orphanRemoval = true` → deleta todas as `chat_message` |
|
|
|
|
## Relação com o Schema Existente
|
|
|
|
A migration V008 é aditiva — não altera nenhuma tabela existente. As novas tabelas referenciam:
|
|
- `binaries(id)` — FK com CASCADE
|
|
- `projects(id)` — FK com CASCADE
|
|
|
|
O `source_id` em `embedding_chunk` referencia logicamente `static_function.id`, mas **sem constraint de FK** para:
|
|
- Permitir valores nulos (chunk METADATA não tem source)
|
|
- Permitir referências futuras a outras tabelas (ex: tabela de strings, tabela de imports)
|