Skip to content

Repository files navigation

atlas-kg

Python Neo4j FastAPI React TypeScript Vite Google Gemini Docker

A queryable knowledge graph of MITRE ATLAS — the adversarial threat taxonomy for AI systems.
Cross-layer traversal from threats to risk mapping to regulatory compliance, in one Cypher query.


What Is This?

MITRE ATLAS catalogs adversarial techniques against AI systems the same way ATT&CK does for traditional cybersecurity. But ATLAS data lives in a YAML file — there is no queryable graph, no way to trace an attack chain step by step, and no bridge to risk frameworks or regulations.

atlas-kg fixes that. It ingests the full ATLAS taxonomy (16 tactics, 170 techniques, 35 mitigations, 57 case studies) into a Neo4j graph and adds cross-layer bridges to the OWASP LLM Top 10 and EU AI Act. An agentic reasoning engine sits on top, investigating the graph at runtime and citing every claim back to the Cypher query that produced it.

The question no existing tool can answer: "Which techniques threaten my Generative AI system, what OWASP risks do they map to, what mitigations apply at deployment, and which EU AI Act articles do those satisfy?"

atlas-kg answers it in one graph traversal.


Features

  • Deterministic ingestion — the entire ATLAS YAML is parsed and loaded without any LLM in the pipeline. Reproducible, auditable, idempotent.
  • Attack chain modeling — case studies are not flat "used these techniques" lists. Each is modeled as a directed graph of AttackStep nodes linked by LEADS_TO edges, preserving the sequential order of real-world attack steps. No prior ATLAS tooling does this.
  • Cross-layer bridges — ATLAS techniques map to OWASP LLM Top 10 risks (MAPS_TO), and ATLAS mitigations map to EU AI Act articles (SATISFIES). One Cypher query traverses threats, risk categorization, and regulatory compliance.
  • Agentic reasoning engine — a graph-investigator agent with 6 generic tools decides at runtime which queries to run and in what order. No predefined query pipelines — the agent adapts to any question.
  • Citation tracing — every factual claim in the agent's answer carries a [Q1], [Q2] marker. The frontend renders these as clickable citations that expand to show the exact Cypher query and raw results.
  • SSE streaming — the API streams tool calls as server-sent events so the frontend shows the agent's investigation live, not just the final answer.
  • 8 pre-built example queries — coverage gaps, attack chains, cross-layer traversals, tactic breakdowns, and more.

Quick Start (Docker Compose)

The fastest way to run everything — Neo4j, backend, and frontend — with a single command. The graph is auto-populated on first startup.

git clone <repo> && cd atlas-kg
cp .env.example .env              # optionally add GEMINI_API_KEY for the reasoning engine
docker compose up --build

That's it. On first launch the backend waits for Neo4j, seeds the graph from the ATLAS YAML, then starts the API. Subsequent restarts skip seeding.

Service URL Description
Frontend localhost React app with query interface and citation UI
API localhost:8000 FastAPI reasoning engine endpoint
Neo4j Browser localhost:7474 Visual graph explorer and Cypher console

Local Development Setup

For development without Docker Compose.

Prerequisites

Requirement Version
Docker Latest (for Neo4j)
Python 3.12+
Node.js 18+
uv Latest

1. Start Neo4j

docker run --name atlas-kg -p 7474:7474 -p 7687:7687 \
  -e NEO4J_AUTH=neo4j/password -d neo4j:5

2. Clone and install

git clone <repo> && cd atlas-kg
uv sync
cp .env.example .env

3. Configure environment

Edit .env with your values:

NEO4J_URI=bolt://localhost:7687
NEO4J_USER=neo4j
NEO4J_PASSWORD=password
GEMINI_API_KEY=your-key-here    # only needed for reasoning engine

4. Load the graph

atlas-load

This runs the full pipeline: schema constraints → YAML parsing → node ingestion → relationship ingestion → validation. Output looks like:

Connecting to Neo4j at bolt://localhost:7687...
Connected.

1. Applying schema constraints and indexes...
2. Loading YAML data...
3. Ingesting nodes...
4. Ingesting relationships...
5. Validating graph...

Done in 4.2s

5. Run example queries

atlas-queries

Executes all 8 documented Cypher queries and prints results as formatted tables.

6. Start the reasoning engine (optional)

uvicorn backend.api.main:app --reload --port 8000

7. Start the frontend (optional)

cd frontend && npm install && npm run dev

Opens at localhost:5173.

Services (local dev)

Service URL Description
Neo4j Browser localhost:7474 Visual graph explorer and Cypher console
API localhost:8000 FastAPI reasoning engine endpoint
Frontend localhost:5173 React app with query interface and citation UI

Architecture

The system runs in three phases:

┌─────────────────────────────────────────────────────────────────────┐
│                         Phase 1: Ingestion                         │
│                                                                     │
│   ATLAS YAML ──▶ parse.py ──▶ Neo4j                                │
│   (v6.0.0)       Deterministic     (Tactics, Techniques,           │
│                  Two-pass for        Mitigations, CaseStudies,     │
│                  attack chains       AttackSteps + all edges)      │
├─────────────────────────────────────────────────────────────────────┤
│                     Phase 2: Bridge Injection                       │
│                                                                     │
│   config/*.json ──▶ bridges.py ──▶ Neo4j                           │
│   (OWASP mapping,    Cross-layer      (OWASPRisk, Regulatory-     │
│    EU AI Act)        bridge nodes      Article + MAPS_TO,          │
│                      and edges         SATISFIES edges)            │
├─────────────────────────────────────────────────────────────────────┤
│                      Phase 3: Reasoning                             │
│                                                                     │
│   User question ──▶ ADK Agent ──▶ Cypher queries ──▶ Cited answer  │
│                     (6 tools)     (4-8 per query)    with [Q] refs  │
│                     Self-discovers schema at runtime                │
└─────────────────────────────────────────────────────────────────────┘

Graph Model

graph TD
    subgraph "ATLAS Core"
        TAC[Tactic]
        TECH[Technique]
        SUB[Sub-technique]
        MIT[Mitigation]
        CS[CaseStudy]
        AS[AttackStep]
    end

    subgraph "Cross-Layer Bridges"
        OWASP[OWASPRisk]
        REG[RegulatoryArticle]
    end

    TECH -->|ACHIEVES| TAC
    SUB -->|SPECIALIZES| TECH
    MIT -->|MITIGATES| TECH
    AS -->|STEP_OF| CS
    AS -->|USES| TECH
    AS -->|IN_TACTIC| TAC
    AS -.->|LEADS_TO| AS
    TECH -->|MAPS_TO| OWASP
    MIT -->|SATISFIES| REG

    style TAC fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style TECH fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style SUB fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style MIT fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style CS fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style AS fill:#1a3a4a,stroke:#389DC6,color:#389DC6
    style OWASP fill:#2a1a3a,stroke:#886FBF,color:#886FBF
    style REG fill:#2a1a3a,stroke:#886FBF,color:#886FBF
Loading

The LEADS_TO edges between AttackStep nodes are the key modelling contribution — they encode the sequential order of attack steps within each case study, turning flat technique lists into traversable attack chain graphs.


Graph Schema

Node Types

Label Count Key Properties
Tactic 16 id, name, description
Technique 170 id, name, platforms[], maturity, is_subtechnique
Mitigation 35 id, name, lifecycle_phases[], categories[]
CaseStudy 57 id, name, type (Exercise / Incident), actor, target
AttackStep ~449 id, step_id, description
OWASPRisk 10 id (LLM01–LLM10), name, description
RegulatoryArticle 2 id (EU_AIA_Art9, EU_AIA_Art15), name, framework

Edge Types

Edge Direction Description
ACHIEVES Technique → Tactic A technique achieves a tactic goal
SPECIALIZES Sub-technique → Technique Sub-technique specializes a parent
MITIGATES Mitigation → Technique A mitigation counters a technique
STEP_OF AttackStep → CaseStudy A step belongs to a case study
USES AttackStep → Technique A step uses a specific technique
IN_TACTIC AttackStep → Tactic A step falls under a tactic
LEADS_TO AttackStep → AttackStep Sequential link in an attack chain
MAPS_TO Technique → OWASPRisk Cross-layer: threat → risk category
SATISFIES Mitigation → RegulatoryArticle Cross-layer: mitigation → regulation

Property Domains

Property Valid Values
platforms Generative AI · Agentic AI · Predictive AI · Enterprise
maturity Realized · Demonstrated · Feasible
lifecycle_phases Deployment · AI Model Engineering · Data Preparation · Business and Data Understanding · Monitoring and Maintenance · AI Model Evaluation
categories Technical - AI · Technical - Cyber · Policy

Reasoning Engine

The reasoning engine is a FastAPI backend + React frontend that exposes the graph through a graph-investigator agent powered by Google Gemini via the ADK framework.

How It Works

 User: "How secure is my RAG pipeline?"
  │
  ▼
 ┌──────────────────────────────────────────────────────────┐
 │  1. get_physical_schema()         → learns graph layout  │
 │  2. count_and_summarize()         → checks data scale    │
 │  3. read_neo4j_cypher()           → finds RAG techniques │
 │  4. get_node_neighbors()          → finds mitigations     │
 │  5. find_paths_between()          → traces to OWASP/reg  │
 │  6. read_neo4j_cypher()           → checks coverage gaps │
 │  7. finished()                    → delivers cited answer │
 └──────────────────────────────────────────────────────────┘
  │
  ▼
 Answer with [Q1]..[Q7] citation markers
 Each marker expands to show the exact Cypher + raw results

Every tool call is logged server-side with a reference ID. The frontend streams these as SSE events, showing the agent's investigation in real time — not just the final answer.

Agent Tools

Tool Purpose
get_physical_schema Self-discover all node labels, relationship types, and properties from the live graph
count_and_summarize Get node counts and property value distributions for a label
get_node_neighbors Traverse relationships from a specific entity (1–3 hops)
find_paths_between Find shortest paths connecting two entities across the graph
read_neo4j_cypher Execute any read-only Cypher query (with safety guards against writes)
finished Signal investigation complete and deliver the final cited answer

Project Structure

atlas-kg/
├── backend/
│   ├── api/
│   │   └── main.py                  # FastAPI SSE endpoint (/query, /health)
│   ├── core/
│   │   ├── config.py                # Settings + .env loading
│   │   └── db.py                    # Neo4j driver singleton
│   ├── ingestion/
│   │   ├── cli.py                   # CLI entrypoints (atlas-load, atlas-reset)
│   │   ├── parse.py                 # YAML → Neo4j ingestion (batched UNWIND)
│   │   ├── schema.py                # Neo4j constraints + indexes
│   │   └── validate.py              # Post-load validation + stats
│   ├── queries/
│   │   ├── example_queries.cypher   # 8 documented Cypher queries
│   │   └── runner.py                # Execute queries, print as rich tables
│   └── reasoning/
│       ├── agent.py                 # ADK agent definition + system prompt
│       ├── tools.py                 # 6 graph query primitives
│       └── citations.py             # Tool call interceptor + citation injection
│
├── frontend/                        # React 19 + TypeScript + Vite 8
│   ├── src/
│   │   ├── App.tsx                  # Main app (landing vs query view)
│   │   ├── components/              # Hero, Stats, ThreeLayer, QueryPage, ...
│   │   ├── hooks/
│   │   │   └── useSSEQuery.ts       # SSE streaming hook
│   │   └── styles/
│   │       └── tokens.css           # Design tokens
│   └── package.json
│
├── data/
│   └── ATLAS-2026.05.yaml           # MITRE ATLAS v6.0.0 source
│
├── config/
│   ├── owasp_atlas_mapping.json     # Technique → OWASP risk mappings
│   └── regulatory_mapping.json      # Mitigation → EU AI Act mappings
│
├── tests/
│   └── test_parse.py                # Unit tests for YAML parsing
│
├── docker-compose.yml               # Full-stack orchestration
├── Dockerfile.backend               # Multi-stage Python build
├── Dockerfile.frontend              # Multi-stage Node+nginx build
├── docker/
│   ├── entrypoint.sh                # Auto-seed Neo4j + start uvicorn
│   └── nginx.conf                   # Reverse proxy for API + static files
│
├── pyproject.toml                   # Python project config + CLI scripts
├── .env.example                     # Environment variable template
└── README.md

CLI Commands

All commands are registered as console scripts via pyproject.toml and available after uv sync:

Command Description
atlas-load Full pipeline: schema → parse → ingest → validate
atlas-reset Wipe Neo4j database and reload from scratch
atlas-queries Run all 8 example queries, print results as tables

Tests

pytest tests/

Data Source

All graph data is sourced from MITRE ATLAS YAML v6.0.0 (data/ATLAS-2026.05.yaml). The ingestion pipeline is fully deterministic — no LLMs are involved in parsing or loading.

Cross-layer mappings (OWASP LLM Top 10, EU AI Act Articles 9 & 15) are derived from published literature and stored as static JSON in config/.


Tech Stack

Layer Technology Role
Graph DB Neo4j 5 Native graph storage, Cypher queries
Backend Python 3.12, FastAPI Ingestion pipeline, SSE API
Agent Google ADK, Gemini Graph-investigator reasoning engine
Frontend React 19, TypeScript, Vite 8 Query interface with live citations
Infra Docker Compose Neo4j, backend, frontend orchestration
Testing pytest Unit and integration tests

Built for the Voyverse AI Governance Engineering assessment.

About

A queryable knowledge graph of MITRE ATLAS ,the adversarial threat taxonomy for AI systems.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages