Skip to content

SriyanRavuri/cve-research-assistant

Repository files navigation

CVE Research Assistant

A production-oriented RAG system over a CVE knowledge base. Ask a plain-language question; get back a grounded answer with inline citations, the supporting CVE snippets, and a per-claim hallucination check that compares the AI's claims against the retrieved source documents.

Designed to be auditable — every claim in every answer is verified against a source, and any unsupported claim raises a hallucination flag instead of being hidden.

Architecture

flowchart TD
    Q[User question] --> R[FAISS semantic retrieval<br/>top-k CVE docs]
    R --> P[Prompt assembly<br/>question + cited sources]
    P --> L[LLM<br/>grounded answer with inline citations]
    L --> A[Answer text + cited CVE ids]
    A --> H[Hallucination detector<br/>per-claim verification]
    R --> H
    H --> O[Final Answer<br/>text · citations · claim verifications<br/>has_hallucinations flag]

    style H fill:#fde68a,stroke:#92400e
    style O fill:#bbf7d0,stroke:#166534
Loading

How hallucination detection works

For every sentence in the model's answer:

  1. Extract any CVE-XXXX-XXXX identifiers cited in that sentence.
  2. If a cited CVE isn't in the retrieved set, flag the claim as unsupported (the model invented it).
  3. Otherwise, compute token overlap between the claim and the cited source document.
  4. If the overlap is below min_supporting_overlap, flag the claim as unsupported (the citation doesn't actually back up the statement).

Every claim ends up in claim_verifications with a boolean supported and the overlap score, so a reviewer can audit exactly which parts of the answer were grounded.

Quick start

git clone <this-repo>
cd cve-research-assistant
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

# Option A — offline fake LLM
USE_FAKE_LLM=1 python run_demo.py

# Option B — real OpenAI
cp .env.example .env  # fill in OPENAI_API_KEY
python run_demo.py

Tests

pytest -q

Onboarding for non-technical users

If you've never used a RAG system before, here's the mental model:

  1. The knowledge base is the source of truth. The assistant can only tell you things that appear in data/cves.jsonl. If you ask about something it's never seen, it should say so rather than guess.
  2. Citations are evidence, not decoration. Every [CVE-XXXX-XXXX] in an answer points to a specific source you can read yourself in the Citations: block printed below the answer.
  3. The hallucination check is your audit trail. Look at Claim verifications — every claim labelled X is something the system thinks may not be backed up by its sources. Treat those claims as needing manual confirmation.

A typical workflow:

1. You ask a question.
2. You read the answer.
3. You glance at the citations - do they look plausible?
4. You check the claim verifications - any 'X' marks?
   - If yes, treat the answer as a starting point only.
   - If all 'OK', the answer is grounded in the indexed CVE corpus.
5. For high-stakes decisions, you still verify against vendor advisories.

Demo walkthrough

run_demo.py runs four in-corpus questions plus one deliberate out-of-corpus question to demonstrate the refusal/grounding behaviour.

Q: What is Log4Shell and how do I mitigate it?
A: Apache Log4j2 versions 2.0-beta9 through 2.15.0 allow attackers to control
   log messages and inject JNDI lookups... [CVE-2021-44228]
Citations:
  - CVE-2021-44228: Apache Log4j2 versions 2.0-beta9 through 2.15.0...
Claim verifications:
  OK  (overlap=0.85) Apache Log4j2 versions 2.0-beta9 through 2.15.0...
Overall: all claims supported

Project layout

cve-research-assistant/
├── src/cve_assistant/
│   ├── config.py        # AssistantConfig
│   ├── models.py        # CVEDocument, Answer, Citation, ClaimVerification
│   ├── backends.py      # OpenAI + deterministic FakeLLM/FakeEmbeddings
│   ├── store.py         # FAISS-backed CVEStore + load_cves()
│   ├── hallucination.py # detect_hallucinations()
│   └── assistant.py     # CVEAssistant: retrieve + answer + verify
├── tests/               # 8 pytest tests
├── data/cves.jsonl      # 8 real notable CVEs (Log4Shell, Heartbleed, etc.)
├── run_demo.py
├── requirements.txt
└── LICENSE

Loading your own CVE corpus

Append JSONL records matching the CVEDocument schema:

{
  "cve_id": "CVE-2024-XXXXX",
  "title": "...",
  "description": "...",
  "affected_software": ["..."],
  "severity": "critical",
  "cvss_score": 9.8,
  "published": "2024-01-15",
  "references": ["https://..."]
}

Then re-run python run_demo.py — the index rebuilds from scratch each run.

Configuration

AssistantConfig (src/cve_assistant/config.py):

Field Default Meaning
top_k 4 Number of CVE docs retrieved per question
min_supporting_overlap 0.15 Lower-bound for claim/source token overlap
use_fake_llm auto True when no OPENAI_API_KEY

Tighten min_supporting_overlap to be more aggressive about flagging weakly-supported claims; loosen it to reduce false positives.

License

MIT — see LICENSE.

About

Auditable RAG system over a CVE knowledge base. Returns grounded answers with inline citations and runs a per-claim hallucination check that verifies every sentence against the retrieved source documents.

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages