DEV Community

Cover image for Threat Modeling for AI Apps | AI Security series
Syed Mohammed Faham
Syed Mohammed Faham

Posted on

Threat Modeling for AI Apps | AI Security series

#ai #aiops #python #security

In the first post of this series, we explored why AI apps need security from the very beginning. Today, let’s dive into something more hands-on: threat modeling.

If you're not familiar with the term, think of threat modeling as the process of asking, “What can go wrong?” before your AI app is exposed to the real world. For AI systems, this means looking beyond traditional vulnerabilities and into the unique risks that come with using models, training data, and user prompts.

Why AI Apps Need a Different Lens

Threat modeling isn’t new. It’s been a common part of security practices for years. But when it comes to AI, we’re dealing with components that behave differently:

  • The model is dynamic and often unpredictable.
  • The data is unstructured and possibly user-generated.
  • The logic isn’t just written in code — it’s embedded in weights, embeddings, and training artifacts.

Because of this, traditional checklists won’t cut it. We need to tailor our threat models to the way AI systems behave.

A Simple Threat Modeling Framework for AI Apps
We don’t need a PhD or a 50-page doc to do threat modeling. A basic 4-step approach works well for most projects:

  1. Identify assets

    What are we trying to protect?

    • The LLM model itself (especially if it’s fine-tuned or proprietary)
    • API keys, secret prompts, and business logic
    • Training or evaluation data
    • User data or input/output logs

  2. Map the architecture

    Sketch out the AI stack. This could be something like:

    • A React or Gradio frontend
    • Backend in FastAPI or Node.js
    • Calls to an external LLM (OpenAI, Gemini, Mistral, etc.)
    • A vector database or a document store
    • Optional fine-tuned model or RAG pipeline

    This step helps us visualize where the weak points are.

  3. Enumerate threats

    Here’s where things get interesting. Ask questions like:

    • What if a user sends a malicious prompt?
    • What if someone tries to extract the model via repeated queries?
    • Could someone inject data during training or fine-tuning?
    • What happens if my API key leaks?

    Some AI-specific threats include:

    • Prompt injection
    • Jailbreaking
    • Model inversion
    • Data poisoning
    • Output manipulation (like leaking PII through summarization)

  4. Plan mitigations
    We won’t be able to stop everything — and that’s okay. Start with the most likely and most damaging risks.

    For example:

    • Sanitize user input before passing it to the model.
    • Limit token responses and set strict output formats.
    • Avoid logging full prompts and responses in plaintext.
    • Use auth and rate limiting on inference endpoints.
    • Randomize or mask data during training to prevent memorization.

Example: Conversational AI with FastAPI + Gemini

Let’s say you’ve built a chatbot using FastAPI and Gemini via OpenRouter. Here’s a basic threat model sketch:

Assets:

  • Prompt structure
  • User chat history
  • API key
  • Response payloads

Threats:

  • Prompt injection to bypass instructions
  • API key abuse, if exposed in frontend
  • Chat history leaking sensitive info
  • Model abuse via extreme prompts

Mitigations:

  • Move keys to the backend only
  • Add prompt pre-processing
  • Use token filtering on Gemini's output
  • Log only anonymized inputs

This isn’t rocket science, it’s just asking the right questions early on.

Tools You Can Use

You can model these threats manually, or use tools like:

Even a whiteboard and sticky notes will do the job if you’re in the early stages.

Final Thoughts

Threat modeling forces you to think like an attacker before the attacker shows up. For AI apps, that mindset is even more critical — because most of the time, the attack surface isn’t obvious until something goes wrong.

In the next post, we’ll get a little more tactical: how to handle training and inference data securely so you can stop worrying about leaks, poisoning, or accidental exposure.

Until then, take 30 minutes and try building a basic threat model for one of your AI projects. You might be surprised at what you find.

Connect & Share

I’m Faham — currently diving deep into AI and security while pursuing my Master’s at the University at Buffalo. Through this series, I’m sharing what I learn as I build real-world AI apps.

If you find this helpful, or have any questions, let’s connect on LinkedIn and X (formerly Twitter).

This is blog post #2 of the Security in AI series. Let's build AI that's not just smart, but safe and secure.
See you guys in the next blog.

Top comments (0)