Skip to content

Commit 16861d0

Browse files
authored
Merge pull request #3598 from akto-api-security/agentic_base_prompt_intent_guardrail
Add Base Prompt Rule functionality to Guardrail Policies
2 parents 51a05a9 + 6a4d873 commit 16861d0

File tree

11 files changed

+270
-21
lines changed

11 files changed

+270
-21
lines changed

apps/dashboard/src/main/java/com/akto/action/GuardrailPoliciesAction.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.List;
2121
import java.util.Map;
2222

23+
2324
public class GuardrailPoliciesAction extends UserAction {
2425
private static final LoggerMaker loggerMaker = new LoggerMaker(GuardrailPoliciesAction.class, LogDb.DASHBOARD);
2526

@@ -55,6 +56,7 @@ public String fetchGuardrailPolicies() {
5556
}
5657
}
5758

59+
5860
public String createGuardrailPolicy() {
5961
try {
6062
User user = getSUser();
@@ -99,6 +101,7 @@ public String createGuardrailPolicy() {
99101
updates.add(Updates.set("regexPatternsV2", policy.getRegexPatternsV2()));
100102
updates.add(Updates.set("contentFiltering", policy.getContentFiltering()));
101103
updates.add(Updates.set("llmRule", policy.getLlmRule()));
104+
updates.add(Updates.set("basePromptRule", policy.getBasePromptRule()));
102105
updates.add(Updates.set("selectedMcpServers", policy.getSelectedMcpServers()));
103106
updates.add(Updates.set("selectedAgentServers", policy.getSelectedAgentServers()));
104107
updates.add(Updates.set("selectedMcpServersV2", policy.getSelectedMcpServersV2()));

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/guardrails/GuardrailPolicies.jsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,8 @@ function GuardrailPolicies() {
446446
contentFiltering: guardrailData.contentFilters || {},
447447
// Add LLM policy if present
448448
...(guardrailData.llmRule ? { llmRule: guardrailData.llmRule } : {}),
449+
// Add Base Prompt Rule if present
450+
...(guardrailData.basePromptRule ? { basePromptRule: guardrailData.basePromptRule } : {}),
449451
applyOnResponse: guardrailData.applyOnResponse || false,
450452
applyOnRequest: guardrailData.applyOnRequest || false,
451453
url: guardrailData.url || '',

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/guardrails/components/CreateGuardrailModal.jsx

Lines changed: 59 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import {
2727
SensitiveInfoConfig,
2828
LlmPromptStep,
2929
LlmPromptConfig,
30+
BasePromptStep,
31+
BasePromptConfig,
3032
ExternalModelStep,
3133
ExternalModelConfig,
3234
ServerSettingsStep,
@@ -76,24 +78,28 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
7678
const [llmPrompt, setLlmPrompt] = useState("");
7779
const [llmConfidenceScore, setLlmConfidenceScore] = useState(0.5);
7880

79-
// Step 7: External model based evaluation
81+
// Step 7: Base Prompt Based Validation (AI Agents)
82+
const [enableBasePromptRule, setEnableBasePromptRule] = useState(false);
83+
const [basePromptConfidenceScore, setBasePromptConfidenceScore] = useState(0.5);
84+
85+
// Step 8: External model based evaluation
8086
const [url, setUrl] = useState("");
8187
const [confidenceScore, setConfidenceScore] = useState(25); // Start with 25 (first checkpoint)
8288

83-
// Step 8: Server settings
89+
// Step 9: Server settings
8490
const [selectedMcpServers, setSelectedMcpServers] = useState([]);
8591
const [selectedAgentServers, setSelectedAgentServers] = useState([]);
8692
const [applyOnResponse, setApplyOnResponse] = useState(false);
8793
const [applyOnRequest, setApplyOnRequest] = useState(false);
88-
94+
8995
// Collections data
9096
const [mcpServers, setMcpServers] = useState([]);
9197
const [agentServers, setAgentServers] = useState([]);
9298
const [collectionsLoading, setCollectionsLoading] = useState(false);
93-
99+
94100
// Get collections from PersistStore
95101
const allCollections = PersistStore(state => state.allCollections);
96-
102+
97103
// Create validation state object
98104
const getStoredStateData = () => ({
99105
// Step 1
@@ -113,9 +119,12 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
113119
llmPrompt,
114120
llmConfidenceScore,
115121
// Step 7
122+
enableBasePromptRule,
123+
basePromptConfidenceScore,
124+
// Step 8
116125
url,
117126
confidenceScore,
118-
// Step 8
127+
// Step 9
119128
selectedMcpServers,
120129
selectedAgentServers,
121130
mcpServers,
@@ -164,6 +173,12 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
164173
summary: LlmPromptConfig.getSummary(storedStateData),
165174
...LlmPromptConfig.validate(storedStateData)
166175
},
176+
{
177+
number: BasePromptConfig.number,
178+
title: BasePromptConfig.title,
179+
summary: BasePromptConfig.getSummary(storedStateData),
180+
...BasePromptConfig.validate(storedStateData)
181+
},
167182
{
168183
number: ExternalModelConfig.number,
169184
title: ExternalModelConfig.title,
@@ -269,6 +284,8 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
269284
setNewRegexPattern("");
270285
setLlmPrompt("");
271286
setLlmConfidenceScore(0.5);
287+
setEnableBasePromptRule(false);
288+
setBasePromptConfidenceScore(0.5);
272289
setUrl("");
273290
setConfidenceScore(25);
274291
setSelectedMcpServers([]);
@@ -282,7 +299,7 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
282299
setDescription(policy.description || "");
283300
setBlockedMessage(policy.blockedMessage || "");
284301
setApplyToResponses(policy.applyToResponses || false);
285-
302+
286303
// Content filters
287304
if (policy.contentFiltering) {
288305
if (policy.contentFiltering.harmfulCategories) {
@@ -301,19 +318,19 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
301318
setPromptAttackLevel(policy.contentFiltering.promptAttacks.level || "HIGH");
302319
}
303320
}
304-
321+
305322
// Denied topics
306323
setDeniedTopics(policy.deniedTopics || []);
307-
324+
308325
// Word filters
309326
setWordFilters({
310327
profanity: policy.wordFilters?.profanity || false,
311328
custom: policy.wordFilters?.custom || []
312329
});
313-
330+
314331
// PII filters
315332
setPiiTypes(policy.piiTypes || []);
316-
333+
317334
// Regex patterns - prefer V2 format with behavior, fallback to old format
318335
if (policy.regexPatternsV2 && policy.regexPatternsV2.length > 0) {
319336
setRegexPatterns(policy.regexPatternsV2);
@@ -337,6 +354,15 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
337354
setLlmConfidenceScore(0.5);
338355
}
339356

357+
// Base Prompt Based Validation (AI Agents)
358+
if (policy.basePromptRule) {
359+
setEnableBasePromptRule(policy.basePromptRule.enabled || false);
360+
setBasePromptConfidenceScore(policy.basePromptRule.confidenceScore !== undefined ? policy.basePromptRule.confidenceScore : 0.5);
361+
} else {
362+
setEnableBasePromptRule(false);
363+
setBasePromptConfidenceScore(0.5);
364+
}
365+
340366
// External model based evaluation
341367
setUrl(policy.url || "");
342368
// Map existing confidence score to nearest checkpoint
@@ -435,6 +461,12 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
435461
confidenceScore: llmConfidenceScore
436462
}
437463
} : {}),
464+
...(enableBasePromptRule ? {
465+
basePromptRule: {
466+
enabled: true,
467+
confidenceScore: basePromptConfidenceScore
468+
}
469+
} : {}),
438470
url: url || null,
439471
confidenceScore: confidenceScore,
440472
selectedMcpServers: selectedMcpServers, // Old format (just IDs)
@@ -629,6 +661,15 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
629661
/>
630662
);
631663
case 7:
664+
return (
665+
<BasePromptStep
666+
enableBasePromptRule={enableBasePromptRule}
667+
setEnableBasePromptRule={setEnableBasePromptRule}
668+
basePromptConfidenceScore={basePromptConfidenceScore}
669+
setBasePromptConfidenceScore={setBasePromptConfidenceScore}
670+
/>
671+
);
672+
case 8:
632673
return (
633674
<ExternalModelStep
634675
url={url}
@@ -637,7 +678,7 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
637678
setConfidenceScore={setConfidenceScore}
638679
/>
639680
);
640-
case 8:
681+
case 9:
641682
return (
642683
<ServerSettingsStep
643684
selectedMcpServers={selectedMcpServers}
@@ -669,7 +710,7 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
669710
});
670711
}
671712

672-
if (currentStep < 8) {
713+
if (currentStep < steps.length) {
673714
actions.push({
674715
content: "Next",
675716
onAction: handleNext
@@ -683,10 +724,10 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
683724
// Check if all steps are valid
684725
const allStepsValid = steps.every(step => step.isValid);
685726

686-
return {
687-
content: isEditMode ? "Update Guardrail" : "Create Guardrail",
688-
onAction: handleSave,
689-
loading: loading,
727+
return {
728+
content: isEditMode ? "Update Guardrail" : "Create Guardrail",
729+
onAction: handleSave,
730+
loading: loading,
690731
disabled: !allStepsValid
691732
};
692733
};
@@ -710,7 +751,7 @@ const CreateGuardrailModal = ({ isOpen, onClose, onSave, editingPolicy = null, i
710751
<Modal.Section>
711752
<Scrollable style={{ height: "600px" }}>
712753
{renderAllSteps()}
713-
</Scrollable>
754+
</Scrollable>
714755
</Modal.Section>
715756
</Modal>
716757
</>
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import { VerticalStack, Text, RangeSlider, FormLayout, Checkbox, Box } from "@shopify/polaris";
2+
3+
export const BasePromptConfig = {
4+
number: 7,
5+
title: "Intent verification using base prompt (AI Agents)",
6+
7+
validate: () => {
8+
return { isValid: true, errorMessage: null };
9+
},
10+
11+
getSummary: ({ enableBasePromptRule, basePromptConfidenceScore }) => {
12+
if (!enableBasePromptRule) return null;
13+
return `Auto-detect from traffic, Confidence: ${basePromptConfidenceScore.toFixed(2)}`;
14+
}
15+
};
16+
17+
const BasePromptStep = ({
18+
enableBasePromptRule,
19+
setEnableBasePromptRule,
20+
basePromptConfidenceScore,
21+
setBasePromptConfidenceScore
22+
}) => {
23+
return (
24+
<VerticalStack gap="4">
25+
<Text variant="bodyMd" tone="subdued">
26+
Verify if agent requests match the intent of the base prompt. The base prompt is automatically detected from traffic, and user inputs filling placeholders like {`{var}`} or {`{}`} are checked against this intent.
27+
</Text>
28+
29+
<FormLayout>
30+
<Checkbox
31+
label="Enable agent intent verification"
32+
checked={enableBasePromptRule}
33+
onChange={setEnableBasePromptRule}
34+
/>
35+
36+
{enableBasePromptRule && (
37+
<Box>
38+
<Box paddingBlockStart="2">
39+
<RangeSlider
40+
label="Confidence Threshold"
41+
value={basePromptConfidenceScore}
42+
min={0}
43+
max={1}
44+
step={0.1}
45+
output
46+
onChange={setBasePromptConfidenceScore}
47+
helpText="Set the confidence threshold (0-1). Higher values require more confidence to block content."
48+
/>
49+
</Box>
50+
</Box>
51+
)}
52+
</FormLayout>
53+
</VerticalStack>
54+
);
55+
};
56+
57+
export default BasePromptStep;
58+

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/guardrails/components/steps/ExternalModelStep.jsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ const validateUrl = (url) => {
88
};
99

1010
export const ExternalModelConfig = {
11-
number: 7,
11+
number: 8,
1212
title: "External model based evaluation",
1313

1414
validate: ({ url }) => {

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/guardrails/components/steps/ServerSettingsStep.jsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import { VerticalStack, Text, FormLayout, Box, Checkbox } from "@shopify/polaris
22
import DropdownSearch from "../../../../components/shared/DropdownSearch";
33

44
export const ServerSettingsConfig = {
5-
number: 8,
5+
number: 9,
66
title: "Server and application settings",
77

88
validate: () => {

apps/dashboard/web/polaris_web/web/src/apps/dashboard/pages/guardrails/components/steps/index.js

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,6 @@ export { default as DeniedTopicsStep, DeniedTopicsConfig } from './DeniedTopicsS
55
export { default as WordFiltersStep, WordFiltersConfig } from './WordFiltersStep';
66
export { default as SensitiveInfoStep, SensitiveInfoConfig } from './SensitiveInfoStep';
77
export { default as LlmPromptStep, LlmPromptConfig } from './LlmPromptStep';
8+
export { default as BasePromptStep, BasePromptConfig } from './BasePromptStep';
89
export { default as ExternalModelStep, ExternalModelConfig } from './ExternalModelStep';
910
export { default as ServerSettingsStep, ServerSettingsConfig } from './ServerSettingsStep';

0 commit comments

Comments
 (0)