An LLM router for C/C++ vulnerability detection. Instead of always calling a large (expensive/slow) code LLM to classify a function as vulnerable or not, this project learns — per code sample — which smaller/cheaper local LLMs already get the right answer, and routes new (unseen) code to the cheapest LLM likely to classify it correctly, based on embedding similarity (KNN) to labeled training data.
All LLM inference (embeddings + chat models) runs locally through Ollama.
- Take a labeled vulnerability dataset (MSR C/C++ functions,
vul= 0/1). - Run several candidate LLMs (e.g.
codestral,qwen2.5-coder:14b,qwen2.5-coder:32b,llama3.1:8b,codegemma:7b) over the training set and record which model(s) got each row right. - Embed all training code with
embeddinggemma(via Ollama) and index the embeddings in FAISS. - At inference time, embed a new (test) code sample, retrieve its top-k nearest neighbors in the training set, and look at which LLM(s) were correct on all of those neighbors (shrinking k until the intersection is non-empty).
- Route the query to the cheapest/smallest model in that intersection (via a priority list) and call it for the actual prediction.
This assumes: if several LLMs already solved similar code correctly, the cheapest of those LLMs will likely solve this new, similar code correctly too — avoiding an expensive model call when a cheap one would suffice.
flowchart TD
A[MSR_data_cleaned.csv\ndata/raw/functions/] --> B[prepare_dataset.py /\ncreate_balanced_dataset.py]
B --> C[msr_cleaned_full.csv / msr_10pct.csv\ndata/processed/]
C --> D[split_dataset.py]
D --> E[train_full.csv / test_full.csv\ndata/processed/full_dataset/split/]
E --> F[generate_embeddings.py\nOllama: embeddinggemma]
F --> G[embeddings_output/train_nochunk/\nembeddings.npy, rows.csv, row_ids.npy]
F --> G2[embeddings_output/test_nochunk/\nembeddings.npy, rows.csv, row_ids.npy]
E --> H[llm_predictions.py\nOllama: codestral / qwen2.5-coder / llama3.1 / codegemma]
H --> I[results/*.csv\nper-model, per-row predictions]
G --> J[training_embedding_map.py]
J --> K[train_master_table.csv\ndata/processed/]
I --> L[llm_list_mapping.py]
K --> L
L --> M[sample_correct_llms.csv\ndata/processed/\nwhich LLM(s) got each sample_id right]
K --> N[emb_faiss.py]
N --> O[train_faiss.index + train_sample_ids.npy\ndata/processed/]
G2 --> P[routing_pipeline.py]
O --> P
M --> P
P --> Q[results/routing_results_k*.csv\nresults/routing_metrics_k*.csv\nresults/routing_confusion_matrix_k*.csv]
src/
dataset/
prepare_dataset.py # 10% random sample of C/C++ rows (quick iteration dataset)
create_balanced_dataset.py # full cleaned C/C++ dataset, labels normalized to 0/1
split_dataset.py # stratified train/test split (80/20) on the cleaned dataset
llm_predictions.py # runs one Ollama model over a CSV, saves per-row predictions
training_embedding_map.py # joins train.csv + rows.csv + embeddings.npy -> master table
llm_list_mapping.py # aggregates results/*.csv -> which LLM(s) were correct per sample_id
cwe_prompt_utils.py # shared helpers for the CWE-type pipeline (label set, prompt, parsing)
build_cwe_label_map.py # scans msr_cleaned_full.csv -> collapsed CWE label set + descriptions
cwe_predictions.py # like llm_predictions.py, but classifies CWE-type (closed-set) not vul
cwe_llm_list_mapping.py # like llm_list_mapping.py, but for the CWE-task results_cwe/*.csv
embeddings/
generate_embeddings.py # batched embedding generation via Ollama (embeddinggemma), resumable
emb_faiss.py # builds a FAISS IndexFlatIP (cosine) over train embeddings
faiss_search_topk.py # ad-hoc top-k nearest neighbor sanity-check script
evaluation/
evaluate_knn.py # LEGACY: KNN baseline using an older chunked-embedding layout
# (embeddings_output/train/{chunks.csv,faiss_index.bin,...}).
# Superseded by faiss_search_topk.py + routing_pipeline.py,
# which use the "nochunk" embedding layout. Kept for reference.
routing_pipeline.py # THE routing pipeline: embed -> FAISS top-k -> shrink-k
# intersection of correct_llms -> pick cheapest -> call it -> score
routing_pipeline_cwe.py # same routing logic, but predicts CWE-type (multi-class) not vul
data/ # gitignored — generated/raw data, not checked in
raw/functions/ # MSR_data_cleaned.csv (source dataset, external)
processed/
msr_cleaned_full.csv # full cleaned dataset (create_balanced_dataset.py output)
msr_10pct.csv # 10% sample (prepare_dataset.py output)
msr_balanced_2k.csv # a smaller/balanced working subset
train.csv / test.csv # working train/test split used by most scripts
train_master_table.csv # sample_id <-> embedding_id <-> label <-> code join table
sample_correct_llms.csv # per-sample_id list of LLMs that predicted correctly
train_faiss.index # FAISS index over train embeddings
train_sample_ids.npy # sample_id order aligned with train_faiss.index
cwe_label_map.csv # raw_cwe_id -> collapsed cwe_label (build_cwe_label_map.py output)
cwe_descriptions.csv # cwe_label -> short description used in the CWE prompt
sample_correct_llms_cwe.csv # per-sample_id list of LLMs correct on the CWE-type task
full_dataset/
split/train_full.csv # stratified 80% split of msr_cleaned_full.csv
split/test_full.csv # stratified 20% split of msr_cleaned_full.csv
embeddings_output/ # gitignored — Ollama embedding runs
train_nochunk/ # embeddings for train set (embeddings.npy, rows.csv, row_ids.npy,
# manifest.json, _checkpoint.json for resumable runs)
test_nochunk/ # embeddings for test set, same layout
test_smoke/ # small smoke-test embedding run
train_chunked_backup/ # older chunked-embedding approach, kept as backup
results/ # gitignored — model outputs & routing evaluation
code_gemma_7b.csv, codestral.csv, llama31_8b_test.csv,
qwen2.5_coder_14b.csv, qwen2.5_coder_32b.csv # per-model predictions (from llm_predictions.py)
faiss_search_top3_nochunk.csv # output of faiss_search_topk.py
routing_results_k*.csv, routing_metrics_k*.csv,
routing_confusion_matrix_k*.csv # output of routing_pipeline.py
results_cwe/ # gitignored — CWE-task model outputs & routing evaluation
<model>.csv # per-model CWE predictions (from cwe_predictions.py)
routing_results_cwe_k*.csv, routing_metrics_cwe_k*.csv,
routing_confusion_matrix_cwe_k*.csv,
routing_classification_report_cwe_k*.csv # output of routing_pipeline_cwe.py
python src/dataset/create_balanced_dataset.py # -> data/processed/msr_cleaned_full.csv
python src/dataset/split_dataset.py # -> data/processed/full_dataset/split/{train_full,test_full}.csv(prepare_dataset.py is an alternative that produces a quick 10% sample instead of the full
cleaned set — useful for fast local iteration.)
The resulting train/test CSVs are then copied/renamed to data/processed/train.csv and
data/processed/test.csv, which most downstream scripts (embeddings, master table, routing)
default to reading from.
python src/embeddings/generate_embeddings.py --input data/processed/train.csv --output embeddings_output/train_nochunk
python src/embeddings/generate_embeddings.py --input data/processed/test.csv --output embeddings_output/test_nochunkThis is checkpointed/resumable (_checkpoint.json) and retries failed rows via batch-halving.
python src/dataset/llm_predictions.py --model codestral --output results/codestral.csv
python src/dataset/llm_predictions.py --model qwen2.5-coder:14b --output results/qwen2.5_coder_14b.csv
python src/dataset/llm_predictions.py --model qwen2.5-coder:32b --output results/qwen2.5_coder_32b.csv
python src/dataset/llm_predictions.py --model llama3.1:8b --output results/llama31_8b.csv
python src/dataset/llm_predictions.py --model codegemma:7b --output results/code_gemma_7b.csvEach run asks the model for a strict 0/1 verdict per row and records correctness.
python src/dataset/training_embedding_map.py # -> data/processed/train_master_table.csv
python src/embeddings/emb_faiss.py # -> data/processed/train_faiss.index, train_sample_ids.npypython src/dataset/llm_list_mapping.py # -> data/processed/sample_correct_llms.csvReads every CSV in results/, aggregates correctness per sample_id, and outputs the list of
LLMs that answered correctly (correct_llms) vs. all LLMs that were tried (all_llms).
python src/evaluation/routing_pipeline.py \
--k 5 \
--model_priority codegemma:7b llama3.1:8b qwen2.5-coder:14b qwen2.5-coder:32b codestral \
--out_dir resultsFor each test row: embed (precomputed) -> FAISS top-k -> take the intersection of
correct_llms across the k neighbors, shrinking k by one until non-empty -> pick the first
model in --model_priority present in that intersection -> call it via Ollama chat -> record
prediction vs. ground truth. Outputs per-k routing results, accuracy/precision/recall, and a
confusion matrix to results/.
src/embeddings/faiss_search_topk.py is a smaller standalone script useful for sanity-checking
retrieval quality (dumps top-3 neighbors + labels for the first 100 test rows) without running
the full routing/LLM-calling pipeline.
The full cleaned dataset is far larger than what's practical to run through the (unbatched,
one-call-per-row) LLM-benchmarking step (llm_predictions.py / cwe_predictions.py) for every
candidate model. Running that step is the actual bottleneck — not embedding generation, which is
batched and comparatively cheap. The FAISS index and correct_llms map must be built from the
exact same row set: a retrieved neighbor with no correct_llms entry contributes an empty set
to the routing intersection, which kills that candidate and any k that includes it — so embedding
more rows than you benchmark buys nothing for the router.
Practical approach:
- Pick a target subsample size (e.g. 20-30k rows) and run
generate_embeddings.pyagainst the full input CSV as usual — sincemsr_cleaned_full.csv/train_full.csvwere already globally shuffled bycreate_balanced_dataset.py+split_dataset.py, walking the file in order and stopping partway already yields a roughly representative sample, not a biased chunk. - It's safe to stop the run at any point (
_checkpoint.json/_embeddings_cache.npyare written atomically after every batch) — Ctrl+C or a network drop loses nothing already completed. - Before committing to the LLM-benchmarking step, verify the partial run is actually
representative:
Reads only the checkpoint (no Ollama needed), maps
python src/dataset/check_embedding_distribution.py \ --input data/processed/full_dataset/split/train_full.csv \ --embeddings-dir embeddings_output/train_full
completed_indicesback tovul/cwe_idin the input CSV, and reports each class's count/percentage in the embedded-so-far sample vs. the full dataset — flagging any CWE class that's missing or thin (--min-count, default 20). If it looks representative, proceed as-is; if specific classes are thin, do a small targeted top-up (embed just those extra rows) rather than restarting. generate_embeddings.pyonly writes the finalembeddings.npy/rows.csv/row_ids.npyafter the entire input CSV finishes — so a deliberately capped partial run needs its own finalization step to become usable byemb_faiss.py/training_embedding_map.py(not yet built as of this writing — flag it if you need it before resuming further).
A second, parallel router that classifies which CWE type a vulnerable function contains
(e.g. CWE-119, CWE-416, ...), or NON_VULNERABLE if it isn't vulnerable, instead of the
binary vulnerable/not-vulnerable label. It reuses the same FAISS index, the same test/train
embeddings, and the same shrink-k-intersection routing logic as routing_pipeline.py — only
the label space, prompt, and per-model correctness map are different.
Label design: CWE-IDs among vulnerable (vul == 1) rows in msr_cleaned_full.csv are
counted; any CWE-ID below --min-count (default 50) rows is collapsed into a catch-all
OTHER class, and a synthetic NON_VULNERABLE class covers non-vulnerable rows. This keeps
the closed label set (and the few-shot prompt listing it) a manageable, fixed size regardless
of how many rare CWE-IDs exist in the raw data.
python src/dataset/build_cwe_label_map.py --min-count 50
# -> data/processed/cwe_label_map.csv, data/processed/cwe_descriptions.csvDescriptions are short, built-in summaries per CWE-ID (cwe_prompt_utils.CWE_DESCRIPTIONS),
with a generic MITRE-CWE-database fallback for any CWE-ID not in that table. These descriptions
are inserted into every prompt so the LLM has context on what each category id means.
python src/dataset/cwe_predictions.py --model codestral --output results_cwe/codestral.csv
python src/dataset/cwe_predictions.py --model qwen2.5-coder:14b --output results_cwe/qwen2.5_coder_14b.csv
python src/dataset/cwe_predictions.py --model qwen2.5-coder:32b --output results_cwe/qwen2.5_coder_32b.csv
python src/dataset/cwe_predictions.py --model llama3.1:8b --output results_cwe/llama31_8b.csv
python src/dataset/cwe_predictions.py --model codegemma:7b --output results_cwe/code_gemma_7b.csvDefaults to reading data/processed/full_dataset/split/train_full.csv. Each row's ground truth
is NON_VULNERABLE (vul==0) or its collapsed CWE label (vul==1); the model must answer with
exactly one label id from the closed set.
python src/dataset/cwe_llm_list_mapping.py # -> data/processed/sample_correct_llms_cwe.csvpython src/evaluation/routing_pipeline_cwe.py \
--k 5 \
--model_priority codegemma:7b llama3.1:8b qwen2.5-coder:14b qwen2.5-coder:32b codestral \
--out_dir results_cweDefaults to data/processed/full_dataset/split/test_full.csv and reuses the existing
embeddings_output/test_nochunk/embeddings.npy, train_faiss.index, and
train_sample_ids.npy from the binary pipeline (same underlying code rows, so the same
embeddings apply — no re-embedding needed). Outputs, per k: routing results, multi-class
accuracy + macro/weighted precision & recall (routing_metrics_cwe_k*.csv), a full confusion
matrix over every label (routing_confusion_matrix_cwe_k*.csv), and a per-class
precision/recall/F1 report (routing_classification_report_cwe_k*.csv).
pandas, numpy, scikit-learn, faiss (cpu or gpu build), requests, ollama (Python
client). No requirements.txt currently exists in the repo — add one if you want reproducible
installs.
sample_idformat isrow_{index:06d}(e.g.row_000042), assigned when a CSV is loaded during embedding generation (generate_embeddings.py:load_data). Several scripts (llm_list_mapping.py,routing_pipeline.py, and theircwe_*/*_cwecounterparts) re-normalize plain integer IDs into this format — keep this consistent if you add new data sources.- The CWE-type pipeline assumes
train_full.csv/test_full.csv(underdata/processed/full_dataset/split/) are the same rows, in the same order, as whatever is currently embedded inembeddings_output/train_nochunk/test_nochunk— it reuses those embeddings and the existingtrain_faiss.indexrather than re-embedding. If that split ever changes independently of what was embedded, re-rungenerate_embeddings.py+emb_faiss.pyfirst. evaluate_knn.pytargets an older directory layout (embeddings_output/train/withchunks.csv+faiss_index.bin+faiss_id_map.json) from an earlier "chunked embedding" design. The current pipeline embeds whole functions without chunking (*_nochunk/dirs), so this script will not run against the current data layout — treat it as historical reference, not a working entry point.- FAISS indexes use
IndexFlatIPover L2-normalized vectors (cosine similarity). - All data/embeddings/results directories are gitignored since they're large and machine/GPU-generated; regenerate them with the commands above rather than expecting them in version control.