Skip to content
Vasco1290Public

About

Domain-specialized RAG system for threat intelligence, ingesting MITRE ATT&CK & CVE datasets to serve source-attributed security answers via a streaming FastAPI & Next.js interface.

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

SecRAG — Cybersecurity Threat Intelligence RAG

Production-grade Retrieval-Augmented Generation for cybersecurity.
Hybrid retrieval · Cross-encoder reranking · Streaming LLM · Swappable providers

SecRAG ingests MITRE ATT&CK reports, CVE bulletins, and threat intelligence PDFs, then serves precise, source-attributed answers over a streaming API — complete with session memory, IOC extraction, and an evaluation framework.

SecRAG Next.js SOC Dashboard

CI Python FastAPI Next.js ChromaDB License: MIT


Features

Feature Details
Hybrid Retrieval Dense (MiniLM) + BM25 sparse, fused via Reciprocal Rank Fusion
Cross-Encoder Reranking BAAI/bge-reranker-base rescores top-20 → top-5 before the LLM
Session Memory Sliding-window conversation history (10-turn deque) per session_id
Provider Abstraction Swap Groq / OpenAI / Ollama via one LLM_PROVIDER env var
Async Ingest PDF upload returns job_id immediately; poll /ingest/status/{job_id}
Document Management GET /documents · DELETE /documents/{source}
Security NER CVE IDs, MITRE T-codes, IPv4, SHA-256/MD5 extracted at ingest
Streaming SSE Token-level streaming with source citations and latency breakdown
Eval Framework Recall@K, MRR, NDCG@K, Faithfulness — CI-gated thresholds
CI Pipeline ruff + pytest + docker build on every push

Architecture

┌──────────────────────────────────────────────────────────────────┐
│                         User Query                                │
│               (session_id → ConversationMemory)                   │
└───────────────────────────┬──────────────────────────────────────┘
                            │
                   ┌────────▼────────┐
                   │  FastAPI /chat  │  X-API-Key auth · SSE stream
                   └────────┬────────┘
                            │
         ┌──────────────────▼──────────────────┐
         │         Hybrid Retrieval             │
         │  ┌─────────────┐  ┌──────────────┐  │
         │  │ Dense (ANN) │  │ Sparse (BM25)│  │
         │  │  ChromaDB   │  │  rank-bm25   │  │
         │  └──────┬──────┘  └──────┬───────┘  │
         │         │   RRF (k=60)   │          │
         │         └───────┬────────┘          │
         │             top-20 candidates        │
         └──────────────────┬──────────────────┘
                            │
         ┌──────────────────▼──────────────────┐
         │     Cross-Encoder Reranking          │
         │     BAAI/bge-reranker-base           │
         │          → top-5 chunks              │
         └──────────────────┬──────────────────┘
                            │
         ┌──────────────────▼──────────────────┐
         │    LLM (Groq / OpenAI / Ollama)      │
         │    context + history + question       │
         │    streaming=True, temperature=0.1   │
         └──────────────────┬──────────────────┘
                            │
            SSE events: chunk · sources · metrics · done

Project Structure

SecRAG/
├── core/
│   ├── config.py           # Pydantic Settings (env var loading)
│   ├── parser.py           # PDF parsing + security NER
│   ├── chunker.py          # Sliding-window sentence-boundary chunker
│   ├── embedder.py         # SentenceTransformer wrapper (lazy-loaded)
│   ├── vector_store.py     # ChromaDB + BM25 + RRF hybrid retrieval
│   ├── reranker.py         # Cross-encoder reranking
│   ├── memory.py           # Session-scoped ConversationMemory
│   ├── providers.py        # LLM provider abstraction (Groq/OpenAI/Ollama)
│   └── engine.py           # End-to-end RAG orchestration
├── api/
│   ├── main.py             # App factory, lifespan, CORS
│   ├── dependencies.py     # Engine singleton + API key auth
│   ├── models.py           # Pydantic v2 schemas
│   └── routers/
│       ├── health.py       # GET  /health             [public]
│       ├── ingest.py       # POST /ingest             [protected, async]
│       │                   # GET  /ingest/status/{id} [protected]
│       ├── documents.py    # GET  /documents          [protected]
│       │                   # DELETE /documents/{src}  [protected]
│       └── chat.py         # POST /chat               [protected, SSE]
├── eval/
│   ├── golden_set.json     # 10 curated cybersecurity QA pairs
│   ├── metrics.py          # Recall@K, MRR, NDCG@K, Faithfulness
│   └── run_eval.py         # CLI eval runner with CI thresholds
├── frontend/               # Next.js 15 SOC dashboard
├── scripts/
│   └── seed_data.py        # Auto-seed MITRE ATT&CK + NVD CVEs
├── tests/
├── .github/workflows/ci.yml
├── CHANGELOG.md
├── CONTRIBUTING.md
└── docker-compose.yml

Quick Start

1. Clone & Configure

git clone https://github.com/Vasco1290/SecRAG.git
cd SecRAG
cp .env.example .env

Edit .env:

LLM_PROVIDER=groq                        # groq | openai | ollama
LLM_MODEL=llama-3.3-70b-versatile
GROQ_API_KEY=your_groq_api_key           # https://console.groq.com (free tier)
SECRAG_API_KEY=your_strong_random_secret

2. Install & Run

python -m venv .venv && source .venv/bin/activate  # Windows: .venv\Scripts\activate
pip install -r requirements.txt

# Terminal 1 — API
uvicorn api.main:app --reload --port 8000

# Terminal 2 — Frontend
cd frontend && npm install && npm run dev

Open http://localhost:3000

3. Seed the Knowledge Base

python scripts/seed_data.py

Downloads the MITRE ATT&CK Enterprise PDF and top-50 CRITICAL CVEs from NVD.


Switching LLM Providers

# Groq (default — fastest, free tier)
LLM_PROVIDER=groq
LLM_MODEL=llama-3.3-70b-versatile
GROQ_API_KEY=gsk_...

# OpenAI
LLM_PROVIDER=openai
LLM_MODEL=gpt-4o-mini
OPENAI_API_KEY=sk-...

# Ollama (local, no API key required)
LLM_PROVIDER=ollama
LLM_MODEL=llama3.2
OLLAMA_BASE_URL=http://localhost:11434

No code changes needed — just restart the API.


API Reference

All protected endpoints require the X-API-Key header. Full interactive docs at http://localhost:8000/docs.

POST /ingest — async PDF ingestion

curl -X POST http://localhost:8000/ingest \
  -H "X-API-Key: $KEY" -F "file=@report.pdf"
# → {"job_id": "abc-123", "status": "queued", "filename": "report.pdf"}

curl http://localhost:8000/ingest/status/abc-123 -H "X-API-Key: $KEY"
# → {"status": "done", "chunks_indexed": 87, "metadata": {...}}

GET /documents — list ingested sources

curl http://localhost:8000/documents -H "X-API-Key: $KEY"
# → {"documents": [{"source": "report.pdf", "chunk_count": 87}], "total_chunks": 87}

DELETE /documents/{source_name} — remove a document

curl -X DELETE "http://localhost:8000/documents/report.pdf" -H "X-API-Key: $KEY"
# → {"status": "deleted", "source": "report.pdf", "chunks_deleted": 87}

POST /chat — streaming RAG query (SSE)

curl -X POST http://localhost:8000/chat \
  -H "X-API-Key: $KEY" \
  -H "Content-Type: application/json" \
  -d '{"query": "What are APT29 initial access techniques?", "session_id": "analyst-1"}' \
  --no-buffer
event: chunk
data: APT29, also known as Cozy Bear...

event: sources
data: [{"doc": "MITRE_ATT&CK.pdf", "page": 42, "excerpt": "...", "iocs": {"cves": [], "mitre_techniques": ["T1566"]}}]

event: metrics
data: {"retrieval_ms": 87, "llm_ms": 1240, "total_ms": 1327, "chunks_used": 5}

event: done
data: [DONE]

Evaluation

# With a running API that has documents seeded:
python eval/run_eval.py --api-url http://localhost:8000 --api-key $KEY --k 5

# Write results to JSON:
python eval/run_eval.py --out eval_results.json

Sample output:

recall_at_k          0.4800   ✓
mrr                  0.3500   ✓
ndcg_at_k            0.4200   ✓
faithfulness         0.3600   ✓

Metrics are keyword-overlap based (no LLM judge dependency) and CI-gated — the runner exits with code 1 if any metric falls below its threshold.


Docker

cp .env.example .env  # fill in keys
docker compose up --build
docker compose exec api python scripts/seed_data.py

Configuration Reference

Variable Default Description
LLM_PROVIDER groq LLM backend: groq | openai | ollama
LLM_MODEL llama-3.3-70b-versatile Model name for the chosen provider
GROQ_API_KEY — Required when LLM_PROVIDER=groq
OPENAI_API_KEY — Required when LLM_PROVIDER=openai
OLLAMA_BASE_URL http://localhost:11434 Required when LLM_PROVIDER=ollama
SECRAG_API_KEY required API auth key (X-API-Key header)
EMBEDDING_MODEL sentence-transformers/all-MiniLM-L6-v2 HuggingFace embedding model
RERANKER_MODEL BAAI/bge-reranker-base Cross-encoder reranker
ENABLE_RERANKING true Toggle reranking stage
CHUNK_SIZE 512 Target tokens per chunk
CHUNK_OVERLAP 64 Overlap tokens between chunks
TOP_K_RETRIEVAL 20 Hybrid search candidate count
TOP_K_FINAL 5 Chunks passed to the LLM after reranking
NVD_API_KEY (optional) Raises NVD rate limit to 50 req/30s

License

MIT — see LICENSE.

About

Domain-specialized RAG system for threat intelligence, ingesting MITRE ATT&CK & CVE datasets to serve source-attributed security answers via a streaming FastAPI & Next.js interface.

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages