VulTriage

KO | EN

gitlite — search

#c #nlp #javascript #video #deep-learning #machine-learning #kotlin #mcp #privacy #postgresql

(★ 9)

The code of VulTriage: Triple-Path Context Augmentation for LLM-Based Vulnerability Detection

File Explorer

ZIP Download

- - profiles_settings.xml
  - Project_Default.xml
- .gitignore
- deployment.xml
- main.iml
- misc.xml
- modules.xml
- 1435_filtered_columns.csv
- - ._control_flow
- - generate_graphs.py
  - README.md
  - vuln_flow_analyzer.py
- cwec_v4.19.1.xml
- main.png
- title.png
.DS_Store
hyperparam_exp.py
main.py
metrics_tracker.py
README.md
requirements.txt
retrieval.py

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

![ChatGPT Image 2026年5月10日 16_51_24](https://raw.githubusercontent.com/vinsontang1/VulTriage/main/fig/title.png)

An intelligent vulnerability detection framework based on Large Language Models and static code analysis, supporting multi-dataset evaluation and ablation studies.

![image-20260510164600403](https://raw.githubusercontent.com/vinsontang1/VulTriage/main/fig/main.png)

## 📋 Table of Contents

- [Introduction](#introduction)
- [Key Features](#key-features)
- [System Requirements](#system-requirements)
- [Quick Start](#quick-start)
- [Usage](#usage)
- [Datasets & Models](#datasets--models)
- [FAQ](#faq)

## Introduction

This project is an intelligent vulnerability detection framework based on GPT-4o and BGE-M3 models, combining:
- **Static Structure Analysis**: Control Flow Graphs (CFG), Data Flow Graphs (DFG), Abstract Syntax Trees (AST)
- **Retrieval-Augmented Generation (RAG)**: Vulnerability pattern matching based on CWE knowledge base
- **Code Understanding**: Semantic analysis using Large Language Models
- **Paired/Single-sample Detection**: Support for multiple evaluation modes

## Key Features

- ✅ Support for PrimeVul and Kotlin datasets
- ✅ Integrated CWE vulnerability knowledge base retrieval
- ✅ Configurable ablation studies
- ✅ Real-time metrics tracking
- ✅ Flexible hyperparameter experiment support

## System Requirements

- Python >= 3.9
- CUDA (optional, for GPU acceleration)

## Quick Start

### 1. Install Dependencies
```bash
pip install -r requirements.txt
```

### 2. Prepare Resources

#### Download BGE-M3 Model
The model will be downloaded automatically, or you can download it manually:
```bash
git clone https://huggingface.co/BAAI/bge-m3
```

#### Prepare CWE Data
Ensure that `cwec_latest.xml/cwec_v4.19.1.xml` and `1435.csv/1435_filtered_columns.csv` (CWE index data) files exist in the project root directory.
> Note: This project uses CWE Top 25 Most Dangerous Software Weaknesses. Related files are already included in the project root directory.

#### Prepare Datasets
Place the following files in the project root directory:
- `primevul_test_paired.jsonl` (PrimeVul dataset)
- `LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv` (Kotlin dataset)

### 3. Configure API Key

Modify the OpenAI API configuration in `main.py` or `hyperparam_exp.py`:
```python
client = OpenAI(
    api_key="YOUR_API_KEY",
    base_url="https://api.poixe.com/v1"
)
```
> Note: The current URL `https://api.poixe.com/v1` failed to parse. It may be an unsupported webpage type or a broken link. Please verify the link validity or switch to another API provider.

### 4. Run Vulnerability Detection

#### Basic Mode (Full Functionality)
```bash
python main.py --wotask all --jsonl_file primevul_test_paired.jsonl
```

#### Ablation Study Mode
```bash
# Without control flow information
python main.py --wotask wo_ctrl_info_pth --jsonl_file primevul_test_paired.jsonl
# Without RAG knowledge
python main.py --wotask wo_rag_pth --jsonl_file primevul_test_paired.jsonl
# Without code understanding
python main.py --wotask wo_exp_pth --jsonl_file primevul_test_paired.jsonl
# Without any auxiliary information
python main.py --wotask nan --jsonl_file primevul_test_paired.jsonl
```

#### Hyperparameter Experiments
```bash
# Pairwise evaluation mode
python hyperparam_exp.py --mode pairwise --input_file primevul_test_paired.jsonl --detail_level C --info_comb all --desc_level concise --max_vuln 2
# Single-sample evaluation mode
python hyperparam_exp.py --mode single --input_file LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv
```

### 5. View Results

Detection results are saved in real-time to `metrics_log.txt` and final evaluation summaries are output to the terminal:
- Evaluation metrics include accuracy, recall, false positive rate, etc.
- Ablation study results compare performance differences across different modes to analyze the contribution of each module.

## Usage

### Main Scripts Overview

| Script | Purpose |
|--------|---------|
| `main.py` | Main vulnerability detection script with ablation study support |
| `hyperparam_exp.py` | Hyperparameter experiments and multi-mode evaluation |
| `retrieval.py` | CWE knowledge base retrieval module (hybrid retrieval based on BGE-M3 model) |
| `metrics_tracker.py` | Real-time metrics tracker with log saving and visualization support |

### Parameter Description

#### main.py Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--wotask` | `all` | Ablation mode: `all` (full functionality), `wo_ctrl_info_pth` (no static structure info), `wo_rag_pth` (no RAG knowledge), `wo_exp_pth` (no code understanding), `nan` (no auxiliary information) |
| `--jsonl_file` | `primevul_test_paired.jsonl` | JSONL dataset path |
| `--metrics_file` | `metrics_log.txt` | Metrics output file |

#### hyperparam_exp.py Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--mode` | `pairwise` | Evaluation mode: `pairwise` or `single` |
| `--input_file` | Required | Input file path (JSONL or CSV format) |
| `--detail_level` | `C` | Code analysis detail level: `C` (complete), `B` (standard), `A` (brief) |
| `--info_comb` | `all` | Information combination: `all`, `no_ast`, `no_cfg`, `no_dfg` |
| `--desc_level` | `concise` | Description level: `concise`, `normal`, `detailed` |
| `--max_vuln` | `2` | Maximum number of vulnerabilities: `2` or `4` |

## Datasets & Models

### 📊 Datasets

#### 1. PrimeVul Dataset
- **Citation**: Yangruibo Ding, Yanjun Fu, Omniyyah Ibrahim, Chawin Sitawarin, Xinyun Chen, Basel Alomair, David Wagner, Baishakhi Ray, and Yizheng Chen. 2025. Vulnerability detection with code language models: How far are we? ISSTA 2025 (2025).
- **Purpose**: Paired vulnerability detection benchmark
- **Format**: JSONL (contains `func` and `target` fields)
- **Source**: ISSTA 2025 Conference ([ISSTA Official Website](https://issta.org/)). Please refer to the original paper for dataset access permissions.

#### 2. Kotlin Dataset
- **Citation**: Triet Huynh Minh Le, M Ali Babar, and Tung Hoang Thai. 2024. Software vulnerability prediction in low-resource languages: An empirical study of codebert and chatgpt. In Proceedings of the 28th International Conference on Evaluation and Assessment in Software Engineering, 679–685.
- **Purpose**: Low-resource language vulnerability prediction research
- **Format**: CSV (contains `code` and `label` columns)
- **Source**: EASE 2024 Conference ([EASE 2024 Official Website](https://conf.researchr.org/home/ease-2024)). Please refer to the original paper for dataset access permissions.

#### 3. CWE Database
- **Version**: CWE v4.19.1
- **Source**: [MITRE CWE](https://cwe.mitre.org/)
- **Purpose**: Vulnerability pattern knowledge base and retrieval augmentation
- **Note**: CWE can be freely used for research, development, and commercial purposes. Copyright attribution to MITRE must be retained.

### 🤖 Models

#### 1. BGE-M3 Embedding Model
- **Developer**: BAAI (Beijing Academy of Artificial Intelligence)
- **HuggingFace**: [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3)
- **License**: MIT License (Note: Current model license page parsing failed. Refer to FlagEmbedding project license which is MIT License)
- **Purpose**: Dense and sparse vector hybrid retrieval, supporting multilingual, multi-granularity (max input length 8192 tokens), and multi-functional retrieval
- **Related Project**: [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding) (Open-source repository for BGE series models)

#### 2. GPT-4o
- **Developer**: OpenAI
- **Model Version**: gpt-4o-2024-11-20
- **License**: Follows [OpenAI Terms of Use](https://openai.com/policies/terms-of-use) (effective January 1, 2026)
- **Purpose**: Code understanding, vulnerability identification, reasoning and judgment
- **Note**: Usage must comply with OpenAI Terms of Service. Illegal or harmful activities and model reverse engineering are prohibited.

### 📦 Dependency Libraries

| Library | Version | License |
|---------|---------|---------|
| FlagEmbedding | 1.4.0 | MIT |
| openai | 2.26.0 | Apache 2.0 |

## FAQ

### Q1: How to change the OpenAI API provider?
Modify the `base_url` and `api_key` in `main.py` or `hyperparam_exp.py`.

### Q2: Why are the retrieval results inaccurate?
1. Check if `1435.csv/1435_filtered_columns.csv` is loaded correctly and ensure the CWE index data format is correct;
2. Confirm that the BGE-M3 model has been downloaded correctly. You can manually download it via `git clone https://huggingface.co/BAAI/bge-m3`;

### Q3: How to customize evaluation metrics?
Modify the `update_pair_metrics` method in `metrics_tracker.py` to add or adjust metric calculation logic.

### Q4: What if BGE-M3 model download fails?
Download the model files directly from the HuggingFace website, extract them to the project root directory, or check your network connection and re-execute the `git clone https://huggingface.co/BAAI/bge-m3` command.

VulTriage

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

VulTriage

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

Download README