feat: improve LLM and prompt node functionality

LLMGenerateNode improvements: - Add collapsible parameters section with temperature and max tokens sliders - Add copy to clipboard button for output text - Clean up UI organization PromptNode improvements: - Add text input handle to receive incoming text connections - Auto-populate prompt when connected to a text source (e.g., another LLM) ConnectionDropMenu: - Add Prompt as a target option for text connections ProjectSetupModal: - Add max tokens slider for LLM default settings Test updates: - Update mocks to include edges and getConnectedInputs - Fix tests for collapsible parameters section - Update wrap-around test for new menu options Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
6 months ago · fa481498cd
7 changed files with 166 additions and 23 deletions
--- a/src/components/ConnectionDropMenu.tsx
+++ b/src/components/ConnectionDropMenu.tsx
@ -73,6 +73,15 @@ const IMAGE_TARGET_OPTIONS: MenuOption[] = [
 ];

 const TEXT_TARGET_OPTIONS: MenuOption[] = [
+  {
+    type: "prompt",
+    label: "Prompt",
+    icon: (
+      <svg className="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={1.5}>
+        <path strokeLinecap="round" strokeLinejoin="round" d="M19.5 14.25v-2.625a3.375 3.375 0 00-3.375-3.375h-1.5A1.125 1.125 0 0113.5 7.125v-1.5a3.375 3.375 0 00-3.375-3.375H8.25m0 12.75h7.5m-7.5 3H12M10.5 2.25H5.625c-.621 0-1.125.504-1.125 1.125v17.25c0 .621.504 1.125 1.125 1.125h12.75c.621 0 1.125-.504 1.125-1.125V11.25a9 9 0 00-9-9z" />
+      </svg>
+    ),
+  },
  {
    type: "nanoBanana",
    label: "Generate Image",
--- a/src/components/ProjectSetupModal.tsx
+++ b/src/components/ProjectSetupModal.tsx
@ -778,6 +778,27 @@ export function ProjectSetupModal({
                    className="flex-1 h-1 bg-neutral-700 rounded-lg appearance-none cursor-pointer accent-neutral-400"
                  />
                </div>
+
+                {/* Max Tokens slider */}
+                <div className="flex items-center gap-2">
+                  <label className="text-xs text-neutral-400 w-20">
+                    Tokens: {(localNodeDefaults.llm?.maxTokens ?? 8192).toLocaleString()}
+                  </label>
+                  <input
+                    type="range"
+                    min="256"
+                    max="16384"
+                    step="256"
+                    value={localNodeDefaults.llm?.maxTokens ?? 8192}
+                    onChange={(e) => {
+                      setLocalNodeDefaults(prev => ({
+                        ...prev,
+                        llm: { ...prev.llm, maxTokens: parseInt(e.target.value, 10) }
+                      }));
+                    }}
+                    className="flex-1 h-1 bg-neutral-700 rounded-lg appearance-none cursor-pointer accent-neutral-400"
+                  />
+                </div>
              </div>
            </div>

--- a/src/components/tests/ConnectionDropMenu.test.tsx
+++ b/src/components/tests/ConnectionDropMenu.test.tsx
@ -226,14 +226,15 @@ describe("ConnectionDropMenu", () => {
    it("should wrap around when navigating past last item", () => {
      render(<ConnectionDropMenu {...defaultProps} handleType="text" connectionType="source" />);

-      // Text options: nanoBanana, generateVideo, llmGenerate (3 items)
-      // Navigate down 3 times to wrap to first
+      // Text target options: Prompt, nanoBanana, generateVideo, llmGenerate (4 items)
+      // Navigate down 4 times to wrap to first
+      fireEvent.keyDown(document, { key: "ArrowDown" });
      fireEvent.keyDown(document, { key: "ArrowDown" });
      fireEvent.keyDown(document, { key: "ArrowDown" });
      fireEvent.keyDown(document, { key: "ArrowDown" });

-      // Should be back on first item
-      const firstButton = screen.getByText("Generate Image").closest("button");
+      // Should be back on first item (Prompt)
+      const firstButton = screen.getByText("Prompt").closest("button");
      expect(firstButton).toHaveClass("bg-neutral-700");
    });
  });
--- a/src/components/tests/LLMGenerateNode.test.tsx
+++ b/src/components/tests/LLMGenerateNode.test.tsx
@ -194,14 +194,18 @@ describe("LLMGenerateNode", () => {
  });

  describe("Temperature Slider", () => {
-    it("should render temperature slider with current value", () => {
+    it("should render temperature slider with current value after expanding parameters", () => {
      render(
        <TestWrapper>
          <LLMGenerateNode {...createNodeProps({ temperature: 0.7 })} />
        </TestWrapper>
      );

-      expect(screen.getByText("Temp: 0.7")).toBeInTheDocument();
+      // Expand the Parameters section
+      const parametersButton = screen.getByText("Parameters");
+      fireEvent.click(parametersButton);
+
+      expect(screen.getByText("Temperature: 0.7")).toBeInTheDocument();
    });

    it("should call updateNodeData when temperature is changed", () => {
@ -211,8 +215,14 @@ describe("LLMGenerateNode", () => {
        </TestWrapper>
      );

-      const slider = screen.getByRole("slider");
-      fireEvent.change(slider, { target: { value: "1.5" } });
+      // Expand the Parameters section
+      const parametersButton = screen.getByText("Parameters");
+      fireEvent.click(parametersButton);
+
+      // Get the first slider (temperature) - there are two sliders: temperature and maxTokens
+      const sliders = screen.getAllByRole("slider");
+      const temperatureSlider = sliders[0];
+      fireEvent.change(temperatureSlider, { target: { value: "1.5" } });

      expect(mockUpdateNodeData).toHaveBeenCalledWith("test-llm-1", {
        temperature: 1.5,
--- a/src/components/tests/PromptNode.test.tsx
+++ b/src/components/tests/PromptNode.test.tsx
@ -17,6 +17,8 @@ vi.mock("@/store/workflowStore", () => ({
      currentNodeId: null,
      groups: {},
      nodes: [],
+      edges: [],
+      getConnectedInputs: vi.fn(() => ({ images: [], videos: [], text: null, dynamicInputs: {} })),
      getNodesWithComments: vi.fn(() => []),
      markCommentViewed: vi.fn(),
      setNavigationTarget: vi.fn(),
--- a/src/components/nodes/LLMGenerateNode.tsx
+++ b/src/components/nodes/LLMGenerateNode.tsx
@ -1,6 +1,6 @@
 "use client";

-import { useCallback } from "react";
+import { useCallback, useState } from "react";
 import { Handle, Position, NodeProps, Node } from "@xyflow/react";
 import { BaseNode } from "./BaseNode";
 import { useCommentNavigation } from "@/hooks/useCommentNavigation";
@ -57,6 +57,15 @@ export function LLMGenerateNode({ id, data, selected }: NodeProps<LLMGenerateNod
    [id, updateNodeData]
  );

+  const handleMaxTokensChange = useCallback(
+    (e: React.ChangeEvent<HTMLInputElement>) => {
+      updateNodeData(id, { maxTokens: parseInt(e.target.value, 10) });
+    },
+    [id, updateNodeData]
+  );
+
+  const [showParams, setShowParams] = useState(false);
+
  const regenerateNode = useWorkflowStore((state) => state.regenerateNode);
  const isRunning = useWorkflowStore((state) => state.isRunning);

@ -68,6 +77,20 @@ export function LLMGenerateNode({ id, data, selected }: NodeProps<LLMGenerateNod
    updateNodeData(id, { outputText: null, status: "idle", error: null });
  }, [id, updateNodeData]);

+  const [copied, setCopied] = useState(false);
+
+  const handleCopyOutput = useCallback(async () => {
+    if (nodeData.outputText) {
+      try {
+        await navigator.clipboard.writeText(nodeData.outputText);
+        setCopied(true);
+        setTimeout(() => setCopied(false), 1500);
+      } catch (err) {
+        console.error("Failed to copy text:", err);
+      }
+    }
+  }, [nodeData.outputText]);
+
  const provider = nodeData.provider || "google";
  const availableModels = MODELS[provider] || MODELS.google;
  const model = availableModels.some(m => m.value === nodeData.model)
@ -145,6 +168,21 @@ export function LLMGenerateNode({ id, data, selected }: NodeProps<LLMGenerateNod
                {nodeData.outputText}
              </p>
              <div className="absolute top-1 right-1 flex gap-1">
+                <button
+                  onClick={handleCopyOutput}
+                  className={`w-5 h-5 ${copied ? "bg-green-600/80" : "bg-neutral-900/80 hover:bg-neutral-700/80"} rounded flex items-center justify-center text-neutral-400 hover:text-white transition-colors`}
+                  title={copied ? "Copied!" : "Copy to clipboard"}
+                >
+                  {copied ? (
+                    <svg className="w-3 h-3 text-white" fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={2}>
+                      <path strokeLinecap="round" strokeLinejoin="round" d="M5 13l4 4L19 7" />
+                    </svg>
+                  ) : (
+                    <svg className="w-3 h-3" fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={2}>
+                      <path strokeLinecap="round" strokeLinejoin="round" d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z" />
+                    </svg>
+                  )}
+                </button>
                <button
                  onClick={handleRegenerate}
                  disabled={isRunning}
@ -201,18 +239,53 @@ export function LLMGenerateNode({ id, data, selected }: NodeProps<LLMGenerateNod
          ))}
        </select>

-        {/* Temperature slider */}
-        <div className="flex flex-col gap-0.5 shrink-0">
-          <label className="text-[9px] text-neutral-500">Temp: {nodeData.temperature.toFixed(1)}</label>
-          <input
-            type="range"
-            min="0"
-            max="2"
-            step="0.1"
-            value={nodeData.temperature}
-            onChange={handleTemperatureChange}
-            className="w-full h-1 bg-neutral-700 rounded-lg appearance-none cursor-pointer accent-neutral-400"
-          />
+        {/* Collapsible parameters section */}
+        <div className="shrink-0">
+          <button
+            onClick={() => setShowParams(!showParams)}
+            className="w-full flex items-center justify-between text-[9px] text-neutral-400 hover:text-neutral-300 py-1"
+          >
+            <span>Parameters</span>
+            <svg
+              className={`w-3 h-3 transition-transform ${showParams ? "rotate-180" : ""}`}
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+              strokeWidth={2}
+            >
+              <path strokeLinecap="round" strokeLinejoin="round" d="M19 9l-7 7-7-7" />
+            </svg>
+          </button>
+          {showParams && (
+            <div className="flex flex-col gap-2 pt-1 border-t border-neutral-700/50">
+              {/* Temperature slider */}
+              <div className="flex flex-col gap-0.5">
+                <label className="text-[9px] text-neutral-500">Temperature: {nodeData.temperature.toFixed(1)}</label>
+                <input
+                  type="range"
+                  min="0"
+                  max="2"
+                  step="0.1"
+                  value={nodeData.temperature}
+                  onChange={handleTemperatureChange}
+                  className="nodrag w-full h-1 bg-neutral-700 rounded-lg appearance-none cursor-pointer accent-neutral-400"
+                />
+              </div>
+              {/* Max tokens slider */}
+              <div className="flex flex-col gap-0.5">
+                <label className="text-[9px] text-neutral-500">Max Tokens: {nodeData.maxTokens.toLocaleString()}</label>
+                <input
+                  type="range"
+                  min="256"
+                  max="16384"
+                  step="256"
+                  value={nodeData.maxTokens}
+                  onChange={handleMaxTokensChange}
+                  className="nodrag w-full h-1 bg-neutral-700 rounded-lg appearance-none cursor-pointer accent-neutral-400"
+                />
+              </div>
+            </div>
+          )}
        </div>
      </div>
    </BaseNode>
--- a/src/components/nodes/PromptNode.tsx
+++ b/src/components/nodes/PromptNode.tsx
@ -1,6 +1,6 @@
 "use client";

-import { useCallback, useState, useEffect } from "react";
+import { useCallback, useState, useEffect, useMemo } from "react";
 import { createPortal } from "react-dom";
 import { Handle, Position, NodeProps, Node } from "@xyflow/react";
 import { BaseNode } from "./BaseNode";
@ -17,12 +17,29 @@ export function PromptNode({ id, data, selected }: NodeProps<PromptNodeType>) {
  const updateNodeData = useWorkflowStore((state) => state.updateNodeData);
  const incrementModalCount = useWorkflowStore((state) => state.incrementModalCount);
  const decrementModalCount = useWorkflowStore((state) => state.decrementModalCount);
+  const getConnectedInputs = useWorkflowStore((state) => state.getConnectedInputs);
+  const edges = useWorkflowStore((state) => state.edges);
  const [isModalOpenLocal, setIsModalOpenLocal] = useState(false);

  // Local state for prompt to prevent cursor jumping during typing
  const [localPrompt, setLocalPrompt] = useState(nodeData.prompt);
  const [isEditing, setIsEditing] = useState(false);

+  // Check if this node has any incoming text connections
+  const hasIncomingTextConnection = useMemo(() => {
+    return edges.some((edge) => edge.target === id && edge.targetHandle === "text");
+  }, [edges, id]);
+
+  // Get connected text input and update prompt when connection provides text
+  useEffect(() => {
+    if (hasIncomingTextConnection) {
+      const { text } = getConnectedInputs(id);
+      if (text !== null && text !== nodeData.prompt) {
+        updateNodeData(id, { prompt: text });
+      }
+    }
+  }, [hasIncomingTextConnection, id, getConnectedInputs, updateNodeData, nodeData.prompt]);
+
  // Sync from props when not actively editing
  useEffect(() => {
    if (!isEditing) {
@ -78,15 +95,25 @@ export function PromptNode({ id, data, selected }: NodeProps<PromptNodeType>) {
        selected={selected}
        commentNavigation={commentNavigation ?? undefined}
      >
+        {/* Text input handle - for receiving text from LLM nodes */}
+        <Handle
+          type="target"
+          position={Position.Left}
+          id="text"
+          data-handletype="text"
+        />
+
        <textarea
          value={localPrompt}
          onChange={handleChange}
          onFocus={handleFocus}
          onBlur={handleBlur}
-          placeholder="Describe what to generate..."
+          placeholder={hasIncomingTextConnection ? "Receiving text from connected node..." : "Describe what to generate..."}
          className="nodrag nopan nowheel w-full flex-1 min-h-[70px] p-2 text-xs leading-relaxed text-neutral-100 border border-neutral-700 rounded bg-neutral-900/50 resize-none focus:outline-none focus:ring-1 focus:ring-neutral-600 focus:border-neutral-600 placeholder:text-neutral-500"
+          readOnly={hasIncomingTextConnection}
        />

+        {/* Text output handle */}
        <Handle
          type="source"
          position={Position.Right}