chore: add extra test cases

Portkey-AI · Nov 27, 2024 · 3881aaa · 3881aaa
1 parent 34fa244
commit 3881aaa
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 3 deletions.
diff --git a/plugins/mistral/index.ts b/plugins/mistral/index.ts
@@ -69,14 +69,16 @@ export const mistralGuardrailHandler: PluginHandler = async (
   const guardrailFunction = fn as GuardrailFunction;
 
   const text = getText(context, eventType);
-  const messages = context.request?.json?.messages;
+  const messages =
+    eventType === 'beforeRequestHook'
+      ? context.request?.json?.messages
+      : context.response?.json?.messages;
 
   // should contain text or should contain messages array
   if (
     (!text && !Array.isArray(messages)) ||
     (Array.isArray(messages) && messages.length === 0)
   ) {
-    console.log(!text, messages);
     return {
       error: 'Mistral: Invalid Request body',
       verdict: false,

diff --git a/plugins/mistral/mistra.test.ts b/plugins/mistral/mistra.test.ts
@@ -8,7 +8,7 @@ function getParameters() {
   };
 }
 
-describe('validateProject handler', () => {
+describe('mistral guardrail handler', () => {
   it('should fail if the apiKey is invalid', async () => {
     const eventType = 'beforeRequestHook';
     const context = {
@@ -106,4 +106,78 @@ describe('validateProject handler', () => {
     expect(result.error).toBeNull();
     expect(result.data).toBeNull();
   });
+
+  it('should give error on invalid request body', async () => {
+    const eventType = 'beforeRequestHook';
+    const context = {
+      request: {},
+    };
+    const parameters = JSON.parse(JSON.stringify(getParameters()));
+
+    const result = await mistralGuardrailHandler(
+      context as unknown as PluginContext,
+      parameters,
+      eventType,
+      { env: {} },
+      'pii'
+    );
+
+    expect(result).toBeDefined();
+    expect(result.verdict).toBe(false);
+    expect(result.error).toBe('Mistral: Invalid Request body');
+    expect(result.data).toBeNull();
+  });
+
+  it('should work for afterRequestHook', async () => {
+    const eventType = 'afterRequestHook';
+    const context = {
+      response: { text: 'this text is safe text' },
+    };
+    const parameters = JSON.parse(JSON.stringify(getParameters()));
+
+    const result = await mistralGuardrailHandler(
+      context as unknown as PluginContext,
+      parameters,
+      eventType,
+      { env: {} },
+      'pii'
+    );
+
+    expect(result).toBeDefined();
+    expect(result.verdict).toBe(false);
+    expect(result.error).toBeNull();
+    expect(result.data).toBeNull();
+  });
+
+  it('should work for afterRequestHook with chatCompletion messages', async () => {
+    const eventType = 'afterRequestHook';
+    const context = {
+      requestType: 'chatComplete',
+      response: {
+        json: {
+          messages: [
+            {
+              role: 'user',
+              content:
+                'Say Hi. My name is Jhon Doe and my email is user@example.com',
+            },
+          ],
+        },
+      },
+    };
+    const parameters = JSON.parse(JSON.stringify(getParameters()));
+
+    const result = await mistralGuardrailHandler(
+      context as unknown as PluginContext,
+      parameters,
+      eventType,
+      { env: {} },
+      'pii'
+    );
+
+    expect(result).toBeDefined();
+    expect(result.verdict).toBe(false);
+    expect(result.error).toBeNull();
+    expect(result.data).toBeNull();
+  });
 });