Skip to content

Commit ccc7f62

Browse files
VisargDVisargD
authored
Merge pull request #1198 from indranil786/main
 [New Plugin] WalledAi
2 parents 8a6086f + 1f97fd0 commit ccc7f62

File tree

5 files changed

+270
-1
lines changed

5 files changed

+270
-1
lines changed

‎conf.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
"patronus",
99
"pangea",
1010
"promptsecurity",
11-
"panw-prisma-airs"
11+
"panw-prisma-airs",
12+
"walledai"
1213
],
1314
"credentials": {
1415
"portkey": {

‎plugins/index.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ import { handler as promptSecurityProtectResponse } from './promptsecurity/prote
4949
import { handler as panwPrismaAirsintercept } from './panw-prisma-airs/intercept';
5050
import { handler as defaultjwt } from './default/jwt';
5151
import { handler as defaultrequiredMetadataKeys } from './default/requiredMetadataKeys';
52+
import { handler as walledaiguardrails } from './walledai/guardrails';
5253

5354
export const plugins = {
5455
default: {
@@ -134,4 +135,7 @@ export const plugins = {
134135
'panw-prisma-airs': {
135136
intercept: panwPrismaAirsintercept,
136137
},
138+
walledai: {
139+
guardrails: walledaiguardrails,
140+
},
137141
};

‎plugins/walledai/guardrails.ts‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
import {
2+
HookEventType,
3+
PluginContext,
4+
PluginHandler,
5+
PluginParameters,
6+
} from '../types';
7+
import { post, getText, getCurrentContentPart } from '../utils';
8+
9+
const API_URL = 'https://services.walled.ai/v1/guardrail/moderate';
10+
11+
export const handler: PluginHandler = async (
12+
context: PluginContext,
13+
parameters: PluginParameters,
14+
eventType: HookEventType
15+
) => {
16+
let error = null;
17+
let verdict = true;
18+
let data = null;
19+
20+
if (!parameters.credentials?.apiKey) {
21+
return {
22+
error: `'parameters.credentials.apiKey' must be set`,
23+
verdict: true,
24+
data,
25+
};
26+
}
27+
28+
const { content, textArray } = getCurrentContentPart(context, eventType);
29+
if (!content) {
30+
return {
31+
error: { message: 'request or response json is empty' },
32+
verdict: true,
33+
data: null,
34+
};
35+
}
36+
let text = textArray.filter((text) => text).join('\n');
37+
38+
// Prepare request body
39+
const requestBody = {
40+
text: text,
41+
text_type: parameters.text_type || 'prompt',
42+
generic_safety_check: parameters.generic_safety_check ?? true,
43+
greetings_list: parameters.greetings_list || ['generalgreetings'],
44+
};
45+
// Prepare headers
46+
const requestOptions = {
47+
headers: {
48+
'Content-Type': 'application/json',
49+
Authorization: `Bearer ${parameters.credentials.apiKey}`, // Uncomment if API key is required
50+
},
51+
};
52+
53+
try {
54+
const response = await post(
55+
API_URL,
56+
requestBody,
57+
requestOptions,
58+
parameters.timeout
59+
);
60+
data = response.data;
61+
if (data.safety[0]?.isSafe == false) {
62+
verdict = false;
63+
}
64+
} catch (e) {
65+
console.log(e);
66+
error = e instanceof Error ? e.message : String(e);
67+
verdict = true;
68+
data = null;
69+
}
70+
return {
71+
error,
72+
verdict,
73+
data,
74+
};
75+
};

‎plugins/walledai/manifest.json‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
{
2+
"id": "walledai",
3+
"description": "Walled AI",
4+
"credentials": {
5+
"type": "object",
6+
"properties": {
7+
"apiKey": {
8+
"type": "string",
9+
"label": "API Key",
10+
"description": "Find your API key in the Walled AI dashboard (https://dev.walled.ai/)",
11+
"encrypted": true
12+
}
13+
},
14+
"required": ["apiKey"]
15+
},
16+
"functions": [
17+
{
18+
"name": "Walled AI Guardrail for checking safety of LLM inputs",
19+
"id": "guardrails",
20+
"supportedHooks": ["beforeRequestHook", "afterRequestHook"],
21+
"type": "guardrail",
22+
"description": [
23+
{
24+
"type": "subHeading",
25+
"text": "Ensure the safety and compliance of your LLM inputs with Walled AI's advanced guardrail system."
26+
}
27+
],
28+
"parameters": {
29+
"type": "object",
30+
"properties": {
31+
"text_type": {
32+
"type": "string",
33+
"label": "Text Type",
34+
"description": [
35+
{
36+
"type": "subHeading",
37+
"text": "Type of Text , defaults to 'prompt'"
38+
}
39+
],
40+
"default": "prompt"
41+
},
42+
"generic_safety_check": {
43+
"type": "string",
44+
"label": "Generic Safety Check",
45+
"description": [
46+
{
47+
"type": "subHeading",
48+
"text": "Boolean value to enable generic safety checks on the text input. Defaults to 'true'."
49+
}
50+
],
51+
"default": true
52+
},
53+
"greetings_list": {
54+
"type": "array",
55+
"label": "Greetings List",
56+
"description": [
57+
{
58+
"type": "subHeading",
59+
"text": "List of greetings to be used in the guardrail check. This can help in identifying and handling greetings appropriately."
60+
}
61+
],
62+
"items": {
63+
"type": "string",
64+
"default": ["generalgreetings"]
65+
}
66+
},
67+
"pii_list": {
68+
"type": "array",
69+
"label": "PII LIST",
70+
"description": [
71+
{
72+
"type": "subHeading",
73+
"text": "Identify all the PII for only the following types of PII will be checked in the text input. Defaults to empty list"
74+
}
75+
],
76+
"items": {
77+
"type": "string",
78+
"enum": [
79+
"Person's Name",
80+
"Address",
81+
"Email Id",
82+
"Contact No",
83+
"Date Of Birth",
84+
"Unique Id",
85+
"Financial Data"
86+
]
87+
}
88+
},
89+
"compliance_list": {
90+
"type": "array",
91+
"label": "List of Compliance Checks",
92+
"description": [
93+
{
94+
"type": "subHeading",
95+
"text": "List of compliance checks to be performed on the text input. This can help in ensuring that the text adheres to specific compliance standards. Defaults to empty"
96+
}
97+
],
98+
"items": {
99+
"type": "string"
100+
}
101+
}
102+
}
103+
}
104+
}
105+
]
106+
}

‎plugins/walledai/walledai.test.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import { handler } from './guardrails';
2+
import testCredsFile from './creds.json';
3+
import { HookEventType, PluginContext, PluginParameters } from '../types';
4+
5+
const testCreds = {
6+
apiKey: testCredsFile.apiKey,
7+
};
8+
9+
describe('WalledAI Guardrail Plugin Handler (integration)', () => {
10+
const baseParams: PluginParameters = {
11+
credentials: testCreds,
12+
text_type: 'prompt',
13+
generic_safety_check: true,
14+
greetings_list: ['generalgreetings'],
15+
pii_list: ["Person's Name", 'Address'],
16+
compliance_list: [],
17+
};
18+
19+
const makeContext = (text: string): PluginContext => ({
20+
requestType: 'chatComplete',
21+
request: {
22+
json: {
23+
messages: [{ role: 'user', content: text }],
24+
},
25+
},
26+
response: {},
27+
});
28+
29+
it('returns verdict=true for safe text', async () => {
30+
const context = makeContext('Hello, how are you');
31+
32+
const result = await handler(context, baseParams, 'beforeRequestHook');
33+
34+
expect(result.verdict).toBe(true);
35+
expect(result.error).toBeNull();
36+
expect(result.data).toBeDefined();
37+
});
38+
39+
it('returns verdict=false for unsafe text', async () => {
40+
const context = makeContext('I want to harm someone.');
41+
42+
const result = await handler(context, baseParams, 'beforeRequestHook');
43+
44+
expect(result.verdict).toBe(false);
45+
expect(result.error).toBeNull();
46+
});
47+
48+
it('returns error if apiKey is missing', async () => {
49+
const context = makeContext('Hello world');
50+
51+
const result = await handler(
52+
context,
53+
{ ...baseParams, credentials: {} },
54+
'beforeRequestHook'
55+
);
56+
57+
expect(result.error).toMatch(/apiKey/i);
58+
expect(result.verdict).toBe(true);
59+
});
60+
61+
it('returns error if text is empty', async () => {
62+
const context = makeContext('');
63+
64+
const result = await handler(context, baseParams, 'beforeRequestHook');
65+
66+
expect(result.error).toBeDefined();
67+
expect(result.verdict).toBe(true);
68+
expect(result.data).toBeNull();
69+
});
70+
71+
it('uses default values for missing optional parameters', async () => {
72+
const context = makeContext('Hello world');
73+
74+
const minimalParams: PluginParameters = {
75+
credentials: testCreds,
76+
};
77+
78+
const result = await handler(context, minimalParams, 'beforeRequestHook');
79+
80+
expect(result.verdict).toBe(true);
81+
expect(result.error).toBeNull();
82+
});
83+
});

0 commit comments

Comments
 (0)