VulTriage

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

📋 Table of Contents

• Introduction
• Key Features
• System Requirements
• Quick Start
• Usage
• Datasets & Models
• 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

pip install -r requirements.txt

2. Prepare Resources

Download BGE-M3 Model

The model will be downloaded automatically, or you can download it manually:

git clone https://huggingface.co/BAAI/bge-m3
# Note: BGE-M3 related dependencies can be installed via FlagEmbedding. For details, refer to: pip install -U FlagEmbedding

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:

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)

python main.py --wotask all --jsonl_file primevul_test_paired.jsonl

Ablation Study Mode

# 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

# 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). 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). Please refer to the original paper for dataset access permissions.

3. CWE Database

• Version: CWE v4.19.1
• Source: MITRE CWE
• 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
• 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 (Open-source repository for BGE series models)

2. GPT-4o

• Developer: OpenAI