From 77f6262290b09da2f7c4297488b28b66d5ed2a27 Mon Sep 17 00:00:00 2001
From: Sean Hatfield <seanhatfield5@gmail.com>
Date: Thu, 5 Jun 2025 14:46:44 -0700
Subject: [PATCH] Direct output for agent flows (#3873)

* wip: create direct output switch on last block and send response to ui

* lint

* Return flow on direct output enabled
prevent new blocks below direct output block
Update executor/aibitat to handle skipping of handler outputs

* dev build

---------

Co-authored-by: Timothy Carambat <rambat1010@gmail.com>
---
 .github/workflows/dev-build.yaml              |  2 +-
 .../Admin/AgentBuilder/AddBlockMenu/index.jsx | 17 +++++++
 .../Admin/AgentBuilder/BlockList/index.jsx    | 49 +++++++++++++++++++
 .../src/pages/Admin/AgentBuilder/index.jsx    |  1 +
 server/utils/agentFlows/executor.js           | 12 +++++
 server/utils/agentFlows/flowTypes.js          | 25 ++++++++++
 server/utils/agentFlows/index.js              | 21 ++++++--
 server/utils/agents/aibitat/index.js          | 39 +++++++++++++--
 8 files changed, 157 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/dev-build.yaml b/.github/workflows/dev-build.yaml
index 0e2e15c3..679cf7e7 100644
--- a/.github/workflows/dev-build.yaml
+++ b/.github/workflows/dev-build.yaml
@@ -6,7 +6,7 @@ concurrency:
 
 on:
   push:
-    branches: ['warn-bad-docker-command'] # put your current branch to create a build. Core team only.
+    branches: ['3872-feat-direct-output-to-chat-from-agent-flows'] # put your current branch to create a build. Core team only.
     paths-ignore:
       - '**.md'
       - 'cloud-deployments/*'
diff --git a/frontend/src/pages/Admin/AgentBuilder/AddBlockMenu/index.jsx b/frontend/src/pages/Admin/AgentBuilder/AddBlockMenu/index.jsx
index 777f521c..02041825 100644
--- a/frontend/src/pages/Admin/AgentBuilder/AddBlockMenu/index.jsx
+++ b/frontend/src/pages/Admin/AgentBuilder/AddBlockMenu/index.jsx
@@ -2,7 +2,23 @@ import React, { useRef, useEffect } from "react";
 import { Plus, CaretDown } from "@phosphor-icons/react";
 import { BLOCK_TYPES, BLOCK_INFO } from "../BlockList";
 
+/**
+ * Check if the last configurable block has direct output disabled or undefined
+ * If this property is true then you cannot add a new block after it.
+ * @param {Array} blocks - The blocks array
+ * @returns {Boolean} True if the last configurable block has direct output disabled, false otherwise
+ */
+function checkIfCanAddBlock(blocks) {
+  const lastConfigurableBlock = blocks[blocks.length - 2];
+  if (!lastConfigurableBlock) return true;
+  return (
+    lastConfigurableBlock?.config?.directOutput === false ||
+    lastConfigurableBlock?.config?.directOutput === undefined
+  );
+}
+
 export default function AddBlockMenu({
+  blocks,
   showBlockMenu,
   setShowBlockMenu,
   addBlock,
@@ -22,6 +38,7 @@ export default function AddBlockMenu({
     };
   }, [setShowBlockMenu]);
 
+  if (checkIfCanAddBlock(blocks) === false) return null;
   return (
     <div className="relative mt-4 w-[280px] mx-auto pb-[50%]" ref={menuRef}>
       <button
diff --git a/frontend/src/pages/Admin/AgentBuilder/BlockList/index.jsx b/frontend/src/pages/Admin/AgentBuilder/BlockList/index.jsx
index eaaf05b8..880eb2fa 100644
--- a/frontend/src/pages/Admin/AgentBuilder/BlockList/index.jsx
+++ b/frontend/src/pages/Admin/AgentBuilder/BlockList/index.jsx
@@ -65,6 +65,7 @@ const BLOCK_INFO = {
       body: "",
       formData: [],
       responseVariable: "",
+      directOutput: false,
     },
     getSummary: (config) =>
       `${config.method || "GET"} ${config.url || "(no URL)"}`,
@@ -116,6 +117,7 @@ const BLOCK_INFO = {
     defaultConfig: {
       instruction: "",
       resultVariable: "",
+      directOutput: false,
     },
     getSummary: (config) => config.instruction || "No instruction",
   },
@@ -128,6 +130,7 @@ const BLOCK_INFO = {
       captureAs: "text",
       querySelector: "",
       resultVariable: "",
+      directOutput: false,
     },
     getSummary: (config) => config.url || "No URL specified",
   },
@@ -152,6 +155,7 @@ export default function BlockList({
   refs,
 }) {
   const renderBlockConfig = (block) => {
+    const isLastConfigurableBlock = blocks[blocks.length - 2]?.id === block.id;
     const props = {
       config: block.config,
       onConfigChange: (config) => updateBlockConfig(block.id, config),
@@ -159,6 +163,51 @@ export default function BlockList({
       onDeleteVariable,
     };
 
+    // Direct output switch to the last configurable block before finish
+    if (
+      isLastConfigurableBlock &&
+      block.type !== BLOCK_TYPES.START &&
+      block.type !== BLOCK_TYPES.FLOW_INFO
+    ) {
+      return (
+        <div className="space-y-4">
+          {renderBlockConfigContent(block, props)}
+          <div className="flex justify-between items-center pt-4 border-t border-white/10">
+            <div>
+              <label className="block text-sm font-medium text-theme-text-primary">
+                Direct Output
+              </label>
+              <p className="text-xs text-theme-text-secondary">
+                The output of this block will be returned directly to the chat.
+                <br />
+                This will prevent any further tool calls from being also being
+                executed.
+              </p>
+            </div>
+            <label className="relative inline-flex cursor-pointer items-center">
+              <input
+                type="checkbox"
+                checked={props.config.directOutput || false}
+                onChange={(e) =>
+                  props.onConfigChange({
+                    ...props.config,
+                    directOutput: e.target.checked,
+                  })
+                }
+                className="peer sr-only"
+                aria-label="Toggle direct output"
+              />
+              <div className="pointer-events-none peer h-6 w-11 rounded-full bg-[#CFCFD0] after:absolute after:left-[2px] after:top-[2px] after:h-5 after:w-5 after:rounded-full after:shadow-xl after:border-none after:bg-white after:box-shadow-md after:transition-all after:content-[''] peer-checked:bg-[#32D583] peer-checked:after:translate-x-full peer-checked:after:border-white peer-focus:outline-none peer-focus:ring-4 peer-focus:ring-transparent"></div>
+            </label>
+          </div>
+        </div>
+      );
+    }
+
+    return renderBlockConfigContent(block, props);
+  };
+
+  const renderBlockConfigContent = (block, props) => {
     switch (block.type) {
       case BLOCK_TYPES.FLOW_INFO:
         return <FlowInfoNode {...props} ref={refs} />;
diff --git a/frontend/src/pages/Admin/AgentBuilder/index.jsx b/frontend/src/pages/Admin/AgentBuilder/index.jsx
index 930a575a..38385298 100644
--- a/frontend/src/pages/Admin/AgentBuilder/index.jsx
+++ b/frontend/src/pages/Admin/AgentBuilder/index.jsx
@@ -338,6 +338,7 @@ export default function AgentBuilder() {
             />
 
             <AddBlockMenu
+              blocks={blocks}
               showBlockMenu={showBlockMenu}
               setShowBlockMenu={setShowBlockMenu}
               addBlock={addBlock}
diff --git a/server/utils/agentFlows/executor.js b/server/utils/agentFlows/executor.js
index 5b9b4c8e..81126e2e 100644
--- a/server/utils/agentFlows/executor.js
+++ b/server/utils/agentFlows/executor.js
@@ -186,6 +186,8 @@ class FlowExecutor {
       this.variables[varName] = result;
     }
 
+    // If directOutput is true, mark this result for direct output
+    if (config.directOutput) result = { directOutput: true, result };
     return result;
   }
 
@@ -210,10 +212,19 @@ class FlowExecutor {
     this.aibitat = aibitat;
     this.attachLogging(aibitat?.introspect, aibitat?.handlerProps?.log);
     const results = [];
+    let directOutputResult = null;
 
     for (const step of flow.config.steps) {
       try {
         const result = await this.executeStep(step);
+
+        // If the step has directOutput, stop processing and return the result
+        // so that no other steps are executed or processed
+        if (result?.directOutput) {
+          directOutputResult = result.result;
+          break;
+        }
+
         results.push({ success: true, result });
       } catch (error) {
         results.push({ success: false, error: error.message });
@@ -225,6 +236,7 @@ class FlowExecutor {
       success: results.every((r) => r.success),
       results,
       variables: this.variables,
+      directOutput: directOutputResult,
     };
   }
 }
diff --git a/server/utils/agentFlows/flowTypes.js b/server/utils/agentFlows/flowTypes.js
index f4bd2a6d..89ec8a67 100644
--- a/server/utils/agentFlows/flowTypes.js
+++ b/server/utils/agentFlows/flowTypes.js
@@ -33,6 +33,11 @@ const FLOW_TYPES = {
         type: "string",
         description: "Variable to store the response",
       },
+      directOutput: {
+        type: "boolean",
+        description:
+          "Whether to return the response directly to the user without LLM processing",
+      },
     },
     examples: [
       {
@@ -60,6 +65,11 @@ const FLOW_TYPES = {
         type: "string",
         description: "Variable to store the result",
       },
+      directOutput: {
+        type: "boolean",
+        description:
+          "Whether to return the result directly to the user without LLM processing",
+      },
     },
   },
   FILE: {
@@ -79,6 +89,11 @@ const FLOW_TYPES = {
         type: "string",
         description: "Variable to store the result",
       },
+      directOutput: {
+        type: "boolean",
+        description:
+          "Whether to return the result directly to the user without LLM processing",
+      },
     },
   },
   CODE: {
@@ -94,6 +109,11 @@ const FLOW_TYPES = {
         type: "string",
         description: "Variable to store the result",
       },
+      directOutput: {
+        type: "boolean",
+        description:
+          "Whether to return the result directly to the user without LLM processing",
+      },
     },
   },
   LLM_INSTRUCTION: {
@@ -122,6 +142,11 @@ const FLOW_TYPES = {
         type: "string",
         description: "Variable to store the scraped content",
       },
+      directOutput: {
+        type: "boolean",
+        description:
+          "Whether to return the scraped content directly to the user without LLM processing",
+      },
     },
   },
 };
diff --git a/server/utils/agentFlows/index.js b/server/utils/agentFlows/index.js
index af0e19eb..f759898e 100644
--- a/server/utils/agentFlows/index.js
+++ b/server/utils/agentFlows/index.js
@@ -218,9 +218,15 @@ class AgentFlows {
                 return `Flow execution failed: ${result.results[0]?.error || "Unknown error"}`;
               }
               aibitat.introspect(`${flow.name} completed successfully`);
-              return typeof result === "object"
-                ? JSON.stringify(result)
-                : String(result);
+
+              // If the flow result has directOutput, return it
+              // as the aibitat result so that no other processing is done
+              if (!!result.directOutput) {
+                aibitat.skipHandleExecution = true;
+                return AgentFlows.stringifyResult(result.directOutput);
+              }
+
+              return AgentFlows.stringifyResult(result);
             },
           });
         },
@@ -228,6 +234,15 @@ class AgentFlows {
       flowName: flow.name,
     };
   }
+
+  /**
+   * Stringify the result of a flow execution or return the input as is
+   * @param {Object|string} input - The result to stringify
+   * @returns {string} The stringified result
+   */
+  static stringifyResult(input) {
+    return typeof input === "object" ? JSON.stringify(input) : String(input);
+  }
 }
 
 module.exports.AgentFlows = AgentFlows;
diff --git a/server/utils/agents/aibitat/index.js b/server/utils/agents/aibitat/index.js
index ff1aab45..6e069def 100644
--- a/server/utils/agents/aibitat/index.js
+++ b/server/utils/agents/aibitat/index.js
@@ -12,6 +12,19 @@ const { Telemetry } = require("../../../models/telemetry.js");
 class AIbitat {
   emitter = new EventEmitter();
 
+  /**
+   * Temporary flag to skip the handleExecution function
+   * This is used to return the result of a flow execution directly to the chat
+   * without going through the handleExecution function (resulting in more LLM processing)
+   *
+   * Setting Skip execution to true will prevent any further tool calls from being executed.
+   * This is useful for flow executions that need to return a result directly to the chat but
+   * can also prevent tool-call chaining.
+   *
+   * @type {boolean}
+   */
+  skipHandleExecution = false;
+
   provider = null;
   defaultProvider = null;
   defaultInterrupt;
@@ -618,19 +631,35 @@ ${this.getHistory({ to: route.to })
       // Execute the function and return the result to the provider
       fn.caller = byAgent || "agent";
 
-      // For OSS LLMs we really need to keep tabs on what they are calling
-      // so we can log it here.
+      // If provider is verbose, log the tool call to the frontend
       if (provider?.verbose) {
         this?.introspect?.(
           `[debug]: ${fn.caller} is attempting to call \`${name}\` tool`
         );
-        this.handlerProps.log(
-          `[debug]: ${fn.caller} is attempting to call \`${name}\` tool`
-        );
       }
 
+      // Always log the tool call to the console for debugging purposes
+      this.handlerProps?.log?.(
+        `[debug]: ${fn.caller} is attempting to call \`${name}\` tool`
+      );
+
       const result = await fn.handler(args);
       Telemetry.sendTelemetry("agent_tool_call", { tool: name }, null, true);
+
+      // If the tool call has direct output enabled, return the result directly to the chat
+      // without any further processing and no further tool calls will be run.
+      if (this.skipHandleExecution) {
+        this.skipHandleExecution = false; // reset the flag to prevent next tool call from being skipped
+        this?.introspect?.(
+          `The tool call has direct output enabled! The result will be returned directly to the chat without any further processing and no further tool calls will be run.`
+        );
+        this?.introspect?.(`Tool use completed.`);
+        this.handlerProps?.log?.(
+          `${fn.caller} tool call resulted in direct output! Returning raw result as string. NO MORE TOOL CALLS WILL BE EXECUTED.`
+        );
+        return result;
+      }
+
       return await this.handleExecution(
         provider,
         [