AI Threat Modeling & Data Flow Architecture
Enterprise architectural blueprint for threat modeling modern Artificial Intelligence (AI), Retrieval-Augmented Generation (RAG), and agentic pipelines. Master the 7-component Data Flow Diagram (DFD), 6 critical Trust Boundaries (TB0 to TB5), STRIDE (Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, Elevation of Privilege) tailored for generative AI, DREAD risk scoring, and a production threat register.
AI System Data Flow & Trust Boundaries
Interactive 7-component data flow diagram with return paths and trust zones
Modern AI systems operate across distributed environments: user interfaces, API gateways, LLMs (Large Language Models), vector databases, and external tool execution servers. Establishing explicit trust boundaries at every seam ensures that untrusted data is sanitized, decisions are governed by code, and actions require verified authorization.
- 7-component end-to-end AI data flow pipeline
- 6 distinct trust boundary crossings (TB0 through TB6)
- Explicit separation between data, model context, and execution
- Return paths: retrieved documents and tool responses re-enter the model context
- The LLM is never an authorization enforcement point; checks live in application code
- TB0 supply chain controls for base models, weights, and third-party tools
AI System Data Flow: Runtime Path & Offline Ingestion Path
Complete architectural data flow showing the online request-response chain (with trust boundaries TB1–TB6), MCP server placement, the offline batch ingestion pipeline, and return paths
Runtime Path: User → Input → LLM → RAG & MCP → Memory → Output → Users
MCP (Model Context Protocol) Server Placement & Trust Boundary (TB3)
Where MCP Servers physically live, how they decouple model reasoning from system execution, and how data returns across TB3.
tools/calltools/list) & validates schemas<mcp_output> to stop Indirect Prompt InjectionOffline Path: Ingestion and Indexing
How an LLM Connects to Different RAG Systems
An LLM (Large Language Model) does not query data stores directly. A retrieval orchestrator sends each question to one or more RAG (Retrieval-Augmented Generation) backends, merges what comes back, and grounds the prompt in that context before the model writes an answer.
Request Lifecycle: Query to Grounded Answer
📚 5 Core Types of RAG Architecture (Naive, Advanced, Graph, Agentic, Modular)
Primary Interview FrameworkRetrieval-Augmented Generation (RAG) enhances LLMs (Large Language Models) by dynamically retrieving relevant facts from external knowledge repositories before generating responses. The five foundational architectures represent an evolution from naive vector lookups to autonomous multi-agent knowledge synthesis:
Optimized retrieval pipeline adding pre-retrieval query rewriting (e.g. HyDE — Hypothetical Document Embeddings, query decomposition), hybrid search (dense vectors + sparse BM25 — Best Matching 25), semantic chunking, and post-retrieval cross-encoder reranking and context compression. Still deterministic in execution order, but delivers dramatically higher precision and lower hallucination than Naive RAG.
The LLM (Large Language Model) acts as an autonomous agent orchestrating multi-step retrieval. Using tools and MCP (Model Context Protocol) servers, the agent evaluates incoming queries, formulates a multi-step search plan, queries heterogeneous data stores (vector DBs, SQL — Structured Query Language, APIs, live web), reflects on returned evidence, and iteratively loops until sufficient context is gathered.
Combines Knowledge Graphs (KGs with entities, relationships, and property attributes) with vector search. Instead of treating text as isolated chunks, Graph RAG extracts structured entity-relationship triples (e.g., Service-A → depends_on → Database-B) and executes multi-hop graph traversals (via Neo4j, Amazon Neptune) to answer relational and topology questions that flat vector search misses.
A composable, service-oriented architecture where routing, retrieval, memory, fusion, and verification are decoupled into interchangeable plug-and-play micro-modules. Systems dynamically assemble specialized pipelines based on query intent — routing simple questions to cached vectors and complex regulatory queries through hybrid graph-vector fusion with guardrails.
The standard 3-step baseline pipeline: (1) chunk and embed raw documents into vector embeddings during ingestion, (2) retrieve top-k chunks matching the raw user query via cosine similarity, and (3) concatenate retrieved chunks directly into the LLM system prompt. While straightforward and fast to prototype, it suffers from context stuffing, semantic mismatch, and zero self-correction.
Side-by-Side Comparison: The 5 RAG Architectural Styles
| Dimension | Naive / Classic | Advanced RAG | Graph RAG | Agentic RAG | Modular RAG |
|---|---|---|---|---|---|
| Retrieval Mechanism | Single-pass dense vector similarity (top-k cosine / dot product) | Hybrid search (dense vector + sparse BM25) + cross-encoder reranker | Graph traversal (Cypher / SPARQL) over entity-relationship triples + vector | Dynamic multi-step tool calls across vector, SQL, APIs, and live web | Dynamic composable routing across swappable retriever microservices |
| Query Pre-Processing | None (raw user string directly embedded) | Query rewriting, HyDE (Hypothetical Document Embeddings), sub-query decomposition | Named Entity Recognition (NER), entity linking, graph schema mapping | Multi-turn intent decomposition, tool selection, plan generation | Configurable per-route pre-processors (classifier, router, re-writer) |
| Self-Correction & Reflection | ❌ None (single forward pass) | ⚠️ Basic (confidence threshold filtering & context compression) | ⚠️ Schema validation & connected path verification | ✅ Full (evaluates chunk relevance, detects gaps, loops to re-fetch) | ✅ Configurable verification and guardrail modules |
| Multi-Hop Reasoning | ❌ Fails (cannot link disconnected document chunks) | ⚠️ Limited (relies on sub-query decomposition and reranking) | ✅ Superior (natively traverses entity relationship edges) | ✅ Strong (agent chains sequential queries across tools) | ✅ Supported via plug-in graph or agentic reasoning modules |
| Typical Latency | < 500 ms (single vector lookup + LLM completion) | 800 ms – 1.5 s (hybrid search + cross-encoder rerank) | 1.0 s – 2.5 s (entity extraction + graph query + vector) | 2.5 s – 8.0 s (multiple LLM reasoning steps & tool roundtrips) | Variable (500 ms – 4.0 s depending on routed pipeline) |
| Complexity & Maintenance | 🟢 Low (minimal infrastructure, 1 vector DB) | 🟡 Medium (vector DB, BM25 index, reranker model) | 🟠 High (graph DB maintenance, entity extraction pipelines) | 🔴 High (orchestrator loops, tool sandboxing, state tracking) | 🟣 Highest (microservice orchestration, interface versioning) |
| Primary Security & Threat Surface | Vector data poisoning, direct prompt injection | Query rewriting manipulation, poisoned chunks passing rerankers | Graph knowledge poisoning, malicious entity relation injection | Indirect Prompt Injection (OWASP LLM01), unauthorized tool execution (LLM06) | Inherits all threat planes; requires boundary validation at each module |
Explain the 5 types of RAG architectures — Naive, Advanced, Graph, Agentic, and Modular — and when would you use each?
Retrieval-Augmented Generation (RAG) grounds LLMs (Large Language Models) in external knowledge.
- Embeds documents into vectors, pulls top-k by cosine similarity, and stuffs them into the prompt
- It is fast to prototype but lacks query optimization, exhibits high irrelevance rates, and cannot handle multi-hop questions
- Use for simple FAQs
- The production standard
- Adds pre-retrieval query rewriting (e.g
- HyDE — Hypothetical Document Embeddings), hybrid search (dense vectors + sparse BM25 — Best Matching 25), and post-retrieval cross-encoder reranking and context compression
- Use for enterprise policy and knowledge-base search
- Extracts entity-relationship triples into Knowledge Graphs (KGs via Neo4j or Amazon Neptune)
- Enables multi-hop relational reasoning across connected entities that flat vector similarity misses
- Use for IAM (Identity and Access Management) permission graphs, IT dependencies, and regulatory lineage
- Treats the LLM as an autonomous agent that plans retrieval steps, queries multiple tools (SQL — Structured Query Language, APIs, vector DBs), evaluates retrieved context sufficiency, and loops dynamically
- Use for multi-system incident triage and deep research
- Decouples routing, search, and validation into plug-and-play microservices, routing simple queries to cache and complex queries to hybrid graph-vector fusion
- In terms of security, Naive RAG is vulnerable to vector data poisoning; Advanced adds query manipulation; Graph adds graph entity poisoning; Agentic introduces severe Indirect Prompt Injection (OWASP LLM01) and Excessive Agency (OWASP LLM06); and Modular inherits all risks across its distributed seams
🔄 LLM vs RAG vs Agentic RAG vs AI Agents vs Multi-Agent AI
Evolutionary SpectrumAs AI systems advance from standalone Large Language Models (LLMs) to multi-agent ecosystems, capability and adaptability increase alongside cost, latency, and attack surface.
| Aspect | 🧠 Base LLM | 📚 RAG | 🔗 Agentic RAG | 🤖 AI Agent | 🏢 Multi-Agent AI |
|---|---|---|---|---|---|
| Information Access | Static training weights only (frozen at pre-training cutoff date). | Dynamic external retrieval from vector DBs and enterprise document repositories. | Strategic multi-source retrieval — decides what, when, and how to query. | Direct read/write access to external tools, APIs, filesystems, and databases. | Collaborative, federated knowledge gathering across specialized agents. |
| Reasoning Capability | Statistical token prediction based on pre-trained patterns. | Context-augmented generation — grounds reasoning in retrieved chunks. | Goal-driven reasoning — plans retrieval steps, validates context sufficiency. | Autonomous goal decomposition, error recovery, and iterative execution. | Collective distributed reasoning with peer critique and consensus mechanisms. |
| Adaptability & State | Static (requires fine-tuning or retraining to learn new facts). | Moderately dynamic (updates instantly when new documents are indexed). | Highly dynamic (adapts retrieval strategy dynamically on runtime failures). | Highly adaptive (maintains short/long-term memory and execution history). | Extremely adaptive (dynamic role delegation and distributed state machines). |
| Problem-Solving Pattern | Single-turn prompt-in, completion-out text generation. | Retrieve-then-generate contextual synthesis. | Looping: Plan → Retrieve → Grade → Refine → Generate. | ReAct cycle: Observe → Orient → Decide → Act via external tools. | Orchestrated division of labor: Supervisor assigns subtasks to specialists. |
| External Interaction | None (pure text processing within context window). | Read-only queries to document indices and vector stores. | Dynamic read queries across vector DBs, SQL, search APIs, and MCP tools. | Full read/write execution (API mutations, bash scripts, database transactions). | Complex inter-agent messaging protocols and distributed workflow execution. |
| Cost & Latency | $ LOW · Single inference pass token cost only. | $$ MEDIUM · Embedding generation + vector DB query + LLM inference. | $$$ HIGH · Multiple LLM reasoning turns + multi-source retrieval calls. | $$$$ VERY HIGH · Multi-turn tool execution + recursive self-correction loops. | $$$$$ HIGHEST · Multiplied token consumption across parallel/sequential agents. |
| Security Risk Profile | Direct prompt injection (jailbreaks), sensitive data memorization leakage. | Corpus data poisoning, contextual leakage, cross-tenant vector contamination. | Indirect Prompt Injection (OWASP LLM01), retrieval query tampering. | Excessive Agency (OWASP LLM06), Server-Side Request Forgery (SSRF), tool abuse. | Cascading delegation hijacks, unlogged autonomous actions, privilege drift. |
How do you explain the transition from Traditional RAG to Agentic RAG and Autonomous AI Agents?
Traditional RAG (Retrieval-Augmented Generation) is a deterministic, linear data pipeline: the user submits a prompt, the system executes one vector similarity lookup, concatenates the top-k chunks into the context window, and the LLM (Large Language Model) generates a completion in a single turn. The model has zero control over what was retrieved or whether the data is sufficient. In AGENTIC RAG, the LLM becomes the active orchestrator: it decomposes complex queries into sub-tasks, selects appropriate retrieval tools (vector search, Text-to-SQL, internal APIs), inspects intermediate retrieved chunks for gaps or hallucinations, and autonomously loops to refine queries until confidence is established. When moving further to FULL AI AGENTS and MULTI-AGENT systems, the scope expands beyond read-only information retrieval into write-capable state mutations — executing API calls, writing and running code, interacting with third-party systems via Model Context Protocol (MCP), and coordinating across specialized agent personas. Crucially, each step up this ladder expands the security attack surface: while Traditional RAG faces data poisoning and retrieval evasion, Agentic and Multi-Agent systems face Indirect Prompt Injection (OWASP LLM01), Insecure Tool Execution (OWASP LLM06), Cascading Delegation Failures, and Denial of Wallet (OWASP LLM10) from runaway recursive loops.
🎛️ 7 RAG Build Patterns (Control Plane) × 8 Data Sources (Data Plane)
Production Control × Data MatrixProduction RAG (Retrieval-Augmented Generation) systems decouple control from data. The build pattern is the control plane (how and when retrieval loops trigger). The data source is the data plane (where structured, unstructured, or operational knowledge lives). Every production application joins one control pattern to one or more data sources through an intelligent retrieval orchestrator:
Connection Matrix: 7 Build Patterns × 8 Data Sources
Rows represent the 7 build patterns; columns represent the 8 enterprise data sources from the LLM Connection Map. Solid circles represent primary architectural fit; ringed circles represent secondary fit. Modular and Agentic RAG span the widest variety of sources because they natively support heterogeneous tool calling.
| Pattern ↓ / Source → | Vector Nearest-neighbor embeddings | Hybrid Keyword + vector | Graph Knowledge Graph | SQL Text-to-SQL | Tools / MCP Function calls | Web Live search | Multimodal Images, PDFs | Memory User history |
|---|---|---|---|---|---|---|---|---|
| Naive RAG Fixed pipeline | ||||||||
| Advanced RAG Pre + post retrieval | ||||||||
| Modular RAG Swappable parts | ||||||||
| Agentic RAG LLM plans retrieval | ||||||||
| Self-RAG Self-critique | ||||||||
| Corrective RAG CRAG | ||||||||
| Adaptive RAG Strategy per query |
Runtime Execution Pipeline: How Control Joins Data
Production Worked Recipes (Pattern + Sources)
How many types of RAG are there, and how do you choose the right one in production?
There is no single fixed number, because production RAG (Retrieval-Augmented Generation) is architecturally decoupled into two axes:
1The Control Plane (how retrieval loops execute): Naive, Advanced, Modular, Agentic, Self-RAG, Corrective RAG (CRAG), and Adaptive RAG.
2The Data Plane (where information lives): Vector embeddings, Hybrid (BM25 — Best Matching 25 + vector), Knowledge Graph (KG), Structured SQL (Structured Query Language), Live Web, MCP (Model Context Protocol) tool execution, Multimodal stores, and Memory buffers. A production system selects one control pattern and joins it to one or more sources via a router. My architectural default is ADVANCED RAG over HYBRID search with cross-encoder reranking. I introduce Graph RAG when relationships and multi-hop paths matter, SQL when exact numerical aggregation is required, and transition to AGENTIC or CORRECTIVE RAG (CRAG) only when offline evaluation using RAGAS (Retrieval-Augmented Generation Assessment) demonstrates that single-pass retrieval fails context precision or recall targets.
⚡ LangGraph vs CrewAI: Orchestration Frameworks for Agentic RAG
Runtime OrchestrationAgent frameworks govern the state machine, control flow loops, and tool execution boundaries of Agentic RAG (Retrieval-Augmented Generation). LangGraph and CrewAI represent the two dominant design paradigms:
- Nodes represent discrete execution steps (LLM call, tool invocation, custom Python function); edges define conditional branching and cyclic loops.
- A centralized, strongly typed state schema flows immutably through every node.
- Native checkpointing enables durable execution state persistence, time-travel debugging, and instant fault resumption.
- Built-in Human-in-the-Loop (HITL) interrupt hooks for approval gates prior to sensitive tool actions.
- Best for production systems: deterministic control, complex cyclic loops (Self-RAG, CRAG), and strict enterprise auditing.
- Agents are defined with explicit human-readable personas: Role, Goal, and Backstory (e.g. Lead Threat Researcher, Compliance Auditor).
- Tasks are assigned to specific agents and composed into a coordinated Crew.
- Supports sequential execution workflows or dynamic hierarchical delegation under a manager agent.
- Flows provide structured, event-driven state control for hybrid deterministic-agentic tasks.
- Best for rapid multi-agent prototyping, exploratory research, and workflows that decompose cleanly by organizational role.
Side-by-Side Comparison: LangGraph vs CrewAI
| Dimension | LangGraph | CrewAI |
|---|---|---|
| Mental Model | Flowchart / cyclical state machine | Collaborative team of role-based specialists |
| Control & Predictability | Fine-grained determinism; every node & edge explicit | High-level abstraction; autonomous task delegation |
| Learning Curve | Steeper (requires state schema & graph theory) | Gentler (intuitive prompt-like agent definitions) |
| State & Fault Tolerance | First-class database checkpointing & replay | Built-in memory buffers; Flows for structured state |
| Human Supervision | Native HITL interrupt primitives at any node | Supported via human-input task flags |
| Primary Enterprise Fit | High-auditability production RAG & banking workflows | Rapid multi-agent prototyping & creative exploration |
When would you choose LangGraph over CrewAI for an enterprise Agentic RAG deployment?
LangGraph models agent workflows as explicit, cyclical state machines where every node represents a deterministic computational or LLM (Large Language Model) step, edges enforce conditional routing, and a strongly typed state schema is checkpointed to persistent storage after each turn. This architecture provides replayability, deterministic audit logging, and first-class Human-in-the-Loop (HITL) interrupt primitives — essential requirements for high-assurance banking or healthcare systems running Corrective RAG (CRAG) or Self-RAG loops. In contrast, CrewAI abstracts agents around roles, goals, and backstories with autonomous hierarchical delegation, making it exceptionally fast to prototype multi-agent collaboration. In enterprise architecture, I use CrewAI to validate multi-agent role decomposition with product stakeholders, then implement the production pipeline in LangGraph to ensure strict deterministic boundary enforcement, scoped Model Context Protocol (MCP) tool access, and end-to-end cryptographic auditability.
🛡️ RAG Trust Boundaries: Why the LLM Is Never an Enforcement Point
Critical Defense ArchitectureIn any RAG (Retrieval-Augmented Generation) system, security failures occur when developers treat the LLM as an authorization or sanitization filter. Because LLMs are probabilistic token predictors, they cannot reliably enforce security boundaries. Enforcement must occur deterministically at the architectural boundaries:
The model produces text that becomes executable queries or tool calls. Enforce strict JSON-schema parameter validation, apply least-privilege API scopes tied to the authenticated user (never a shared superuser service account), and restrict tools with explicit allowlists.
Retrieved chunks and tool outputs flow directly back into the model context window. This return path is the #1 vector for Indirect Prompt Injection (OWASP LLM01). Enforce XML/delimiter encapsulation, dual-LLM verification, and document-level Access Control List (ACL) pre-filtering.
Writing conversation turns or extracted entities into long-term vector stores or session buffers. Enforce strict tenant namespace isolation, sanitize memory writes to prevent persistent memory poisoning, and apply cryptographic tenant tags to every stored vector embedding.
Where do you draw trust boundaries in a retrieval-augmented generation (RAG) agent, and why is the LLM not one of them?
A trust boundary is any seam in the architecture where identity, privilege, data sensitivity, ownership, or tenant contexts change. In a RAG (Retrieval-Augmented Generation) agent, we establish six explicit boundaries: TB-1 at public ingress (WAF — Web Application Firewall, OAuth rate limiting); TB-2 between input and context assembly (PII — Personally Identifiable Information redaction, delimiter tagging); TB-3 between the orchestrator and tool execution (least-privilege scopes, strict JSON schemas); TB-4 on the critical retrieval return path where retrieved documents re-enter the context window (document ACL — Access Control List filtering, dual-LLM indirect injection inspection); TB-5 at persistent memory storage (tenant-isolated namespaces, cryptographic context tagging); and TB-6 at action egress (Human-in-the-Loop approval for financial or data mutations). The LLM itself CANNOT be a trust boundary because it is a probabilistic token predictor, not a deterministic reference monitor. If an attacker injects an instruction into an ingested RAG document, system prompt instructions like 'ignore malicious commands' fail probabilistically. Therefore, all security policy enforcement — authentication, authorization, schema validation, and sanitization — must live in deterministic application code and infrastructure before data enters and after outputs leave the model.
Setting Up Trust Boundaries (TB0 through TB6)
A boundary is any place where the level of trust changes, and each one needs a written enforcement contract
Draw a boundary wherever any of these change between two components:
| Boundary | Crossing | Trust Change | Must Validate | Enforcement Controls |
|---|---|---|---|---|
| TB1 | User → Input Layer | Internet / untrusted → authenticated app | Who is calling, how much, how big | OAuth / OpenID Connect (OIDC) login at the gateway, Web Application Firewall (WAF), per-user rate limits and token quotas, input length caps |
| TB2 | Input Layer → LLM | User data mixes with trusted system instructions | Instructions vs data separation | Fixed system prompt, delimiting / "spotlighting" of user content, prompt-injection classifier, Personally Identifiable Information (PII) redaction before the model |
| TB3 | LLM ↔ RAG & Tools | Model text becomes executable tool calls; retrieved text flows back in | Is this call allowed for this user? Is returned content safe to read? | Tool allowlist, JSON-schema validation of arguments, authorization with the end user's identity (not a shared service account), treat retrieved docs and tool output as untrusted |
| TB4 | Tools → Memory / Data | Shared stores across users and tenants | Tenant isolation and write integrity | Per-tenant namespaces in the vector database, document-level access control (ACL) filtering at query time, validation on memory writes to block memory poisoning |
| TB5 | Memory → Output Layer | Internal data heads back toward a user | Is the output safe to show this user? | Output Data Loss Prevention (DLP), PII and secret filters, safe rendering (strip markdown image links used for exfiltration), grounding / citation checks |
| TB6 | Action Layer → External Systems | Agent acts in the real world with real credentials | Destination, scope, and approval | Egress allowlist, Server-Side Request Forgery (SSRF) protection, short-lived scoped credentials, Human-in-the-Loop (HITL) approval for destructive or financial actions |
| TB0 | Supply chain → Runtime | Third-party weights, datasets, and plugins enter production | Provenance and integrity | Model scanning, SafeTensors only, Cosign signature check at load, AI Bill of Materials (AIBOM), pinned plugin / MCP server versions |
AI Threat Modeling with STRIDE (STRIDE-AI)
Operational threat modeling framework adapted specifically for Generative AI, RAG, and Agentic systems
STRIDE (Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, Elevation of Privilege) provides the foundational risk taxonomy for assessing AI applications. By evaluating every system interface against the 6 threat classes, security teams establish verifiable controls across prompt ingestion, tool execution, and vector storage.
Applying STRIDE to AI Components
Each category with the AI-specific version of the threat, the questions to ask, and the control that works
- Prompt claims "I am the admin" and the model believes it
- Forged agent-to-agent messages or a fake tool / Model Context Protocol (MCP) server
- Poisoned document impersonating an authoritative source
- Can identity be faked?
- Can inputs be disguised?
- Can attackers impersonate users, systems, or data sources?
- Direct and indirect prompt injection
- Training-data or RAG-corpus poisoning
- Swapped or backdoored model weights
- Can data be changed?
- Can prompts or instructions be altered?
- Can the model, tools, or outputs be manipulated?
- Agent takes an action and nobody can prove who asked for it
- No record of which model version or prompt produced an answer
- Logs missing tool arguments and retrieved sources
- Are actions logged?
- Can users deny their requests?
- Is there enough audit information?
- System prompt extraction
- Cross-tenant leakage through shared vector stores or memory
- Training-data memorization and markdown-image exfiltration
- Can confidential data leak?
- Can prompts, system info, or retrieved data be exposed?
- Are outputs safe to share?
- Huge or recursive prompts that max out context and cost ("denial of wallet")
- Agent loops calling tools forever
- Flooding expensive embedding or GPU endpoints
- Can the system be overloaded?
- Can expensive prompts or loops drain resources?
- Are rate limits and quotas in place?
- Injected instructions make the agent call admin tools
- Tool runs with a broad service account instead of the user's rights
- Jailbreak bypasses role or policy restrictions
- Can a user access tools or data they should not?
- Can the model be tricked into privileged actions?
- Are permissions enforced outside the model?
STRIDE × Trust Boundary Heat Map
Severity of each STRIDE threat at each boundary: red = high risk, yellow = medium, gray = lower
| Trust Boundary | S | T | R | I | D | E | Top Risk to Address |
|---|---|---|---|---|---|---|---|
| TB1 User → Input | H | M | M | L | H | L | Account takeover, prompt flooding |
| TB2 Input → LLM | M | H | L | M | M | M | Direct prompt injection |
| TB3 LLM ↔ Tools | M | H | H | M | M | H | Indirect injection → privileged tool call |
| TB4 Tools → Memory | L | H | M | H | L | M | Cross-tenant leak, memory poisoning |
| TB5 Memory → Output | L | M | L | H | L | L | PII / secret leakage in responses |
| TB6 Action → External | M | M | H | H | M | H | Unauthorized email, payment, or data export |
How to Build an AI Threat Model
A repeatable 5-step method for any AI feature, product, or agent
Draw the data flow diagram (DFD): users, input layer, model, retrieval, tools, memory, outputs, and every external system. Include return paths (retrieved docs and tool results flowing back to the model) and build-time flows (model weights, datasets, plugins).
Mark every point where data enters, changes trust level, or leaves the system. Label each boundary (TB0-TB6) and write down what must be checked when data crosses it.
Walk each component and each boundary through all six STRIDE categories using the "ask yourself" questions. Add AI-specific threats from MITRE ATLAS (Adversarial Threat Landscape for AI Systems) and the OWASP Top 10 for Large Language Model (LLM) Applications.
Score each threat for impact and likelihood. DREAD (Damage, Reproducibility, Exploitability, Affected users, Discoverability) rates each factor 1-10 and averages them: 8+ is critical, 5-7 high, below 5 medium or low.
Assign a control and an owner to each high threat, add detections for what you cannot prevent, and re-run the model when you add a tool, data source, or model version.
DREAD for AI (DREAD AI): Quantitative Risk Scoring Framework
AI Risk QuantificationAdapting Damage, Reproducibility, Exploitability, Affected Users, and Discoverability for Non-Deterministic & Autonomous LLM Workloads
Traditional Common Vulnerability Scoring System (CVSS) often struggles with Artificial Intelligence (AI) and Large Language Model (LLM) architectures because attacks rely on semantic manipulation, non-deterministic model states, and tool-use privilege escalation rather than binary memory overflows. DREAD for AI (DREAD AI) scores each threat on a calibrated 1–10 scale across 5 dimensions, providing an empirical average to prioritize remediation backlogs.
(D + R + E + A + D) / 5Worked Example: Threat Register for a RAG Support Agent
What the output of an AI threat modeling session looks like
| ID | Boundary | STRIDE | Threat | DREAD | Mitigation | Owner |
|---|---|---|---|---|---|---|
| T-01 | TB3 | T / E | Hidden instructions in a retrieved support article tell the agent to issue a refund | 8.4 | Retrieved content tagged as data, refund tool requires user-scoped token + HITL above $50 | ML Platform |
| T-02 | TB4 | I | Customer A sees Customer B's tickets via a shared vector index | 8.0 | Per-tenant namespace + ACL filter applied before similarity search | Data Eng |
| T-03 | TB2 | I | User extracts the system prompt and internal API names | 6.2 | No secrets in prompt, canary token in prompt with SIEM alert on leak | AI AppSec |
| T-04 | TB5 | I | Model echoes a credit card number from an old ticket | 7.6 | PII redaction at ingest + output DLP on responses | Data Eng |
| T-05 | TB6 | R / E | Agent emails ticket exports to an external address with no audit trail | 7.8 | Recipient domain allowlist, signed action log with user ID and prompt hash | AI AppSec |
| T-06 | TB1 | D | Scripted 100k-token prompts run up inference cost | 6.0 | Per-user token quota, input cap, budget alert at 80% | ML Platform |
| T-07 | TB0 | T | Fine-tuned model replaced with a backdoored copy in the registry | 7.0 | Cosign signature verified at load, admission control blocks unsigned models | MLOps |
Interview Preparation: AI Threat Modeling
Questions on trust boundaries, STRIDE for AI, and prioritization
How do you threat-model an AI system?
I use five steps. First I map the system as a data flow diagram: users, the prompt and input layer, the large language model (LLM), retrieval and tools, memory, the output and action layer, and external systems. I include the return paths, because retrieved documents and tool results flow back into the model. Second I mark trust boundaries wherever identity, privilege, data sensitivity, ownership, execution context, or tenant changes, and give each one an ID. Third I apply STRIDE (Spoofing, Tampering, Repudiation, Information Disclosure, Denial of Service, Elevation of Privilege) to every component and boundary, then add AI-specific threats from MITRE ATLAS (Adversarial Threat Landscape for AI Systems) and the OWASP Top 10 for LLM Applications. Fourth I score threats with DREAD (Damage, Reproducibility, Exploitability, Affected users, Discoverability) and rank them. Fifth I assign a control and an owner to each high threat, add detections, and re-run the model whenever we add a tool, data source, or model version.
Give an example of each STRIDE category in an AI system.
Spoofing: a prompt injection says 'I am the admin' or a fake tool server impersonates an approved Model Context Protocol (MCP) server. Tampering: indirect prompt injection inside a retrieved document alters what the LLM instructs an agent to do. Repudiation: an agent takes an action like issuing a refund or deleting a file, but the logs only record the final HTTP call with no record of which prompt, model version, or user identity triggered it. Information Disclosure: a user extracts the system prompt and internal API keys via prompt injection, or a shared vector index leaks Customer A's data to Customer B. Denial of Service: scripted 100k-token prompts burn inference budget ('denial of wallet'), or an agent loop gets stuck calling tools recursively. Elevation of Privilege: an unprivileged user's prompt tricks the model into calling an admin-only tool like a refund API or database write.
How do you adapt the traditional DREAD methodology specifically to AI systems (DREAD AI)?
Traditional DREAD scores Damage, Reproducibility, Exploitability, Affected users, and Discoverability from 1 to 10. For AI, we calibrate each dimension to handle machine learning realities:
We score based on autonomous tool execution (e.g. issuing unauthorized refunds, dropping tables via MCP) and proprietary IP exfiltration.
- Because generative models are non-deterministic, we account for stochasticity
- A prompt jailbreak that requires 20 tries at high temperature scores lower (4-5), whereas a deterministic prefix attack or universal adversarial suffix (Greedy Coordinate Gradient / GCG) that reliably executes across temperatures scores an 8-9
Plain English natural language prompt injection scores high (9-10) because any user can type it without coding skills, while gradient-based evasion (PGD/FGSM) requiring white-box tensor access scores low (2-3).
Ephemeral in-session injection scores low (1-2), but poisoning a shared vector database (RAG) or backdooring a base model registry scores an 8-10 because it poisons the experience for all tenants.
- Default LangChain/LlamaIndex prompt templates or exposed OpenAPI specifications score high (8-9)
- We average the 5 dimensions; anything scoring 8.0 or above is a critical production blocker