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>
2025-06-05 14:46:44 -07:00 · 2025-06-05 14:46:44 -07:00 · 77f6262290
commit 77f6262290
parent 8dd328cbea
8 changed files with 157 additions and 9 deletions
--- 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/*'
--- 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
--- 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} />;
--- 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}
--- 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,
    };
  }
 }
--- 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",
+      },
    },
  },
 };
--- 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;
--- 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,
        [