diff --git a/plugins/pangea/pangea.test.ts b/plugins/pangea/pangea.test.ts index 3c6ab583..af3f1b27 100644 --- a/plugins/pangea/pangea.test.ts +++ b/plugins/pangea/pangea.test.ts @@ -8,7 +8,7 @@ const options = { describe('textGuardContentHandler', () => { it('should return an error if hook type is not supported', async () => { const context = { - request: { text: 'this is a test string for moderations' }, + request: { text: 'This is a message' }, }; const eventType = 'unsupported'; const parameters = {}; @@ -26,7 +26,7 @@ describe('textGuardContentHandler', () => { it('should return an error if fetch request fails', async () => { const context = { - request: { text: 'this is a test string for moderations' }, + request: { text: 'This is a message' }, }; const eventType = 'beforeRequestHook'; const parameters = { @@ -39,16 +39,19 @@ describe('textGuardContentHandler', () => { options ); expect(result.error).toBeDefined(); - expect(result.verdict).toBe(false); + expect(result.verdict).toBe(true); expect(result.data).toBeNull(); }); it('should return an error if no apiKey', async () => { const context = { - request: { text: 'this is a test string for moderations' }, + request: { text: 'This is a message' }, }; const eventType = 'beforeRequestHook'; - const parameters = { credentials: { domain: testCreds.domain } }; + const parameters = { + credentials: { domain: testCreds.domain }, + recipe: undefined, + }; const result = await textGuardContentHandler( context, parameters, @@ -62,7 +65,7 @@ describe('textGuardContentHandler', () => { it('should return an error if no domain', async () => { const context = { - request: { text: 'this is a test string for moderations' }, + request: { text: 'This is a message' }, }; const eventType = 'beforeRequestHook'; const parameters = { credentials: { apiKey: testCreds.apiKey } }; @@ -77,9 +80,9 @@ describe('textGuardContentHandler', () => { expect(result.data).toBeNull(); }); - it('should return false verdict and data if fetch request succeeds', async () => { + it('should return true verdict and data if fetch request succeeds', async () => { const context = { - request: { text: 'this is a test string for moderations' }, + request: { text: 'This is a message' }, }; const eventType = 'beforeRequestHook'; const parameters = { @@ -91,10 +94,31 @@ describe('textGuardContentHandler', () => { eventType, options ); - expect(result.error).toBeNull(); expect(result.verdict).toBeDefined(); expect(result.verdict).toBe(true); expect(result.data).toBeDefined(); + expect(result.error).toBeDefined(); + }); + + it('should return true verdict and data if fetch request succeeds. Empty `recipe`', async () => { + const context = { + request: { text: 'This is a message' }, + }; + const eventType = 'beforeRequestHook'; + const parameters = { + credentials: testCreds, + recipe: ' ', + }; + const result = await textGuardContentHandler( + context, + parameters, + eventType, + options + ); + expect(result.verdict).toBeDefined(); + expect(result.verdict).toBe(true); + expect(result.data).toBeDefined(); + expect(result.error).toBeDefined(); }); it('should return verdict as false if text is flagged', async () => { @@ -113,9 +137,9 @@ describe('textGuardContentHandler', () => { eventType, options ); - expect(result.error).toBeNull(); expect(result.verdict).toBe(false); expect(result.data).toBeDefined(); + expect(result.error).toBeNull(); }); it('should return true verdict and error if no text', async () => { @@ -137,4 +161,25 @@ describe('textGuardContentHandler', () => { expect(result.verdict).toBe(true); expect(result.data).toBeNull(); }); + + it('should return true verdict and error if invalid `recipe`', async () => { + const context = { + request: { text: 'This is a message' }, + }; + const eventType = 'beforeRequestHook'; + const parameters = { + credentials: testCreds, + recipe: 'invalid_recipe', + }; + const result = await textGuardContentHandler( + context, + parameters, + eventType, + options + ); + expect(result.error).toBeDefined(); + expect(result.verdict).toBeDefined(); + expect(result.verdict).toBe(true); + expect(result.data).toBeNull(); + }); }); diff --git a/plugins/pangea/textGuard.ts b/plugins/pangea/textGuard.ts index 589c1fd0..9c38c34f 100644 --- a/plugins/pangea/textGuard.ts +++ b/plugins/pangea/textGuard.ts @@ -4,7 +4,7 @@ import { PluginHandler, PluginParameters, } from '../types'; -import { post, getText } from '../utils'; +import { post, getText, HttpError } from '../utils'; import { VERSION } from './version'; export const handler: PluginHandler = async ( @@ -13,61 +13,82 @@ export const handler: PluginHandler = async ( eventType: HookEventType ) => { let error = null; - let verdict = false; + let verdict = true; let data = null; + if (!parameters.credentials?.domain) { + return { + error: `'parameters.credentials.domain' must be set`, + verdict: true, + data, + }; + } + + if (!parameters.credentials?.apiKey) { + return { + error: `'parameters.credentials.apiKey' must be set`, + verdict: true, + data, + }; + } + + // TODO: Update to v1 once released + const url = `https://ai-guard.${parameters.credentials.domain}/v1beta/text/guard`; + + const text = getText(context, eventType); + if (!text) { + return { + error: 'request or response text is empty', + verdict: true, + data, + }; + } + + const requestOptions = { + headers: { + 'Content-Type': 'application/json', + 'User-Agent': 'portkey-ai-plugin/' + VERSION, + Authorization: `Bearer ${parameters.credentials.apiKey}`, + }, + }; + + // Avoid sending empty strings + const recipe = parameters.recipe ? parameters.recipe : undefined; + + const request = { + text: text, + recipe: recipe, + debug: parameters.debug, + overrides: parameters.overrides, + }; + + let response; try { - if (!parameters.credentials?.domain) { - return { - error: `'parameters.credentials.domain' must be set`, - verdict: true, - data, - }; - } - - if (!parameters.credentials?.apiKey) { - return { - error: `'parameters.credentials.apiKey' must be set`, - verdict: true, - data, - }; - } - - // TODO: Update to v1 once released - const url = `https://ai-guard.${parameters.credentials.domain}/v1beta/text/guard`; - - const text = getText(context, eventType); - if (!text) { - return { - error: 'request or response text is empty', - verdict: true, - data, - }; - } - - const requestOptions = { - headers: { - 'Content-Type': 'application/json', - 'User-Agent': 'portkey-ai-plugin/' + VERSION, - Authorization: `Bearer ${parameters.credentials.apiKey}`, - }, - }; - const request = { - text: text, - recipe: parameters.recipe, - debug: parameters.debug, - overrides: parameters.overrides, - }; - - const response = await post(url, request, requestOptions); - data = response.result; - const si = response.result.findings; - if ( - !(si.prompt_injection_count || si.malicious_count || si.artifact_count) - ) { - verdict = true; - } + response = await post(url, request, requestOptions); } catch (e) { - error = e as Error; + if (e instanceof HttpError) { + error = `${e.message}. body: ${e.response.body}`; + } else { + error = e as Error; + } + } + + if (response) { + if (response.status == 'Success') { + data = response.result; + const detectors = data.detectors; + if ( + detectors?.prompt_injection?.detected || + detectors?.pii_entity?.detected || + detectors?.malicious_entity?.detected || + detectors?.custom_entity?.detected || + detectors?.secrets_detection?.detected || + detectors?.profanity_and_toxicity?.detected + ) { + verdict = false; + } + } else { + error = errorToString(response); + } } return { @@ -76,3 +97,15 @@ export const handler: PluginHandler = async ( data, // any additional data you want to return }; }; + +function errorToString(response: any): string { + let ret = `Summary: ${response.summary}\n`; + ret += `status: ${response.status}\n`; + ret += `request_id: ${response.request_id}\n`; + ret += `request_time: ${response.request_time}\n`; + ret += `response_time: ${response.response_time}\n`; + (response.result?.errors || []).forEach((ef: any) => { + ret += `\t${ef.source} ${ef.code}: ${ef.detail}\n`; + }); + return ret; +} diff --git a/plugins/utils.ts b/plugins/utils.ts index 89965d3a..0652e4fe 100644 --- a/plugins/utils.ts +++ b/plugins/utils.ts @@ -4,13 +4,13 @@ interface PostOptions extends RequestInit { headers?: Record; } -interface ErrorResponse { +export interface ErrorResponse { status: number; statusText: string; body: string; } -class HttpError extends Error { +export class HttpError extends Error { response: ErrorResponse; constructor(message: string, response: ErrorResponse) {