Merge branch 'main' of github.com:Portkey-AI/gateway into fix/webhook-plugin-errors

Author: b4s36t4

cookbook/integrations/Sutra_with_Portkey.ipynb

@@ -706,6 +706,49 @@
         },
         "required": ["jwksUri", "headerKey"]
       }
+    },
+    {
+      "name": "Required Metadata Keys",
+      "id": "requiredMetadataKeys",
+      "type": "guardrail",
+      "supportedHooks": ["beforeRequestHook"],
+      "description": [
+        {
+          "type": "subHeading",
+          "text": "Checks if the metadata contains all the required keys"
+        }
+      ],
+      "parameters": {
+        "type": "object",
+        "properties": {
+          "metadataKeys": {
+            "type": "array",
+            "label": "Metadata Keys",
+            "description": [
+              {
+                "type": "subHeading",
+                "text": "Enter the metadata keys to check for."
+              }
+            ],
+            "items": {
+              "type": "string"
+            }
+          },
+          "operator": {
+            "type": "string",
+            "label": "Operator",
+            "description": [
+              {
+                "type": "subHeading",
+                "text": "Select the operator to use."
+              }
+            ],
+            "enum": ["all", "any", "none"],
+            "default": "all"
+          }
+        },
+        "required": ["metadataKeys", "operator"]
+      }
     }
   ]
 }

plugins/default/manifest.json

@@ -0,0 +1,120 @@
+import {
+  HookEventType,
+  PluginContext,
+  PluginHandler,
+  PluginHandlerResponse,
+  PluginParameters,
+} from '../types';
+
+export const handler: PluginHandler = async (
+  context: PluginContext,
+  parameters: PluginParameters,
+  eventType: HookEventType
+) => {
+  const pluginResponse: PluginHandlerResponse = {
+    error: null,
+    verdict: true,
+    data: null,
+    transformedData: {
+      request: {
+        json: null,
+      },
+      response: {
+        json: null,
+      },
+    },
+    transformed: false,
+  };
+
+  try {
+    // Only process before request hooks
+    if (eventType !== 'beforeRequestHook') {
+      pluginResponse.error = {
+        message: 'This plugin only works for before_request_hooks',
+      };
+      return pluginResponse;
+    }
+    const metadataKeys = parameters.metadataKeys;
+    if (!(Array.isArray(metadataKeys) && metadataKeys.length > 0)) {
+      pluginResponse.error = {
+        message: 'metadataKeys must be an array and not empty',
+      };
+      return pluginResponse;
+    }
+    if (metadataKeys.some((key: string) => typeof key !== 'string')) {
+      pluginResponse.error = {
+        message: 'metadataKeys must be an array of strings',
+      };
+      return pluginResponse;
+    }
+    if (!['all', 'any', 'none'].includes(parameters.operator)) {
+      pluginResponse.error = {
+        message: 'operator must be one of: all, any, none',
+      };
+      return pluginResponse;
+    }
+
+    const foundMetadataKeys: string[] = [];
+    const missingMetadataKeys: string[] = [];
+
+    const metadata = context.metadata || {};
+
+    metadataKeys.forEach((key: string) => {
+      if (key in metadata) {
+        foundMetadataKeys.push(key);
+      } else {
+        missingMetadataKeys.push(key);
+      }
+    });
+
+    switch (parameters.operator) {
+      case 'any':
+        pluginResponse.verdict = foundMetadataKeys.length > 0;
+        break;
+      case 'all':
+        pluginResponse.verdict = missingMetadataKeys.length === 0;
+        break;
+      case 'none':
+        pluginResponse.verdict = foundMetadataKeys.length === 0;
+        break;
+    }
+
+    pluginResponse.data = {
+      explanation: `Check ${pluginResponse.verdict ? 'passed' : 'failed'} for '${parameters.operator}' metadata keys.`,
+      missing_metadata_keys: missingMetadataKeys,
+      found_metadata_keys: foundMetadataKeys,
+      operator: parameters.operator,
+    };
+
+    if (pluginResponse.verdict) {
+      switch (parameters.operator) {
+        case 'any':
+          pluginResponse.data.explanation += ` At least one metadata key was found.`;
+          break;
+        case 'all':
+          pluginResponse.data.explanation += ` All metadata keys were found.`;
+          break;
+        case 'none':
+          pluginResponse.data.explanation += ` No metadata keys were found.`;
+          break;
+      }
+    } else {
+      switch (parameters.operator) {
+        case 'any':
+          pluginResponse.data.explanation += ` No metadata keys were found.`;
+          break;
+        case 'all':
+          pluginResponse.data.explanation += ` Some metadata keys were missing.`;
+          break;
+        case 'none':
+          pluginResponse.data.explanation += ` Some metadata keys were found.`;
+          break;
+      }
+    }
+  } catch (e: any) {
+    pluginResponse.error = e;
+    pluginResponse.data = null;
+  }
+
+  return pluginResponse;
+};

plugins/default/requiredMetadataKeys.ts

@@ -37,8 +37,8 @@ const performExaSearch = async (
     query,
     numResults,
     useAutoprompt: true,
-    contents: {
-      text: true,
+    text: {
+      includeHtmlTags: false,
     },
   };
 

plugins/exa/online.ts

@@ -48,6 +48,7 @@ import { handler as promptSecurityProtectPrompt } from './promptsecurity/protect
 import { handler as promptSecurityProtectResponse } from './promptsecurity/protectResponse';
 import { handler as panwPrismaAirsintercept } from './panw-prisma-airs/intercept';
 import { handler as defaultjwt } from './default/jwt';
+import { handler as defaultrequiredMetadataKeys } from './default/requiredMetadataKeys';
 
 export const plugins = {
   default: {
@@ -67,6 +68,7 @@ export const plugins = {
     endsWith: defaultendsWith,
     modelWhitelist: defaultmodelWhitelist,
     jwt: defaultjwt,
+    requiredMetadataKeys: defaultrequiredMetadataKeys,
   },
   portkey: {
     moderateContent: portkeymoderateContent,

plugins/index.ts

@@ -27,6 +27,36 @@ export interface PIIResult {
   redactedText: string;
 }
 
+export interface LogObjectRequest {
+  url: string;
+  method: string;
+  headers: Record<string, any>;
+  body: any;
+  provider?: string;
+}
+
+export interface LogObjectResponse {
+  status: number;
+  headers: Record<string, any>;
+  body: any;
+  response_time: number;
+  streamingMode: boolean;
+}
+
+export interface LogObjectMetadata extends Record<string, any> {
+  traceId?: string;
+  spanId?: string;
+  parentSpanId?: string;
+  spanName?: string;
+}
+
+export interface LogObject {
+  request: LogObjectRequest;
+  response: LogObjectResponse;
+  metadata: LogObjectMetadata;
+  createdAt: string;
+}
+
 interface PIIParameters extends PluginParameters {
   categories: string[];
   credentials: Record<string, any>;
@@ -40,22 +70,70 @@ export const fetchPortkey = async (
   credentials: any,
   data: any,
   timeout?: number
-) => {
+): Promise<{ response: PIIResponse[]; log: LogObject }> => {
   const options = {
     headers: {
       'x-portkey-api-key': credentials.apiKey,
     },
   };
 
-  if (getRuntimeKey() === 'workerd' && env.portkeyGuardrails) {
-    return postWithCloudflareServiceBinding(
-      `${BASE_URL}${endpoint}`,
-      data,
-      env.portkeyGuardrails,
-      options,
-      timeout
-    );
+  const url = `${BASE_URL}${endpoint}`;
+  const startTime = Date.now();
+  let responseData: any;
+  let responseStatus: number = 200;
+  let responseHeaders: Record<string, any> = {};
+  let error: any = null;
+
+  try {
+    if (getRuntimeKey() === 'workerd' && env.portkeyGuardrails) {
+      responseData = await postWithCloudflareServiceBinding(
+        url,
+        data,
+        env.portkeyGuardrails,
+        options,
+        timeout
+      );
+    } else {
+      responseData = await post(url, data, options, timeout);
+    }
+  } catch (e: any) {
+    error = e;
+    responseStatus = e.response?.status || 500;
+    responseHeaders = {};
+    responseData = e.response?.body || { error: e.message };
   }
 
-  return post(`${BASE_URL}${endpoint}`, data, options, timeout);
+  const endTime = Date.now();
+  const responseTime = endTime - startTime;
+
+  const log: LogObject = {
+    request: {
+      url,
+      method: 'POST',
+      headers: options.headers,
+      body: data,
+    },
+    response: {
+      status: responseStatus,
+      headers: responseHeaders,
+      body: responseData,
+      response_time: responseTime,
+      streamingMode: false,
+    },
+    metadata: {
+      spanId: generateSpanId(),
+      spanName: 'Portkey PII Check',
+    },
+    createdAt: new Date().toISOString(),
+  };
+
+  if (error) {
+    throw { ...error, log };
+  }
+
+  return { response: responseData, log };
 };
+
+function generateSpanId(): string {
+  return Math.random().toString(36).substring(2, 10);
+}