Add request, response log for the Portkey PII plugin to the main requ…

plugins/portkey/globals.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);
+}

plugins/portkey/pii.ts

@@ -10,6 +10,7 @@ import {
   setCurrentContentPart,
 } from '../utils';
 import {
+  LogObject,
   PIIResponse,
   PIIResult,
   PORTKEY_ENDPOINTS,
@@ -20,8 +21,8 @@ export async function detectPII(
   textArray: Array<string> | string,
   parameters: any,
   env: Record<string, any>
-): Promise<Promise<PIIResult[]>> {
-  const result: PIIResponse[] = await fetchPortkey(
+): Promise<{ results: PIIResult[]; log: LogObject }> {
+  const { response, log } = await fetchPortkey(
     env,
     PORTKEY_ENDPOINTS.PII,
     parameters.credentials,
@@ -32,7 +33,7 @@ export async function detectPII(
     parameters.timeout
   );
 
-  return result.map((item) => ({
+  const results = response.map((item) => ({
     detectedPIICategories: [
       ...new Set(item.entities.flatMap((entity) => Object.keys(entity.labels))),
     ],
@@ -42,6 +43,8 @@ export async function detectPII(
     })),
     redactedText: item.processed_text,
   }));
+
+  return { results, log };
 }
 
 export const handler: PluginHandler = async (
@@ -53,6 +56,7 @@ export const handler: PluginHandler = async (
   let error = null;
   let verdict = false;
   let data: any = null;
+  let log: LogObject | null = null;
   let transformedData: Record<string, any> = {
     request: {
       json: null,
@@ -70,6 +74,7 @@ export const handler: PluginHandler = async (
         data: null,
         transformedData,
         transformed,
+        log: null,
       };
     }
 
@@ -83,6 +88,7 @@ export const handler: PluginHandler = async (
         data: null,
         transformedData,
         transformed,
+        log: null,
       };
     }
 
@@ -93,6 +99,7 @@ export const handler: PluginHandler = async (
         data: null,
         transformedData,
         transformed,
+        log: null,
       };
     }
 
@@ -103,14 +110,16 @@ export const handler: PluginHandler = async (
         data: null,
         transformedData,
         transformed,
+        log: null,
       };
     }
 
-    let mappedResult = await detectPII(
+    const { results: mappedResult, log: piiLog } = await detectPII(
       textArray,
       parameters,
       options?.env || {}
     );
+    log = piiLog;
 
     const categoriesToCheck = parameters.categories || [];
     const not = parameters.not || false;
@@ -169,6 +178,10 @@ export const handler: PluginHandler = async (
     };
   } catch (e) {
     error = e as Error;
+    // If there's a log in the error, capture it
+    if ((e as any).log) {
+      log = (e as any).log;
+    }
     const text = getText(context, eventType);
     data = {
       explanation: `An error occurred while checking for PII: ${error.message}`,
@@ -182,5 +195,5 @@ export const handler: PluginHandler = async (
     };
   }
 
-  return { error, verdict, data, transformedData, transformed };
+  return { error, verdict, data, transformedData, transformed, log };
 };