(EAI-1117): Guardrail fix (#781)

* thx claude

* Remove unused guardrail fxn

---------

Co-authored-by: Ben Perlmutter <[email protected]>
Co-authored-by: Helen Schawe <[email protected]>

Author: 3 people

packages/chatbot-server-mongodb-public/src/processors/generateResponseWithSearchTool.test.ts

@@ -272,7 +272,34 @@ describe("generateResponseWithSearchTool", () => {
       ) as UserMessage;
       expect(userMessage.customData).toMatchObject(searchToolMockArgs);
     });
+    it("should not generate until guardrail has resolved (reject)", async () => {
+      const generateResponse = makeGenerateResponseWithSearchTool({
+        ...makeMakeGenerateResponseWithSearchToolArgs(),
+        inputGuardrail: async () => {
+          // sleep for 2 seconds
+          await new Promise((resolve) => setTimeout(resolve, 2000));
+          return mockGuardrailRejectResult;
+        },
+      });
 
+      const result = await generateResponse(generateResponseBaseArgs);
+
+      expectGuardrailRejectResult(result);
+    });
+    it("should not generate until guardrail has resolved (pass)", async () => {
+      const generateResponse = makeGenerateResponseWithSearchTool({
+        ...makeMakeGenerateResponseWithSearchToolArgs(),
+        inputGuardrail: async () => {
+          // sleep for 2 seconds
+          await new Promise((resolve) => setTimeout(resolve, 2000));
+          return mockGuardrailPassResult;
+        },
+      });
+
+      const result = await generateResponse(generateResponseBaseArgs);
+
+      expectSuccessfulResult(result);
+    });
     describe("non-streaming", () => {
       test("should handle successful generation non-streaming", async () => {
         const generateResponse = makeGenerateResponseWithSearchTool(

packages/chatbot-server-mongodb-public/src/processors/generateResponseWithSearchTool.ts

@@ -26,7 +26,6 @@ import {
   MakeReferenceLinksFunc,
   makeDefaultReferenceLinks,
   GenerateResponse,
-  withAbortControllerGuardrail,
   GenerateResponseReturnValue,
   InputGuardrailResult,
 } from "mongodb-chatbot-server";
@@ -128,14 +127,26 @@ export function makeGenerateResponseWithSearchTool({
 
       const references: References = [];
       let userMessageCustomData: Partial<MongoDbSearchToolArgs> = {};
-      const { result, guardrailResult } = await withAbortControllerGuardrail(
-        async (controller) => {
-          // Pass the tools as a separate parameter
+
+      // Create an AbortController for the generation
+      const generationController = new AbortController();
+      let guardrailRejected = false;
+
+      // Start guardrail check immediately and monitor it
+      const guardrailMonitor = inputGuardrailPromise?.then((result) => {
+        if (result?.rejected) {
+          guardrailRejected = true;
+          generationController.abort();
+        }
+        return result;
+      });
+
+      // Start generation immediately (in parallel with guardrail)
+      const generationPromise = (async () => {
+        try {
           const result = streamText({
             ...generationArgs,
-            // Abort the stream if the guardrail AbortController is triggered
-            abortSignal: controller.signal,
-            // Add the search tool results to the references
+            abortSignal: generationController.signal,
             onStepFinish: async ({ toolResults, toolCalls }) => {
               toolCalls?.forEach((toolCall) => {
                 if (toolCall.toolName === SEARCH_TOOL_NAME) {
@@ -159,10 +170,13 @@ export function makeGenerateResponseWithSearchTool({
             },
           });
 
+          // Process the stream
           for await (const chunk of result.fullStream) {
-            if (controller.signal.aborted) {
+            // Check if we should abort due to guardrail rejection
+            if (generationController.signal.aborted) {
               break;
             }
+
             switch (chunk.type) {
               case "text-delta":
                 if (shouldStream) {
@@ -185,22 +199,32 @@ export function makeGenerateResponseWithSearchTool({
                 break;
             }
           }
-          try {
-            if (references.length > 0) {
-              if (shouldStream) {
-                dataStreamer?.streamData({
-                  data: references,
-                  type: "references",
-                });
-              }
+
+          // Stream references if we have any and weren't aborted
+          if (references.length > 0 && !generationController.signal.aborted) {
+            if (shouldStream) {
+              dataStreamer?.streamData({
+                data: references,
+                type: "references",
+              });
             }
-            return result;
-          } catch (error: unknown) {
-            throw new Error(typeof error === "string" ? error : String(error));
           }
-        },
-        inputGuardrailPromise
-      );
+
+          return result;
+        } catch (error: unknown) {
+          // If aborted due to guardrail, return null
+          if (generationController.signal.aborted && guardrailRejected) {
+            return null;
+          }
+          throw new Error(typeof error === "string" ? error : String(error));
+        }
+      })();
+
+      // Wait for both to complete
+      const [guardrailResult, result] = await Promise.all([
+        guardrailMonitor ?? Promise.resolve(undefined),
+        generationPromise,
+      ]);
 
       // If the guardrail rejected the query,
       // return the LLM refusal message

packages/mongodb-chatbot-server/src/processors/InputGuardrail.test.ts

@@ -19,40 +19,3 @@ export type InputGuardrail<
 > = (
   generateResponseParams: GenerateResponseParams
 ) => Promise<InputGuardrailResult<Metadata>>;
-
-export function withAbortControllerGuardrail<T>(
-  fn: (abortController: AbortController) => Promise<T>,
-  guardrailPromise?: Promise<InputGuardrailResult>
-): Promise<{
-  result: T | null;
-  guardrailResult: InputGuardrailResult | undefined;
-}> {
-  const abortController = new AbortController();
-  return (async () => {
-    try {
-      // Run both the main function and guardrail function in parallel
-      const [result, guardrailResult] = await Promise.all([
-        fn(abortController),
-        guardrailPromise
-          ?.then((guardrailResult) => {
-            if (guardrailResult.rejected) {
-              abortController.abort();
-            }
-            return guardrailResult;
-          })
-          .catch((error) => {
-            abortController.abort();
-            return { ...guardrailFailedResult, metadata: { error } };
-          }),
-      ]);
-
-      return { result, guardrailResult };
-    } catch (error) {
-      // If an unexpected error occurs, abort any ongoing operations
-      if (!abortController.signal.aborted) {
-        abortController.abort();
-      }
-      throw error;
-    }
-  })();
-}

packages/mongodb-chatbot-server/src/processors/InputGuardrail.ts

@@ -288,7 +288,5 @@ export const defaultConversationConstants: ConversationConstants = {
 
 Please try to rephrase your message. Adding more details can help me respond with a relevant answer.`,
   LLM_NOT_WORKING: `Unfortunately, my chat functionality is not working at the moment,
-so I cannot respond to your message. Please try again later.
-
-However, here are some links that might provide some helpful information for your message:`,
+so I cannot respond to your message. Please try again later.`,
 };