Medical Device Threat Modeling Agent

Build a multi-agent AI tool that generates FDA-aligned STRIDE threat models.

Introduction

30 Second Summary

Every time a doctor checks your heart rhythm through a mobile app or a nurse monitors your vitals on a bedside dashboard, software is making medical decisions or handling sensitive health data. A single cybersecurity vulnerability in that software could put patients at risk.

In this project, you will build a command-line Python tool that uses three specialist AI agents to threat model medical devices. You will feed it a device description and get back a complete, structured threat model aligned with FDA cybersecurity guidance.

What You'll Build

Picture running a single command with a medical device description and watching three AI agents collaborate in real time to classify the device, question its security posture, and produce a threat register you could hand to a regulatory reviewer.

By the end of this project, you'll have:

  • A tool that takes any device scenario and classifies it as SaMD, SiMD, or MDDS, with the regulatory pathway and cybersecurity focus areas printed to your terminal.
  • Targeted probing questions that expose real attack surfaces, plus a full STRIDE threat register with severity ratings, mitigations, and FDA security control mappings.
  • A single-command pipeline powered by multi-agent orchestration, where one orchestrator agent automatically chains three specialist agents together.
  • Secret Mission: Run a comparative analysis across all three device types and generate a side-by-side threat report revealing how threat surfaces shift between classifications.

Are there any prerequisites?

This project assumes you have basic Python experience and some familiarity with cybersecurity concepts. You will also need Python 3.12, VS Code, and an OpenAI API key ready on your machine.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Your Development Environment

Before you can build any AI agents, you need the right foundation in place. Your threat modeling tool depends on the OpenAI Agents SDK to create and orchestrate specialist agents, and the SDK needs a valid API key to communicate with OpenAI's models.

In this step, you will create an isolated project workspace, install the SDK, and confirm everything connects correctly.

In this step, get ready to:
  • Create a project directory with a Python virtual environment.
  • Install the OpenAI Agents SDK and configure your API key.
  • Verify the SDK loads correctly with a quick import test.
Create a project directory and virtual environment
  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder, move into it, and set up a virtual environment by running these commands:
mkdir threat-modeler
cd threat-modeler
python3 -m venv .venv
  • Confirm the virtual environment was created by running:
ls

You should see a .venv folder listed in the output.

What is a virtual environment?

A virtual environment is an isolated Python workspace. Packages you install inside it stay separate from your system Python, so different projects can use different package versions without conflicts.

The .venv folder holds a dedicated copy of the Python interpreter and a clean package directory where your SDK will be installed.

  • Activate the virtual environment by running the command for your operating system:

macOS

  • Activate the environment by running:
source .venv/bin/activate

Windows

  • Activate the environment by running:
.venv\Scripts\activate

Your terminal prompt should now start with (.venv). This confirms you are working inside the virtual environment.

Don't see (.venv) in your prompt?

  • Make sure you are inside the threat-modeler folder. Run pwd to confirm your location.
  • On Windows, if you get a "running scripts is disabled" error, open PowerShell as Administrator and run Set-ExecutionPolicy RemoteSigned, then try the activate command again.
  • Check that the .venv folder exists by running ls. If it is missing, re-run the python3 -m venv .venv command.

Still stuck? Ask for help.

Install the SDK and set your API key

With your virtual environment active, any packages you install will stay isolated inside this project.

  • Install the OpenAI Agents SDK by running:
pip install openai-agents

What is the OpenAI Agents SDK?

The openai-agents package is OpenAI's framework for building multi-agent systems. It provides primitives like Agent, Runner, and function_tool that let you define specialist agents and wire them together into pipelines.

This single install also brings in pydantic and the openai client library as dependencies, so you do not need to install them separately.

You should see output scrolling through the terminal as pip downloads and installs the SDK and its dependencies. The final line will confirm the packages were installed successfully.

Installation failed?

  • Make sure your virtual environment is active. You should see (.venv) at the start of your terminal prompt.
  • Check your internet connection. pip needs to download packages from PyPI.
  • If you see a permissions error, confirm you are running inside the .venv environment. Installing into a virtual environment should never require administrator privileges.

Need more help debugging?

The SDK authenticates with OpenAI using an environment variable called OPENAI_API_KEY. You need to set this in your terminal session so the agents can reach the API.

  • Set your API key by running the command for your operating system. Replace your-api-key-here with your actual OpenAI API key.

macOS

export OPENAI_API_KEY=your-api-key-here

Windows

set OPENAI_API_KEY=your-api-key-here

Keep your API key safe

This export (or set) command only lasts for the current terminal session. If you close the terminal and reopen it, you will need to set the key again.

Never paste your API key directly into a Python file or commit it to version control.

Verify the SDK loads correctly

Everything is installed and configured. Before you move on, take a moment to predict: if you try to import the SDK's Agent class right now, do you think it will succeed or fail?

  • Test that the SDK loads by running this command:
python -c "from agents import Agent; print('SDK ready')"

You should see SDK ready printed in the terminal. This confirms that the SDK is installed, the import path works, and Python can find the agents package inside your virtual environment.

Seeing ModuleNotFoundError?

  • Confirm your virtual environment is still active. Look for (.venv) at the start of your prompt. If it is missing, reactivate with the command from the previous substep.
  • Try running pip install openai-agents again. If the package is already installed, pip will confirm that and skip the download.
  • Make sure you are using python (not python3) when running the test. Inside the virtual environment, python always points to the correct interpreter.

Still not working? Ask for help.

Your development environment is ready to go. Next up, you will build your first AI agent: a device classifier that determines whether a medical device is SaMD, SiMD, or MDDS.

Build the Device Classifier Agent

Your development environment is set up and the OpenAI Agents SDK is ready to go. Now it's time to build the first agent in your threat modeling pipeline.

Before you can threat model a medical device, you need to know what type it is. SaMD, SiMD, and MDDS each have fundamentally different threat surfaces and cybersecurity requirements, so your first agent handles this classification.

In this step, get ready to:
  • Define sample medical device scenarios for SaMD, SiMD, and MDDS.
  • Build a Device Classifier agent with FDA-aligned classification instructions.
  • Run the classifier on a sample cardiac arrhythmia app and examine its output.
Create the device scenarios and classifier agent

The FDA recognizes three broad classifications for medical device software. Each one changes which regulatory pathway a product follows and which cybersecurity risks matter most.

What are SaMD, SiMD, and MDDS?

The FDA uses a "verb test" to sort medical device software. If the software analyzes, interprets, predicts, or controls, it is likely SaMD (Software as a Medical Device). If it only transfers, stores, converts format, or displays data, it is likely MDDS (Medical Device Data System).

SiMD (Software in a Medical Device) is software embedded inside hardware. Think firmware in an insulin pump or a pacemaker. The hardware cannot achieve its medical purpose without the software.

  • In VS Code, create a new file called threat_modeler.py inside your project folder.
  • Add the imports and sample device scenarios by pasting this code into the file:
import asyncio
import sys

from agents import Agent, Runner


# Sample medical device descriptions, one for each classification type
SAMPLE_SCENARIOS = {
    "samd": (
        "Cloud-based cardiac arrhythmia detection application. "
        "The software receives ECG waveform data from wearable devices via a REST API, "
        "processes it using a machine learning model hosted on AWS, and sends alerts to "
        "clinicians when it detects atrial fibrillation or other arrhythmias. "
        "The application runs as a web service with a clinician dashboard."
    ),
    "simd": (
        "Firmware for a smart insulin infusion pump. "
        "The embedded software controls insulin delivery rates based on continuous glucose "
        "monitor (CGM) readings received via Bluetooth Low Energy. It has a touchscreen "
        "interface for manual overrides and connects to a hospital network for remote "
        "monitoring by healthcare providers. Firmware updates are delivered via USB."
    ),
    "mdds": (
        "Patient vital signs data aggregation and display system. "
        "The software collects readings from bedside monitors (heart rate, blood pressure, "
        "SpO2, temperature) via HL7 FHIR interfaces, stores them in a database, and "
        "displays them on a central nursing station dashboard. It does not analyze, "
        "interpret, or generate any alerts based on the data."
    ),
}

What does this code do?

  • Agent and Runner are the two core classes from the OpenAI Agents SDK. Agent defines what an agent knows and can do. Runner executes it.
  • SAMPLE_SCENARIOS is a dictionary with three device descriptions. Each one is designed to clearly fit one classification: SaMD (the arrhythmia app analyzes data), SiMD (the pump firmware is embedded in hardware), and MDDS (the vital signs display only stores and shows data).

Now you need the agent that will read these scenarios and classify each device. The classifier's instructions embed the FDA's verb test and the three classification definitions.

  • Add the classifier agent definition below the SAMPLE_SCENARIOS dictionary by pasting this code:
# Agent with instructions grounded in FDA classification criteria
classifier_agent = Agent(
    name="Device Classifier",
    instructions=(
        "You are a medical device regulatory classification expert. "
        "Given a description of a medical device or software system, classify it as one of:\n\n"
        "- SaMD (Software as a Medical Device): Standalone software that performs a medical "
        "purpose independently, without being part of a hardware medical device. Examples: "
        "cloud-based diagnostic apps, AI-powered imaging analysis, clinical decision support.\n\n"
        "- SiMD (Software in a Medical Device): Software embedded in or integral to a hardware "
        "medical device. The hardware cannot achieve its medical purpose without the software. "
        "Examples: infusion pump firmware, pacemaker software, MRI image reconstruction.\n\n"
        "- MDDS (Medical Device Data System): Software that transfers, stores, converts format, "
        "or displays medical device data WITHOUT analyzing, interpreting, or alarming. "
        "Examples: patient data display systems, FHIR data pass-through, secondary ECG displays.\n\n"
        "Use the 'verb test': if the software analyzes, interprets, predicts, triages, "
        "prioritizes, alarms, or controls, it is likely SaMD. If it only transfers, stores, "
        "converts format, or displays, it is likely MDDS. If it is embedded in hardware, "
        "it is SiMD.\n\n"
        "Provide your classification with clear rationale, the likely FDA regulatory pathway, "
        "and the key cybersecurity focus areas specific to that classification."
    ),
)

What does this agent definition do?

  • name gives the agent an identity the SDK uses for tracing and logging.
  • instructions is the system prompt. It contains the full SaMD/SiMD/MDDS definitions and the verb test so the agent can classify devices without any extra context.
  • The agent has no output_type yet, so it returns plain text. You'll add structured output in a later step.
  • Save threat_modeler.py.
  • Verify the file loads without errors by running this quick import check:
python -c "from threat_modeler import classifier_agent; print(classifier_agent.name)"

You should see Device Classifier printed in your terminal.

Seeing a syntax error?

  • Check that you have the closing parenthesis ) on the last line of the Agent(...) call. A missing parenthesis is the most common cause.
  • Make sure your virtual environment is still active. You should see (.venv) at the start of your terminal prompt.
  • Help me debug this import error.
Add a main function to run the classifier

The agent definition tells the SDK what the classifier knows. To actually run it, you need a main() function that sends a device scenario to the agent and prints its response.

  • Add the main function at the bottom of threat_modeler.py by pasting this code below the classifier_agent definition:
# Run the classifier on the SaMD cardiac arrhythmia scenario
async def main() -> None:
    scenario = SAMPLE_SCENARIOS["samd"]
    print(f"Analyzing: SAMD")
    print(f"Description: {scenario}\n")

    # Send the scenario to the classifier agent and print the result
    result = await Runner.run(classifier_agent, scenario)
    print(result.final_output)


if __name__ == "__main__":
    asyncio.run(main())

What does this code do?

  • Runner.run() is an async method that sends the scenario text to the agent and waits for its response. The SDK handles the API call to the default model (gpt-5.4-mini) behind the scenes.
  • result.final_output contains the agent's response text. Since there's no structured output_type yet, this is plain text the model generated.
  • asyncio.run(main()) is the standard Python pattern for running an async function from a synchronous entry point.
  • Save threat_modeler.py.
Run the script and examine the classification output

The sample scenario describes a cloud-based app that receives ECG data from wearables and sends arrhythmia alerts to clinicians. Before you run the classifier, think about which classification that should receive.

Based on what you know about the verb test, do you think this cardiac arrhythmia detection app will be classified as SaMD, SiMD, or MDDS? What verbs in the description give it away?

  • Run the classifier by executing this command:
python threat_modeler.py

The classifier should identify the device as SaMD. The key verbs are "detects" and "sends alerts." The app independently analyzes ECG data and generates clinical alerts, so it performs a medical purpose without being part of hardware.

You should also see a regulatory pathway (likely De Novo or 510(k)) and a list of cybersecurity focus areas such as API security, cloud infrastructure, and data protection.

Seeing an authentication error?

  • Make sure your OPENAI_API_KEY environment variable is still set in the current terminal session. If you opened a new terminal, you need to export the key again.
  • Confirm your virtual environment is active. You should see (.venv) in your terminal prompt.
  • Help me fix this API authentication error when running my threat modeler.

✔️ Awesome, I've got everything!

Great. Double check you've saved threat_modeler.py and move on to the next step.

ⓧ I'd like to double check the full code

Your threat_modeler.py should contain these three sections in order. Scroll up and confirm each is present:

  • Imports (asyncio, sys, Agent, Runner) followed by the SAMPLE_SCENARIOS dictionary with three device descriptions.
  • classifier_agent definition using Agent() with a name of "Device Classifier" and long instructions covering all three classifications.
  • An async def main() function that runs the classifier on the "samd" scenario and prints the result, plus the if __name__ == "__main__" block.

Your classifier can determine what type of medical device you're dealing with. Next, you'll build a second agent that generates targeted security questions to expose the real attack surface.

Add the Probing Questions Agent

In the previous step, your classifier agent identified the cardiac arrhythmia app as SaMD. That classification tells you what the device is, but not where the security gaps are.

A threat modeler needs targeted questions that probe the real attack surface. In this step, you will build a second agent that generates classification-specific cybersecurity questions and manually chain it to the classifier. This manual handoff will work, but pay attention to how you are threading data between agents.

In this step, get ready to:
  • Add a probing questions agent with classification-specific security question instructions.
  • Chain the two agents together using manual text concatenation.
  • Observe how unstructured text handoffs between agents introduce fragility.
Add the probing questions agent definition

The classifier answers "what type of device is this?" The probing questions agent answers a different question: "what should a threat modeler investigate about this device?" Its instructions are tailored to each classification type, asking different questions for SaMD versus SiMD versus MDDS.

  • In threat_modeler.py, find the closing parenthesis of the classifier_agent definition.
  • Add the probing questions agent definition below it by pasting this code after the classifier agent's closing line:


probing_agent = Agent(
    name="Probing Questions Generator",
    instructions=(
        "You are a cybersecurity threat modeling expert specializing in medical devices.\n\n"
        "Given a device description and its classification (SaMD, SiMD, or MDDS), generate "
        "targeted probing questions that a threat modeler should ask to identify security risks.\n\n"
        "Tailor your questions to the device classification:\n\n"
        "For SaMD: focus on API security, cloud infrastructure, data protection, software "
        "updates, third-party components, and data flow trust boundaries.\n\n"
        "For SiMD: focus on secure boot, firmware integrity, physical access controls, "
        "network interface hardening, safe failure modes, and hardware-software integration.\n\n"
        "For MDDS: focus on data integrity during transfer and storage, access control, "
        "audit logging, interoperability security, and scope creep prevention.\n\n"
        "Always include questions about supply chain security, SBOM considerations, "
        "vulnerability disclosure processes, and environment assumptions. Per FDA guidance, "
        "assume hospital networks are hostile.\n\n"
        "Generate 8-10 specific questions. Each should explain why it matters for the threat model."
    ),
)

What does this agent do?

The probing agent's instructions are split into three classification-specific blocks. Each block tells the agent which security areas to focus on for that device type.

  • For SaMD, the focus is cloud and API risks (REST endpoints, ML model hosting, data in transit).
  • For SiMD, the focus shifts to firmware and physical risks (secure boot, Bluetooth, USB updates).
  • For MDDS, the focus is data integrity and scope creep prevention (making sure the system stays "display only").
  • The hostile network assumption comes directly from FDA cybersecurity guidance, which recommends manufacturers assume adversaries control the network.
  • Save threat_modeler.py.
Update main() to chain both agents with manual handoff

Right now, main() only runs the classifier. You need it to run the classifier first, capture its text output, and then pass that text into the probing agent as input.

  • In threat_modeler.py, find the existing main() function. It currently looks like this:
async def main() -> None:
    scenario = SAMPLE_SCENARIOS["samd"]
    print(f"Analyzing: SAMD")
    print(f"Description: {scenario}\n")

    result = await Runner.run(classifier_agent, scenario)
    print(result.final_output)
  • Replace the entire main() function with this updated version that runs both agents:
async def main() -> None:
    scenario = SAMPLE_SCENARIOS["samd"]
    print(f"Analyzing: SAMD")
    print(f"Description: {scenario}\n")

    # Run the classifier agent first
    classifier_result = await Runner.run(classifier_agent, scenario)
    classification_text = classifier_result.final_output
    print("DEVICE CLASSIFICATION:")
    print(classification_text)
    print()

    # Manually concatenate the classifier's text output into a prompt for the probing agent
    probing_input = f"Device description: {scenario}\n\nClassification result: {classification_text}"
    probing_result = await Runner.run(probing_agent, probing_input)
    print("PROBING QUESTIONS:")
    print(probing_result.final_output)

What does this code do?

  • classifier_result.final_output captures the classifier's plain-text response into a variable called classification_text.
  • probing_input is built by manually concatenating the original scenario description with the classifier's text output using an f-string.
  • The second Runner.run() call passes that concatenated string to the probing agent. The probing agent has no idea what format the classifier used. It just receives a blob of text.
  • Save threat_modeler.py.

Seeing a syntax error on save?

  • Make sure you replaced the entire main() function, including the async def main() -> None: line. Leaving the old function header alongside the new body causes indentation errors.
  • Check that the probing_agent definition sits above main(), not inside it. Python needs the agent defined before main() references it.

Help me fix a syntax error in my threat_modeler.py after adding the probing agent and updating main()

✔️ Awesome, I've got everything!

Make sure you have saved threat_modeler.py with both the new probing_agent definition and the updated main() function.

ⓧ I'd like to double check the full code

Your threat_modeler.py should contain these sections in order. Scroll up and confirm each one is present:

  • Imports: asyncio, sys, and from agents import Agent, Runner.
  • SAMPLE_SCENARIOS dict: three entries for samd, simd, and mdds (unchanged from Step 2).
  • classifier_agent: the Device Classifier agent definition (unchanged from Step 2).
  • probing_agent: the Probing Questions Generator agent definition you added in this step, placed directly after the classifier agent.
  • main() function: the updated version that runs both agents sequentially, with probing_input built by f-string concatenation.
  • if __name__ block: asyncio.run(main()) at the bottom (unchanged).
Run the pipeline and observe the text handoff

You now have two agents wired together. Before you run the script, think about what the probing agent will receive as input. It gets a raw string you assembled with an f-string. What happens if the classifier's output format changes between runs?

  • Run the two-agent pipeline by running this command:
python threat_modeler.py

You should see two sections printed in your terminal. First, a DEVICE CLASSIFICATION section with the SaMD classification. Then, a PROBING QUESTIONS section with 8-10 numbered security questions tailored to the SaMD classification.

What should I see?

The classification output should identify the cardiac arrhythmia app as SaMD with rationale about its diagnostic function. The probing questions should focus on SaMD-specific concerns: REST API authentication, cloud infrastructure security, ML model integrity, and data-in-transit encryption.

Each question should include a short explanation of why it matters for the threat model. The agent may take 15-30 seconds to produce both outputs since it makes two separate LLM calls.

Seeing an API error or timeout?

  • Confirm your OPENAI_API_KEY environment variable is still set in your current terminal session. If you opened a new terminal, you need to activate your virtual environment and re-export the key.
  • If you see a rate limit error, wait 30 seconds and try again. Two back-to-back LLM calls can occasionally hit rate limits on lower-tier API plans.

Help me debug an API error when running my two-agent threat modeling pipeline

The pipeline works, but look at the line that builds probing_input. You are manually gluing the scenario text and the classifier's output together with an f-string. There is no structure, no schema, and no guarantee about what fields the classifier included.

Why is manual text handoff a problem?

The classifier's final_output is plain text. One run might label the classification as "SaMD" while the next uses "Software as a Medical Device". The probing agent has no way to reliably extract the classification type from that free-form string.

When you add a third agent in the next step, the problem compounds. Each agent's output feeds into the next, and any inconsistency in format cascades through the whole pipeline. Structured output with typed data models solves this.

Your two-agent pipeline produces real, classification-specific security questions. Next up, you will add typed data models with Pydantic and build the STRIDE threat analyst. Structured output replaces the fragile string glue you just saw with typed contracts that every agent can rely on.

Build the STRIDE Analyst with Structured Output

In the last step, you saw the fragility of passing raw text between agents. The classifier returned free-form prose, and you had to manually concatenate it into a prompt string for the probing agent. If the classifier's format ever changes, that string concatenation silently breaks.

In this step, you'll replace those unstructured text handoffs with typed Pydantic data models. Each agent will return a structured object with named fields instead of a wall of text. You'll also build the third and final specialist: a STRIDE threat analyst that produces a full threat register.

In this step, get ready to:
  • Create models.py with Pydantic data models for classification, probing questions, and the threat register.
  • Update all three agents to return structured output using output_type.
  • Add the STRIDE analyst agent and run the full three-agent pipeline.
Create the Pydantic data models

Each agent in the pipeline needs a typed contract defining the shape of its output. Instead of hoping the LLM returns text in a format the next agent can parse, a Pydantic model enforces exact field names and types.

  • Create a new file called models.py in your project directory.
  • Add the classification and probing question models by pasting this code:
from pydantic import BaseModel


# Structured output for the device classifier agent
class DeviceClassification(BaseModel):
    device_name: str
    classification: str
    rationale: str
    regulatory_pathway: str
    cybersecurity_focus_areas: list[str]


# Individual probing question with context
class ProbingQuestion(BaseModel):
    category: str
    question: str
    why_it_matters: str


# Full output for the probing questions agent
class ProbingQuestionsOutput(BaseModel):
    device_name: str
    classification: str
    questions: list[ProbingQuestion]

What do these models define?

DeviceClassification captures what the classifier agent returns: the device name, its SaMD/SiMD/MDDS classification, the reasoning, the regulatory pathway, and the cybersecurity focus areas.

ProbingQuestion holds a single security question, while ProbingQuestionsOutput wraps a list of them into the probing agent's full output. The questions field uses the ProbingQuestion model as a nested type, so Pydantic validates each item in the list.

  • Add the threat register models below the code you just pasted in models.py:
# Single threat entry in the STRIDE register
class Threat(BaseModel):
    stride_category: str
    threat_description: str
    affected_component: str
    severity: str
    likelihood: str
    mitigation: str
    fda_alignment_note: str


# Full threat model output for the STRIDE analyst agent
class ThreatModel(BaseModel):
    device_name: str
    classification: str
    threats: list[Threat]
    assumptions: list[str]
    summary: str

What does the Threat model capture?

Each Threat maps to one row in a STRIDE threat register. The fda_alignment_note field connects each threat to one of the FDA's eight security control categories (Authentication, Authorization, Cryptography, and so on).

ThreatModel wraps the full analysis: a list of threats, the environmental assumptions the agent made, and an overall summary. This is the structured output the STRIDE analyst agent will return.

  • Save models.py.
  • Confirm the file exists by running:
ls models.py

You should see models.py listed in the output.

Don't see models.py?

Make sure you created the file in your project directory (the same folder as threat_modeler.py). Check that your terminal is in that directory by running ls and confirming you see threat_modeler.py listed too.

help me find where models.py was created

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

ⓧ I'd like to double check the full code

Your models.py should contain these sections in order. Scroll up and confirm each one is present:

  • The from pydantic import BaseModel import at the top.
  • DeviceClassification class with 5 fields (device_name, classification, rationale, regulatory_pathway, cybersecurity_focus_areas).
  • ProbingQuestion class with 3 fields (category, question, why_it_matters).
  • ProbingQuestionsOutput class with 3 fields, including a questions: list[ProbingQuestion] field.
  • Threat class with 7 fields (stride_category through fda_alignment_note).
  • ThreatModel class with 5 fields (device_name, classification, threats, assumptions, summary).
Update the agents to use structured output

Now you need to tell each agent which Pydantic model to return. The OpenAI Agents SDK's output_type parameter handles this. When you set it, the SDK forces the LLM to return data matching that model's schema exactly.

  • In threat_modeler.py, find this import line near the top:
from agents import Agent, Runner
  • Add the models import directly below it, so the two import blocks look like this:
from agents import Agent, Runner

from models import (
    DeviceClassification,
    ProbingQuestionsOutput,
    ThreatModel,
)
  • Scroll down to the classifier_agent definition. Find the closing parenthesis of the Agent() call. It currently looks like this:
        "and the key cybersecurity focus areas specific to that classification."
    ),
)
  • Add output_type=DeviceClassification, before the closing ), so the end of the classifier definition now looks like this:
        "and the key cybersecurity focus areas specific to that classification."
    ),
    output_type=DeviceClassification,
)
  • Do the same for probing_agent. Find its closing parenthesis, which currently looks like this:
        "Generate 8-10 specific questions. Each should explain why it matters for the threat model."
    ),
)
  • Add output_type=ProbingQuestionsOutput, before the closing ):
        "Generate 8-10 specific questions. Each should explain why it matters for the threat model."
    ),
    output_type=ProbingQuestionsOutput,
)
  • Save threat_modeler.py.
  • Run the script to see how structured output changes what the agents return:
python threat_modeler.py

The classification output now shows named fields like device_name= and classification= instead of free-form prose. The probing questions output looks similar, with each question's fields visible as structured data.

What changed?

By setting output_type, the SDK tells the LLM to return structured data matching your Pydantic model's schema. The final_output is now a Python object with typed fields instead of a raw text string. This means downstream agents can access fields like classification.regulatory_pathway instead of parsing free text.

Seeing an import error?

If you see ModuleNotFoundError: No module named 'models', make sure models.py is in the same folder as threat_modeler.py. Also check that the filename is exactly models.py with no extra characters.

help me fix my models import error

Add the STRIDE analyst and run the full pipeline

The classifier tells you what the device is. The probing agent exposes where to look for vulnerabilities. Now you need an agent that turns all of that context into a concrete threat register using the STRIDE methodology.

  • In threat_modeler.py, add the STRIDE analyst agent definition directly below the probing_agent block:
stride_agent = Agent(
    name="STRIDE Threat Analyst",
    instructions=(
        "You are a STRIDE threat modeling expert for medical devices, aligned with "
        "FDA cybersecurity guidance (June 2025) and AAMI TIR57.\n\n"
        "Given a device description, its classification, and probing questions, "
        "perform a comprehensive STRIDE analysis. Assume reasonable default answers to "
        "the probing questions based on common architectures for the device type.\n\n"
        "STRIDE categories:\n"
        "- Spoofing: Can an attacker impersonate a legitimate user, device, or system?\n"
        "- Tampering: Can data or code be modified without detection?\n"
        "- Repudiation: Can actions be denied without proper logging?\n"
        "- Information Disclosure: Can sensitive data be exposed?\n"
        "- Denial of Service: Can the device or system be made unavailable?\n"
        "- Elevation of Privilege: Can an attacker gain unauthorized access levels?\n\n"
        "For each threat, provide:\n"
        "1. STRIDE category\n"
        "2. Concrete threat description specific to this device\n"
        "3. The affected component or data flow\n"
        "4. Severity (Critical/High/Medium/Low) considering patient safety impact\n"
        "5. Likelihood (High/Medium/Low) based on attack surface\n"
        "6. A specific mitigation\n"
        "7. FDA alignment: which FDA security control category this maps to "
        "(Authentication, Authorization, Cryptography, Code/Data/Execution Integrity, "
        "Confidentiality, Event Detection and Logging, Resiliency and Recovery, or "
        "Updatability and Patchability)\n\n"
        "Generate at least 6 threats covering all STRIDE categories. State your "
        "environment assumptions. Per FDA guidance, assume hospital networks are hostile.\n\n"
        "Provide a summary of the overall threat landscape."
    ),
    output_type=ThreatModel,
)

What makes this agent's instructions special?

The STRIDE analyst's instructions are grounded in authoritative sources. They reference the FDA's June 2025 cybersecurity guidance and AAMI TIR57, the standard for medical device security risk management. The agent is told to assume hospital networks are hostile, which comes directly from the FDA's recommended threat assumptions.

The instructions also require each threat to map to one of the FDA's eight security control categories. This is what makes the output professionally useful: a real threat modeler could hand this register to a regulatory reviewer and it would speak their language.

Now you need to update main() so it runs all three agents in sequence, using structured output at every handoff.

  • Replace your entire main() function with this updated version:
async def main() -> None:
    scenario = SAMPLE_SCENARIOS["samd"]
    print(f"Analyzing: SAMD")
    print(f"Description: {scenario}\n")

    # Run the classifier and get structured output
    classifier_result = await Runner.run(classifier_agent, scenario)
    classification: DeviceClassification = classifier_result.final_output
    print("DEVICE CLASSIFICATION:")
    print(classification.model_dump_json(indent=2))
    print()

    # Pass structured classification to the probing agent
    probing_input = f"Device description: {scenario}\n\nClassification: {classification.model_dump_json()}"
    probing_result = await Runner.run(probing_agent, probing_input)
    probing: ProbingQuestionsOutput = probing_result.final_output
    print("PROBING QUESTIONS:")
    print(probing.model_dump_json(indent=2))
    print()

    # Feed everything to the STRIDE analyst
    stride_input = (
        f"Device description: {scenario}\n\n"
        f"Classification: {classification.model_dump_json()}\n\n"
        f"Probing questions: {probing.model_dump_json()}"
    )
    stride_result = await Runner.run(stride_agent, stride_input)
    threat_model: ThreatModel = stride_result.final_output
    print("STRIDE THREAT REGISTER:")
    print(threat_model.model_dump_json(indent=2))

How does the structured handoff work?

Each agent's final_output is now a Pydantic object, not a string. The code calls model_dump_json() to convert it to clean JSON when passing it to the next agent. This means each agent receives reliably formatted input with consistent field names.

For the display output, model_dump_json(indent=2) produces pretty-printed JSON so you can read the threat register clearly in the terminal.

  • Save threat_modeler.py.

Before you run this, think about what you expect the STRIDE analyst to find. A cloud-based cardiac arrhythmia detection app receives sensitive ECG data via a REST API, processes it with ML, and sends alerts to clinicians. What STRIDE categories do you think will be most critical?

  • Run the full three-agent pipeline:
python threat_modeler.py

You should see three sections of structured JSON output. The DEVICE CLASSIFICATION section shows the SaMD classification with its rationale and cybersecurity focus areas. The PROBING QUESTIONS section lists 8-10 targeted security questions. The STRIDE THREAT REGISTER section contains at least 6 threats covering all six STRIDE categories, each with a severity, mitigation, and FDA alignment note.

What are the six STRIDE categories?

  • Spoofing: impersonating a user, device, or system.
  • Tampering: modifying data or code without detection.
  • Repudiation: denying actions without proper logging.
  • Information Disclosure: exposing sensitive data.
  • Denial of Service: making the system unavailable.
  • Elevation of Privilege: gaining unauthorized access levels.

Seeing a validation error or timeout?

If you see a Pydantic validation error, the LLM returned data that didn't match the model's schema. Run the command again. LLM output can occasionally be inconsistent.

If the script seems to hang, it's making three sequential API calls. Each one can take 10-30 seconds depending on response length. Give it up to 2 minutes to finish.

help me debug my threat modeling pipeline output

✔️ Awesome, I've got everything!

Great. Double check you've saved both files before moving on.

ⓧ I'd like to double check the full code

Your threat_modeler.py should contain these sections in order. Scroll through and confirm each one is present:

  • Imports: asyncio, sys, Agent, Runner from agents, and DeviceClassification, ProbingQuestionsOutput, ThreatModel from models.
  • SAMPLE_SCENARIOS dictionary with three entries: samd, simd, mdds.
  • classifier_agent with output_type=DeviceClassification.
  • probing_agent with output_type=ProbingQuestionsOutput.
  • stride_agent with output_type=ThreatModel.
  • main() function that runs all three agents sequentially and prints model_dump_json(indent=2) for each result.

Your three specialist agents now produce clean, typed data at every step of the pipeline. Next up, you'll wire them together with an orchestrator agent that calls each specialist automatically, so you can run the entire threat model with a single command.

Orchestrate the Full Pipeline

You now have three specialist agents that produce clean, typed output. However, your main() function manually chains them together. It runs the classifier, formats the result into a string, feeds it to the probing agent, formats again, and feeds it to the STRIDE analyst.

The OpenAI Agents SDK has a pattern called agents-as-tools that eliminates this manual wiring. You wrap each specialist as a callable tool, then hand them to an orchestrator agent that decides what to call, what context to pass, and how to combine the results. One call to the orchestrator runs the entire pipeline.

In this step, get ready to:
  • Add a function tool and orchestrator agent that wires the three specialists together.
  • Update main() to run the full pipeline with a single orchestrator call.
  • Run the pipeline against a different device scenario to see the orchestrator in action.
Add a function tool for retrieving device scenarios

The orchestrator agent needs a way to look up sample device scenarios by name. The SDK's @function_tool decorator turns a plain Python function into a tool that any agent can call.

  • In threat_modeler.py, find the import line near the top of the file:
from agents import Agent, Runner
  • Add function_tool to the import so the line looks like this:
from agents import Agent, Runner, function_tool
  • Scroll down past the stride_agent definition. Add the get_device_scenario function tool below it by pasting this code:
@function_tool
def get_device_scenario(scenario_type: str) -> str:
    """Retrieve a sample medical device scenario. Valid types: samd, simd, mdds."""
    # Look up the scenario from the dictionary
    scenario = SAMPLE_SCENARIOS.get(scenario_type.lower())
    if scenario:
        return scenario
    return f"Unknown type: {scenario_type}. Available: {', '.join(SAMPLE_SCENARIOS.keys())}"

What does this code do?

The @function_tool decorator wraps your Python function so the SDK can expose it to any agent as a callable tool. The agent sees the function's name, type hints, and docstring, then calls it when it needs a scenario.

The function itself is simple: it looks up a key in SAMPLE_SCENARIOS and returns the matching description string. If the key is invalid, it returns a helpful error listing the available options.

  • Save threat_modeler.py and verify the function tool loads correctly by running:
python -c "from threat_modeler import get_device_scenario; print('Function tool ready')"

You should see Function tool ready printed in the terminal.

Seeing an ImportError?

  • Make sure you added function_tool to the import line. The decorator comes from the agents package.
  • Check that the @function_tool decorator sits directly above def get_device_scenario with no blank line between them.
  • Confirm your virtual environment is still activated. You should see (.venv) in your terminal prompt.

Still stuck? Help me debug my function_tool import error in threat_modeler.py

Wire the specialists into an orchestrator agent

With the function tool ready, you can now build the orchestrator. It uses agent.as_tool() to wrap each specialist agent as a callable tool that the orchestrator can invoke on its own.

Why agents-as-tools instead of manual chaining?

Manual chaining means you write Python code to run each agent, format its output, and pass it to the next agent. Every new scenario or pipeline change requires editing your orchestration code.

With agents-as-tools, the orchestrator agent decides which tools to call, in what order, and what context to pass. You describe the workflow in the orchestrator's instructions, and the SDK handles the rest. The orchestrator stays in control of the final answer while each specialist runs independently and returns its result.

The orchestrator definition is large, so you will add it in two parts. Do not save until both parts are in place.

  • In threat_modeler.py, add the first part of the orchestrator agent below the get_device_scenario function. This defines the agent's name and instructions:
orchestrator_agent = Agent(
    name="Threat Model Orchestrator",
    instructions=(
        "You coordinate end-to-end threat modeling for medical devices.\n\n"
        "When given a device description:\n"
        "1. Call classify_device with the device description\n"
        "2. Call generate_probing_questions with the device description and "
        "the classification result\n"
        "3. Call analyze_threats with the device description, classification, "
        "and probing questions\n\n"
        "After all three tools return, present the combined results as a "
        "formatted threat model report with these sections:\n\n"
        "DEVICE CLASSIFICATION: the classification, rationale, and focus areas\n"
        "PROBING QUESTIONS: numbered list of security questions generated\n"
        "STRIDE THREAT REGISTER: each threat with its category, severity, "
        "mitigation, and FDA alignment\n"
        "ASSUMPTIONS: the environmental assumptions made\n"
        "SUMMARY: overall threat landscape assessment"
    ),

What do the instructions tell the orchestrator?

The instructions lay out a three-step workflow: classify first, then generate probing questions, then run STRIDE analysis. The orchestrator follows this order because each step's output feeds into the next.

The instructions also specify the output format: five named sections that the orchestrator compiles from the three tool results. The orchestrator owns the final report.

  • Immediately below the instructions (still inside the same Agent() call), add the tools list that wires in all four tools:
    tools=[
        # Wrap each specialist agent as a callable tool
        classifier_agent.as_tool(
            tool_name="classify_device",
            tool_description=(
                "Classify a medical device description as SaMD, SiMD, or MDDS. "
                "Returns classification with rationale and cybersecurity focus areas."
            ),
        ),
        probing_agent.as_tool(
            tool_name="generate_probing_questions",
            tool_description=(
                "Generate targeted cybersecurity probing questions for a medical "
                "device based on its description and classification."
            ),
        ),
        stride_agent.as_tool(
            tool_name="analyze_threats",
            tool_description=(
                "Perform STRIDE threat analysis on a medical device given its "
                "description, classification, and security context."
            ),
        ),
        # Also include the function tool for scenario lookup
        get_device_scenario,
    ],
)

What does as_tool() do?

Each .as_tool() call converts a specialist agent into a tool the orchestrator can invoke. The tool_name is the identifier the orchestrator uses in its tool calls. The tool_description tells the orchestrator what the tool does and when to use it.

When the orchestrator calls classify_device, the SDK spins up the classifier_agent in a nested run. The specialist runs independently and returns its result to the orchestrator. The orchestrator stays in control the whole time.

  • Save threat_modeler.py and verify the orchestrator loads by running:
python -c "from threat_modeler import orchestrator_agent; print('Orchestrator ready:', orchestrator_agent.name)"

You should see Orchestrator ready: Threat Model Orchestrator printed in the terminal.

Seeing a SyntaxError?

  • Make sure the two code blocks you pasted form one continuous Agent() call. The instructions=(...) section ends with a comma, and the tools=[...] section follows directly after it.
  • Check for a missing closing parenthesis ) at the very end of the orchestrator_agent definition.
  • Verify there is no duplicate orchestrator_agent definition in the file.

Still stuck? Help me debug my orchestrator_agent SyntaxError in threat_modeler.py

Update main() to run the orchestrator end-to-end

The last piece is replacing your manual three-step main() with a single orchestrator call. The new version also accepts a command-line argument so you can choose which device scenario to analyze.

  • In threat_modeler.py, find the existing main() function. It starts with:
async def main() -> None:
    scenario = SAMPLE_SCENARIOS["samd"]
    print(f"Analyzing: SAMD")
  • Delete the entire main() function body (everything from scenario = SAMPLE_SCENARIOS["samd"] down to print(threat_model.model_dump_json(indent=2))).
  • Replace it with the following code:
async def main() -> None:
    # Default to the SaMD scenario if no argument is provided
    scenario_key = "samd"
    if len(sys.argv) > 1:
        scenario_key = sys.argv[1].lower()

    # Validate the scenario argument
    if scenario_key not in SAMPLE_SCENARIOS:
        print(f"Unknown scenario: {scenario_key}")
        print(f"Available: {', '.join(SAMPLE_SCENARIOS.keys())}")
        return

    # Look up and display the selected scenario
    scenario = SAMPLE_SCENARIOS[scenario_key]
    print(f"\nAnalyzing: {scenario_key.upper()}")
    print(f"Description: {scenario}\n")
    print("Running threat model pipeline...\n")

    # Run the full pipeline with a single orchestrator call
    result = await Runner.run(
        orchestrator_agent,
        f"Perform a complete threat model for this medical device:\n\n{scenario}",
    )

    print(result.final_output)

What changed in main()?

The old main() ran three agents manually, formatting and passing output between each one. The new version calls Runner.run() once with the orchestrator_agent. The orchestrator handles the entire classify, question, analyze workflow internally.

The sys.argv check lets you pass samd, simd, or mdds as a command-line argument. If you provide nothing, it defaults to the SaMD cardiac arrhythmia scenario.

  • Save threat_modeler.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved threat_modeler.py.

ⓧ I'd like to double check the full code

Your threat_modeler.py is too long to display in one block, so scroll through the file and confirm these sections appear in order:

  • The updated import: from agents import Agent, Runner, function_tool.
  • The SAMPLE_SCENARIOS dictionary with samd, simd, and mdds keys.
  • Three agent definitions: classifier_agent, probing_agent, stride_agent.
  • The @function_tool decorated get_device_scenario function.
  • The orchestrator_agent definition with instructions and a tools list containing four items (three .as_tool() calls plus get_device_scenario).
  • The updated main() function that reads sys.argv and calls Runner.run(orchestrator_agent, ...).
  • The if __name__ == "__main__": block at the bottom.

Your models.py is unchanged from the previous step.

Time to see the orchestrator in action. In the previous steps, you always ran the SaMD cardiac arrhythmia scenario. This time, you will point the tool at a completely different device type.

Before you run this, think about what you expect. The insulin pump scenario describes embedded firmware that controls insulin delivery. Will the orchestrator classify it the same way as the cloud-based cardiac app?

  • Run the orchestrator on the SiMD insulin pump scenario:
python threat_modeler.py simd

The orchestrator calls all three specialist agents automatically. This takes a moment as each agent runs in sequence. You should see a complete formatted threat model report with five sections: DEVICE CLASSIFICATION, PROBING QUESTIONS, STRIDE THREAT REGISTER, ASSUMPTIONS, and SUMMARY.

The device should be classified as SiMD (Software in a Medical Device), with cybersecurity focus areas specific to embedded firmware like secure boot, firmware integrity, and physical access controls.

Pipeline taking too long or timing out?

  • The orchestrator makes multiple API calls (one for itself plus one for each specialist). A full run can take 30-60 seconds. Wait for it to complete.
  • If you see an authentication error, confirm your OPENAI_API_KEY environment variable is still set in this terminal session.
  • If you see a NameError for orchestrator_agent, check that the orchestrator definition appears after all three specialist agents and the get_device_scenario function in your file.

Still stuck? Help me debug my orchestrator pipeline run

Your threat modeling tool now runs a complete pipeline with a single command. You can point it at any of the three device types and get back an FDA-aligned STRIDE threat model. Next up, a secret mission to take this even further.

Secret mission

Run Comparative Analysis Across Device Types

Your threat modeling tool currently analyzes one device at a time. What if you could run all three device types through the pipeline and get a side-by-side comparison revealing patterns a single analysis would miss? In this secret mission, you will add a comparative analysis mode that highlights how threat surfaces shift across SaMD, SiMD, and MDDS classifications.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely on your local machine, so there are no ongoing cloud infrastructure costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • Your project folder on the Desktop containing threat_modeler.py, models.py, and the .venv virtual environment.
  • OPENAI_API_KEY environment variable (if you added it to your shell configuration file).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with different device scenarios or extending the threat modeler.

  • Your project folder stays on your Desktop with all files intact.
  • Your virtual environment and installed packages remain ready to use.
  • Your OPENAI_API_KEY stays configured for future runs.

Pause - I'll come back to this later

Deactivate your virtual environment and close the project. Your files stay on disk so you can pick up where you left off.

  • Deactivate your virtual environment by running this in your terminal:
deactivate
  • Close the project in VS Code.
  • When you come back, reactivate the virtual environment with source .venv/bin/activate (macOS/Linux) or .venv\Scripts\activate (Windows).

Delete - I don't want to use this again

Remove all project files, the virtual environment, and the API key configuration.

  • Deactivate your virtual environment first (if it is still active) by running:
deactivate
  • Delete your project folder from the Desktop. On macOS or Linux, run:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your project folder name"]]
  • On Windows, run:
cd ~\Desktop
rmdir /s /q [[PROJECT_FOLDER="your project folder name"]]
  • If you added OPENAI_API_KEY to your shell configuration file (such as .zshrc or .bashrc), open the file and remove the line that exports it.
  • If you only set it for the terminal session with export, it is already gone once you close that terminal window.

Nice Work!

Nice Work!

You just built an AI-powered threat modeling tool that classifies medical devices and generates FDA-aligned STRIDE threat models. Point it at any device description and get back a structured security analysis in seconds.

What you learned:

  • Built three specialist AI agents with domain-specific instructions grounded in FDA cybersecurity guidance, and orchestrated them into a single pipeline using the OpenAI Agents SDK's agents-as-tools pattern.
  • Learned how SaMD, SiMD, and MDDS differ in classification criteria, regulatory pathways, and cybersecurity focus areas.
  • Used Pydantic structured output to enforce typed contracts between agents, replacing fragile text handoffs with clean, machine-readable data models.
  • Secret Mission: Generated a comparative threat analysis across all three medical device classifications, revealing how threat surfaces shift from cloud/API risks (SaMD) to firmware/physical risks (SiMD) to data integrity risks (MDDS).

Ready to quiz yourself?