How the Rule-First, LLM-Augmented Approach Works in Tencent's AI-Infra-Guard
TLDR: AI-Infra-Guard runs deterministic YAML-based security rules through a Go scanner engine first, then optionally augments each matched finding with a large-language-model (LLM) plugin that assesses nuances, adds severity ratings, and generates remediation insights — finally merging both outputs into a single advisory for a hybrid detection pipeline.
Tencent's AI-Infra-Guard repository combines two complementary detection philosophies: a classically engineered rule-first scanner and a flexible LLM-augmented analysis layer. The goal isn't to replace static detection with AI noise, but to let each layer correct the other's weaknesses. As implemented across pkg/vulstruct/scanner.go and the Python-based deepteam plugin system, this approach guarantees known vulnerabilities are caught with zero false negatives, while LLMs contribute context-aware judgment that pure signatures cannot express.
The Rule-First Core: Deterministic Scanning Engine
The foundation of AI-Infra-Guard is an exhaustive set of versioned, static rules defined in YAML files. These rules encode all of the following:
- Vulnerability fingerprints for known attack patterns
- Exact-match signatures for prompt injection, toxicity, and data leakage
- MCP (Model Context Protocol) server plugin definitions
- Regex-free structural patterns that describe HTTP responses, code snippets, and container images
The rule engine in pkg/vulstruct/scanner.go is the execution backbone. It loads these YAML rules, walks through the target assets (code files, HTTP responses, container images), and checks each byte stream against every pattern in the rule set. When a rule matches, the engine emits a preliminary advisory object — defined in pkg/vulstruct/advisory.go — that captures the rule ID, severity, evidence snippet, and context.
Because the rules are static and immutable, the scanner produces reproducible results every time. That guarantee matters for CI/CD pipelines, compliance audits, and regression testing.
The LLM-Augmented Layer: Adding Context-Aware Intelligence
After the rule engine finishes its deterministic pass, AI‑Ingested Governance optionally invokes an LLM to enrich the findings. This augmentation isn't bolted on ad hoc — it's a structured plugin architecture that maps each vulnerability type to a dedicated LLM wrapper.
For example, in the Prompt Security sub-module, each vulnerability category (toxicity, bias, prompt leakage, etc.) is implemented as a separate plugin. These plugins rely on the generic infrastructure in deepteam/plugin_system/plugin_manager.py and are registered via deepteam/plugin_system/plugin_registry.py. The plugin registry is the key link: it maintains a mapping from a rule identifier to its corresponding LLM plugin.
Inside each plugin, the flow is:
- The plugin receives the raw input — a prompt text, code snippet, or full HTTP payload.
- It attaches the rule engine's context (which specific rule matched, what severity was assigned).
- It constructs a tailored LLM prompt using a per-vulnerability template found under
deepteam/vulnerabilities/*/template.py. - It sends that prompt to the configured LLM provider (OpenAI, Anthropic, or a local model).
- The LLM returns a structured response that the plugin parses and normalizes — typically a severity score, risk category, or a human‑readable justification.
- The
deepteam/metrics/*/*.pyfiles implement the exact API calls and response‑parsing logic for each metric type.
The final LLM output is then merged back into the original advisory object, which already contains the deterministic rule evidence.
How the Rule-First + LLM Workflow Executes Step by Step
Here's the exact execution order in code:
Step 1 — Rule Matching (Go)
The scanner iterates over all static rules loaded from the YAML directory. When a rule matches the target, it creates a preliminary finding (an Advisory instance) via advisory.go. No LLM is called yet.
Step 2 — LLM Invocation (Python, conditional)
For each advisory, the scanner checks if the matched rule is flagged as requiring LLM augmentation (that flag lives in the rule definition itself). If so, the plugin registry (plugin_registry.py) is consulted to resolve the correct plugin for the vulnerability type. The plugin then constructs a prompt from its template and sends the request to the configured LLM provider.
Step 3 — Result Fusion (Python + Go)
The LLM's response — a severity rating, generated remediation text, non‑technical explanation — is attached directly to the advisory object. If the LLM assigns a more granular severity (e.g., "critical" vs the rule's binary "high"), that richer value overwrites the coarse rule severity. The final report contains all evidence.
This three‑step pipeline yields a hybrid advisory that has both:
- Deterministic evidence — the exact rule text, matched snippet, line numbers
- Generative insight — an LLM-authored explanation, a root‑cause analysis, or even a suggested patch
Benefits of the Rule-First + LLM Combination
The design solves several known trade‑offs between pure static analysis and pure AI‑driven detection:
- Determinism against drift — Rules are versioned and reproducible, so any deployment or scan runs yields identical findings — families.
- Coverage that scales with nuance — LLMs understand context. They can flag subtle prompt leakage or ambiguous phrasing in a prompt that a regex-based rule would miss entirely.
- Explainability — The rule engine retains the exact match artifacts, which serve as the human‑verifiable evidence chain. The LLM's output adds a readable justification that auditors and developers can trust.
When the LLM layer is disabled (e.g., in an offline environment), the scanner still functions as a full rule‑first engine — the augmentation is an optional, composable step that you can enable per rule or globally.
Code Examples: The Pipeline in Action
Example 1: Python — Rule‑First Scan with Optional LLM Augmentation (Prompt Security)
# Example: Running a rule‑first scan with optional LLM augmentation
from pkg.vulstruct import scanner, advisory
# 1️⃣ Load static rules
scanner_obj = scanner.Scanner()
scanner_obj.load_rules(dir_path="data/vuln") # ← rule‑first engine
# 2️⃣ Scan a target (e.g., a prompt file)
findings = scanner_obj.scan_file("samples/prompt.txt")
# 3️⃣ For each finding, optionally invoke the LLM plugin
for adv in findings:
if adv.requires_llm:
# Load the LLM plugin for the specific vulnerability type
from deepteam.plugin_system import plugin_manager
llm_plugin = plugin_manager.get_plugin(adv.vuln_type)
# Run the LLM and attach its result
llm_result = llm_plugin.run(adv.input_content)
adv.llm_result = llm_result
adv.severity = llm_result["severity"]
print(adv.to_dict())
Example 2: Go — Rule Engine Usage in the Core Scanner
// Example: Using the Go rule engine to detect a known vulnerability
package main
import (
"fmt"
"github.com/Tencent/AI-Infra-Guard/pkg/vulstruct"
)
func main() {
// Initialize scanner and load compressed YAML rules
s := vulstruct.NewScanner()
if err := s.LoadRules("data/vuln"); err != nil {
panic(err)
}
// Scan a URL response
advs, err := s.ScanURL("http://example.com")
if err != nil {
panic(err)
}
// Print matched advisories (rule‑first results)
for _, a := range advs {
fmt.Printf("Rule: %s, Severity: %s\n", a.ID, a.Severity)
}
}
Key Files Implementing the Approach
Summary
- Rule-first guarantees reproducibility and zero false negatives: every scan runs identically, and known patterns are always caught.
- LLM augmentation adds AI‑powered nuance: it assesses context, explains findings in plain language, and lets you set severity thresholds informed by LLM judgment.
- The pipeline is strictly sequential — rule matching (
scanner.go) runs first, then optional plugin invocation (Pythonplugin_*files) runs second, and finally result fusion (advisory.go) merges both. - You control the LLM: disable it entirely, enable it only for specific rule categories, or use it to override rule-severity with logical LLM severity scores.
Frequently Asked Questions
When should I enable the LLM-augmented layer?
Enable it when you need context-aware justification for the rule findings — such as for a security analyst workflow that requires human-readable remediation suggestions for each alert. Disable it when you need fastest scans and fully deterministic CI/CD gates with no external API calls.
Does the LLM replace the rule engine?
No. The rule engine remains the source of evidence and always runs first. The LLM only adds enrichment to the findings and can optionally adjust the severity. According to the architecture in scanner.go and plugin_registry.py, a deterministic rule timestamp exists for every validated finding before LLM logic ever runs.
How does the registry know which LLM plugin to use?
The plugin_registry.py file maintains a rule_id-to-plugin mapping. When adv.requires_llm is True, that resolver is called with the vulnerability type (adv.vuln_type) and returns the correct metaclass. Each plugin then wraps a specific template.py that generates the risk prompt.
Can I use a local LLM with the augmentation layer?
Yes. The plugin architecture depends on the configured provider fixed in the Python metric implementation (deepteam/metrics/). Each metric file abstracts the vendor API — so you can point to OpenAI, Anthropic, or a self-hosted model without changing the rule‑first pipeline or the Go scanner.
Have a question about this repo?
These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →