mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-24 03:21:47 +08:00
feat(ui): Add "Connect other models" option for model-type handles (#12466)
* Display proper error messages and strip null params from tool calls * add connect other models option * fix model provider selection state * fix tests and ruff * revert mcp changes * [autofix.ci] apply automated fixes * fix handle not appearing properly * fix jest test * [autofix.ci] apply automated fixes * address qa fixes * add tests to validate be and fe * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * ruff style and checker * fix tests failuer --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Cristhian Zanforlin Lousa
committed by
GitHub
GitHub
parent
4bdc87dcc1
commit
e6846ea6e7
@ -0,0 +1,232 @@
|
||||
/**
|
||||
* Unit tests for muted handle state in HandleRenderComponent.
|
||||
*
|
||||
* Model-type handles should be visually muted (invisible) when:
|
||||
* - No wire is connected to the handle
|
||||
* - No connection filter is active (drag/connection mode)
|
||||
*
|
||||
* They should be visible when:
|
||||
* - A wire is connected
|
||||
* - A connection filter is active (e.g., "Connect other models")
|
||||
* - The handle's own node initiated connection mode
|
||||
*/
|
||||
|
||||
describe("model handle muted state", () => {
|
||||
describe("isMuted computation logic", () => {
|
||||
// Pure logic extracted from HandleRenderComponent useMemo
|
||||
const computeIsMuted = (
|
||||
idType: string | undefined,
|
||||
connectedEdge: boolean,
|
||||
filterPresent: boolean,
|
||||
isInConnectionMode: boolean = false,
|
||||
): boolean => {
|
||||
const isModelType = idType === "model";
|
||||
return (
|
||||
isModelType && !connectedEdge && !filterPresent && !isInConnectionMode
|
||||
);
|
||||
};
|
||||
|
||||
it("should be muted when model type has no connection, no filter, and not in connection mode", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsMuted("model", false, false, false);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
|
||||
it("should not be muted when model type has a connected edge", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsMuted("model", true, false, false);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should not be muted when model type has an active filter", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsMuted("model", false, true, false);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should not be muted when model type has both connection and filter", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsMuted("model", true, true, false);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should not be muted when in connection mode (Connect other models selected)", () => {
|
||||
// Arrange & Act — no edge, no filter, but connection mode is active
|
||||
const result = computeIsMuted("model", false, false, true);
|
||||
|
||||
// Assert — handle must stay visible so user can drag a wire to it
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should never be muted for non-model types", () => {
|
||||
// Arrange & Act & Assert
|
||||
expect(computeIsMuted("str", false, false, false)).toBe(false);
|
||||
expect(computeIsMuted("Message", false, false, false)).toBe(false);
|
||||
expect(computeIsMuted(undefined, false, false, false)).toBe(false);
|
||||
expect(computeIsMuted("", false, false, false)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("isOwnModelConnectionMode logic", () => {
|
||||
// Pure logic extracted from HandleRenderComponent useMemo
|
||||
const computeIsOwnModelConnectionMode = (
|
||||
idType: string | undefined,
|
||||
left: boolean,
|
||||
filterTargetNodeId: string | undefined,
|
||||
nodeId: string,
|
||||
): boolean => {
|
||||
return idType === "model" && left && filterTargetNodeId === nodeId;
|
||||
};
|
||||
|
||||
const NODE_ID = "agent-node-123";
|
||||
|
||||
it("should be true when filter targets this node's model input handle", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsOwnModelConnectionMode(
|
||||
"model",
|
||||
true,
|
||||
NODE_ID,
|
||||
NODE_ID,
|
||||
);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
|
||||
it("should be false when filter targets a different node", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsOwnModelConnectionMode(
|
||||
"model",
|
||||
true,
|
||||
"other-node-456",
|
||||
NODE_ID,
|
||||
);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should be false for non-model types", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsOwnModelConnectionMode(
|
||||
"str",
|
||||
true,
|
||||
NODE_ID,
|
||||
NODE_ID,
|
||||
);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should be false for right-side (output) handles", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsOwnModelConnectionMode(
|
||||
"model",
|
||||
false,
|
||||
NODE_ID,
|
||||
NODE_ID,
|
||||
);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it("should be false when no filter is active", () => {
|
||||
// Arrange & Act
|
||||
const result = computeIsOwnModelConnectionMode(
|
||||
"model",
|
||||
true,
|
||||
undefined,
|
||||
NODE_ID,
|
||||
);
|
||||
|
||||
// Assert
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("muted handle visual styling", () => {
|
||||
// Pure logic extracted from HandleContent contentStyle useMemo
|
||||
const computeContentStyle = (
|
||||
isNullHandle: boolean,
|
||||
isMuted: boolean,
|
||||
): { width: string; height: string; opacity: number } => {
|
||||
return {
|
||||
width: isMuted && !isNullHandle ? "6px" : "10px",
|
||||
height: isMuted && !isNullHandle ? "6px" : "10px",
|
||||
opacity: isMuted && !isNullHandle ? 0 : 1,
|
||||
};
|
||||
};
|
||||
|
||||
it("should render invisible when muted and not null", () => {
|
||||
// Arrange & Act
|
||||
const style = computeContentStyle(false, true);
|
||||
|
||||
// Assert
|
||||
expect(style.width).toBe("6px");
|
||||
expect(style.height).toBe("6px");
|
||||
expect(style.opacity).toBe(0);
|
||||
});
|
||||
|
||||
it("should render full size when not muted and not null", () => {
|
||||
// Arrange & Act
|
||||
const style = computeContentStyle(false, false);
|
||||
|
||||
// Assert
|
||||
expect(style.width).toBe("10px");
|
||||
expect(style.height).toBe("10px");
|
||||
expect(style.opacity).toBe(1);
|
||||
});
|
||||
|
||||
it("should use null handle styling when isNullHandle overrides isMuted", () => {
|
||||
// isNullHandle takes priority — size stays 10px, opacity stays 1
|
||||
// (actual null handle styling uses different background/border, tested elsewhere)
|
||||
|
||||
// Arrange & Act
|
||||
const style = computeContentStyle(true, true);
|
||||
|
||||
// Assert
|
||||
expect(style.width).toBe("10px");
|
||||
expect(style.height).toBe("10px");
|
||||
expect(style.opacity).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("muted handle neon shadow", () => {
|
||||
const computeNeonShadow = (
|
||||
isNullHandle: boolean,
|
||||
isMuted: boolean,
|
||||
isActive: boolean,
|
||||
): string => {
|
||||
if (isNullHandle || isMuted) return "none";
|
||||
if (!isActive) return "ring";
|
||||
return "glow";
|
||||
};
|
||||
|
||||
it("should return none when muted", () => {
|
||||
expect(computeNeonShadow(false, true, false)).toBe("none");
|
||||
expect(computeNeonShadow(false, true, true)).toBe("none");
|
||||
});
|
||||
|
||||
it("should return none when null handle", () => {
|
||||
expect(computeNeonShadow(true, false, false)).toBe("none");
|
||||
});
|
||||
|
||||
it("should return ring when not muted and not active", () => {
|
||||
expect(computeNeonShadow(false, false, false)).toBe("ring");
|
||||
});
|
||||
|
||||
it("should return glow when not muted and active", () => {
|
||||
expect(computeNeonShadow(false, false, true)).toBe("glow");
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -3,6 +3,7 @@ import { memo, useCallback, useEffect, useMemo, useState } from "react";
|
||||
import { useShallow } from "zustand/react/shallow";
|
||||
import { useDarkStore } from "@/stores/darkStore";
|
||||
import useFlowStore from "@/stores/flowStore";
|
||||
import type { NodeDataType } from "@/types/flow";
|
||||
import { nodeColorsName } from "@/utils/styleUtils";
|
||||
import ShadTooltip from "../../../../components/common/shadTooltipComponent";
|
||||
import {
|
||||
@ -24,6 +25,7 @@ const BASE_HANDLE_STYLES = {
|
||||
|
||||
const HandleContent = memo(function HandleContent({
|
||||
isNullHandle,
|
||||
isMuted,
|
||||
handleColor,
|
||||
accentForegroundColorName,
|
||||
isHovered,
|
||||
@ -35,6 +37,7 @@ const HandleContent = memo(function HandleContent({
|
||||
nodeId,
|
||||
}: {
|
||||
isNullHandle: boolean;
|
||||
isMuted: boolean;
|
||||
handleColor: string;
|
||||
accentForegroundColorName: string;
|
||||
isHovered: boolean;
|
||||
@ -97,7 +100,7 @@ const HandleContent = memo(function HandleContent({
|
||||
|
||||
const getNeonShadow = useCallback(
|
||||
(color: string, isActive: boolean) => {
|
||||
if (isNullHandle) return "none";
|
||||
if (isNullHandle || isMuted) return "none";
|
||||
if (!isActive) return `0 0 0 3px ${color}`;
|
||||
return [
|
||||
"0 0 0 1px hsl(var(--border))",
|
||||
@ -110,27 +113,29 @@ const HandleContent = memo(function HandleContent({
|
||||
`0 0 20px ${color}`,
|
||||
].join(", ");
|
||||
},
|
||||
[isNullHandle],
|
||||
[isNullHandle, isMuted],
|
||||
);
|
||||
|
||||
const contentStyle = useMemo(
|
||||
() => ({
|
||||
background: isNullHandle ? "hsl(var(--border))" : handleColor,
|
||||
width: "10px",
|
||||
height: "10px",
|
||||
width: isMuted && !isNullHandle ? "6px" : "10px",
|
||||
height: isMuted && !isNullHandle ? "6px" : "10px",
|
||||
transition: "all 0.2s",
|
||||
opacity: isMuted && !isNullHandle ? 0 : 1,
|
||||
boxShadow: getNeonShadow(
|
||||
accentForegroundColorName,
|
||||
isHovered || openHandle,
|
||||
),
|
||||
animation:
|
||||
(isHovered || openHandle) && !isNullHandle
|
||||
(isHovered || openHandle) && !isNullHandle && !isMuted
|
||||
? `pulseNeon-${nodeId} 1.1s ease-in-out infinite`
|
||||
: "none",
|
||||
border: isNullHandle ? "2px solid hsl(var(--muted))" : "none",
|
||||
}),
|
||||
[
|
||||
isNullHandle,
|
||||
isMuted,
|
||||
handleColor,
|
||||
getNeonShadow,
|
||||
accentForegroundColorName,
|
||||
@ -186,6 +191,18 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
|
||||
const edges = useFlowStore((state) => state.edges);
|
||||
|
||||
// Check if this node is in "connect other models" mode
|
||||
const isInConnectionMode = useFlowStore(
|
||||
useCallback(
|
||||
(state) => {
|
||||
if (id?.type !== "model" || !left) return false;
|
||||
const node = state.nodes.find((n) => n.id === nodeId);
|
||||
return (node?.data as NodeDataType)?._connectionMode === true;
|
||||
},
|
||||
[nodeId, id?.type, left],
|
||||
),
|
||||
);
|
||||
|
||||
const {
|
||||
setHandleDragging,
|
||||
setFilterType,
|
||||
@ -234,6 +251,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
filterPresent,
|
||||
currentFilter,
|
||||
isNullHandle,
|
||||
isMuted,
|
||||
handleColor,
|
||||
accentForegroundColorName,
|
||||
} = useMemo(() => {
|
||||
@ -276,8 +294,18 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
const outputType = connectedEdge?.data?.sourceHandle?.output_types?.[0];
|
||||
const connectedColor = (outputType && nodeColorsName[outputType]) || "gray";
|
||||
|
||||
// Model handles that initiated connection mode on this node should not be nulled
|
||||
const isOwnModelConnectionMode =
|
||||
id?.type === "model" && left && filterType?.target === nodeId;
|
||||
|
||||
const isNullHandle =
|
||||
filterPresent && !(openHandle || ownDraggingHandle || ownFilterHandle);
|
||||
filterPresent &&
|
||||
!(
|
||||
openHandle ||
|
||||
ownDraggingHandle ||
|
||||
ownFilterHandle ||
|
||||
isOwnModelConnectionMode
|
||||
);
|
||||
|
||||
// Create a Set from colorName to remove duplicates
|
||||
const colorNameSet = new Set(colorName || []);
|
||||
@ -325,6 +353,10 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
color: handleColorName,
|
||||
};
|
||||
|
||||
const isModelType = id?.type === "model";
|
||||
const isMuted =
|
||||
isModelType && !connectedEdge && !filterPresent && !isInConnectionMode;
|
||||
|
||||
return {
|
||||
sameNode: sameDraggingNode || sameFilterNode,
|
||||
ownHandle: ownDraggingHandle || ownFilterHandle,
|
||||
@ -334,6 +366,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
filterPresent,
|
||||
currentFilter,
|
||||
isNullHandle,
|
||||
isMuted,
|
||||
handleColor,
|
||||
};
|
||||
}, [
|
||||
@ -347,6 +380,8 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
colorName,
|
||||
tooltipTitle,
|
||||
edges,
|
||||
id,
|
||||
isInConnectionMode,
|
||||
]);
|
||||
|
||||
const handleMouseDown = useCallback(
|
||||
@ -441,6 +476,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
|
||||
>
|
||||
<HandleContent
|
||||
isNullHandle={isNullHandle ?? false}
|
||||
isMuted={isMuted ?? false}
|
||||
handleColor={handleColor}
|
||||
accentForegroundColorName={accentForegroundColorName}
|
||||
isHovered={isHovered}
|
||||
|
||||
@ -31,17 +31,20 @@ jest.mock("@/hooks/use-refresh-model-inputs", () => ({
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock("@/stores/flowStore", () => ({
|
||||
__esModule: true,
|
||||
default: {
|
||||
getState: () => ({
|
||||
getNode: jest.fn(),
|
||||
setFilterEdge: jest.fn(),
|
||||
setFilterType: jest.fn(),
|
||||
nodes: [],
|
||||
}),
|
||||
},
|
||||
}));
|
||||
jest.mock("@/stores/flowStore", () => {
|
||||
const state = {
|
||||
getNode: jest.fn(),
|
||||
setNode: jest.fn(),
|
||||
setFilterEdge: jest.fn(),
|
||||
setFilterType: jest.fn(),
|
||||
nodes: [],
|
||||
edges: [],
|
||||
};
|
||||
const hook = (selector?: (s: any) => any) =>
|
||||
selector ? selector(state) : state;
|
||||
hook.getState = () => state;
|
||||
return { __esModule: true, default: hook };
|
||||
});
|
||||
|
||||
jest.mock("@/stores/typesStore", () => ({
|
||||
useTypesStore: {
|
||||
|
||||
@ -0,0 +1,271 @@
|
||||
/**
|
||||
* Unit tests for the connection mode logic in useModelConnectionLogic.
|
||||
*
|
||||
* The hook activates "Connect other models" mode purely on the frontend
|
||||
* by setting filter state in the flow store. No backend calls are made.
|
||||
* These tests validate the pure logic decisions without React rendering.
|
||||
*/
|
||||
|
||||
describe("useModelConnectionLogic", () => {
|
||||
describe("effective input types defaulting", () => {
|
||||
const getEffectiveInputTypes = (inputTypes: string[]): string[] => {
|
||||
return inputTypes.length > 0 ? inputTypes : ["LanguageModel"];
|
||||
};
|
||||
|
||||
it("should default to LanguageModel when input_types is empty", () => {
|
||||
// Arrange & Act
|
||||
const result = getEffectiveInputTypes([]);
|
||||
|
||||
// Assert
|
||||
expect(result).toEqual(["LanguageModel"]);
|
||||
});
|
||||
|
||||
it("should use provided input_types when available", () => {
|
||||
// Arrange & Act
|
||||
const result = getEffectiveInputTypes(["Embeddings"]);
|
||||
|
||||
// Assert
|
||||
expect(result).toEqual(["Embeddings"]);
|
||||
});
|
||||
|
||||
it("should preserve multiple input types", () => {
|
||||
// Arrange & Act
|
||||
const result = getEffectiveInputTypes(["LanguageModel", "ChatModel"]);
|
||||
|
||||
// Assert
|
||||
expect(result).toEqual(["LanguageModel", "ChatModel"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("filter object construction", () => {
|
||||
const NODE_ID = "agent-node-123";
|
||||
|
||||
const buildFilterObj = (nodeId: string, targetHandle: string) => ({
|
||||
source: undefined,
|
||||
sourceHandle: undefined,
|
||||
target: nodeId,
|
||||
targetHandle,
|
||||
type: "LanguageModel",
|
||||
color: "datatype-fuchsia",
|
||||
});
|
||||
|
||||
it("should set target to the node ID", () => {
|
||||
// Arrange & Act
|
||||
const filterObj = buildFilterObj(NODE_ID, "handle-123");
|
||||
|
||||
// Assert
|
||||
expect(filterObj.target).toBe(NODE_ID);
|
||||
expect(filterObj.source).toBeUndefined();
|
||||
expect(filterObj.sourceHandle).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should use LanguageModel as the filter type", () => {
|
||||
// Arrange & Act
|
||||
const filterObj = buildFilterObj(NODE_ID, "handle-123");
|
||||
|
||||
// Assert
|
||||
expect(filterObj.type).toBe("LanguageModel");
|
||||
expect(filterObj.color).toBe("datatype-fuchsia");
|
||||
});
|
||||
});
|
||||
|
||||
describe("option value routing", () => {
|
||||
const shouldActivateConnectionMode = (optionValue: string): boolean => {
|
||||
return optionValue === "connect_other_models";
|
||||
};
|
||||
|
||||
it("should activate for connect_other_models", () => {
|
||||
expect(shouldActivateConnectionMode("connect_other_models")).toBe(true);
|
||||
});
|
||||
|
||||
it("should not activate for other option values", () => {
|
||||
expect(shouldActivateConnectionMode("some_other_option")).toBe(false);
|
||||
expect(shouldActivateConnectionMode("")).toBe(false);
|
||||
expect(shouldActivateConnectionMode("connect")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("selectedModel in connection mode", () => {
|
||||
// Pure logic: connection mode is determined by local state (isConnectionMode),
|
||||
// NOT by a special string value. The template value is cleared to [] when
|
||||
// entering connection mode so the backend doesn't retain stale config.
|
||||
const computeSelectedModel = (
|
||||
isConnectionMode: boolean,
|
||||
value: any,
|
||||
flatOptions: Array<{ name: string; icon?: string; provider?: string }>,
|
||||
externalDisplayName?: string,
|
||||
externalIcon?: string,
|
||||
) => {
|
||||
if (isConnectionMode) {
|
||||
return {
|
||||
name: externalDisplayName || "Connect other models",
|
||||
icon: externalIcon || "CornerDownLeft",
|
||||
provider: "",
|
||||
};
|
||||
}
|
||||
const currentName = value?.[0]?.name;
|
||||
if (!currentName) return null;
|
||||
return flatOptions.find((o) => o.name === currentName) || null;
|
||||
};
|
||||
|
||||
it("should return connection mode display when isConnectionMode is true", () => {
|
||||
// value is [] (cleared) when in connection mode
|
||||
const result = computeSelectedModel(true, [], []);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result!.name).toBe("Connect other models");
|
||||
expect(result!.icon).toBe("CornerDownLeft");
|
||||
});
|
||||
|
||||
it("should use external display name when available", () => {
|
||||
const result = computeSelectedModel(
|
||||
true,
|
||||
[],
|
||||
[],
|
||||
"OpenAI Compatible",
|
||||
"Brain",
|
||||
);
|
||||
|
||||
expect(result!.name).toBe("OpenAI Compatible");
|
||||
expect(result!.icon).toBe("Brain");
|
||||
});
|
||||
|
||||
it("should return normal model when not in connection mode", () => {
|
||||
const options = [{ name: "gpt-4", icon: "Bot", provider: "OpenAI" }];
|
||||
const result = computeSelectedModel(false, [{ name: "gpt-4" }], options);
|
||||
|
||||
expect(result!.name).toBe("gpt-4");
|
||||
});
|
||||
|
||||
it("should return null when no model selected and not in connection mode", () => {
|
||||
const result = computeSelectedModel(false, null, []);
|
||||
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("connection mode clears template value (stale config prevention)", () => {
|
||||
it("should clear value to empty array when entering connection mode", () => {
|
||||
const previousValue = [
|
||||
{ name: "gpt-4", icon: "Bot", provider: "OpenAI" },
|
||||
];
|
||||
const clearedValue: any[] = [];
|
||||
|
||||
expect(clearedValue).toEqual([]);
|
||||
expect(clearedValue).not.toEqual(previousValue);
|
||||
});
|
||||
|
||||
it("should never send a string value to handleOnNewValue", () => {
|
||||
// The old bug: handleOnNewValue({ value: "connect_other_models" })
|
||||
// caused "string indices must be integers, not 'str'" on the backend.
|
||||
const newValue: any[] = [];
|
||||
expect(typeof newValue).not.toBe("string");
|
||||
expect(Array.isArray(newValue)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("credential clearing in connection mode", () => {
|
||||
// Pure logic: useModelConnectionLogic.setNode clears password fields
|
||||
const clearCredentialFields = (
|
||||
template: Record<string, any>,
|
||||
): Record<string, any> => {
|
||||
const updated = { ...template };
|
||||
if (updated.model) {
|
||||
updated.model = {
|
||||
...updated.model,
|
||||
value: [],
|
||||
_connection_mode: true,
|
||||
};
|
||||
}
|
||||
for (const [key, field] of Object.entries(updated)) {
|
||||
if (field?.password || field?._input_type === "SecretStrInput") {
|
||||
updated[key] = { ...field, value: "", load_from_db: false };
|
||||
}
|
||||
}
|
||||
return updated;
|
||||
};
|
||||
|
||||
it("should clear model value and set _connection_mode flag", () => {
|
||||
const template = {
|
||||
model: {
|
||||
type: "model",
|
||||
value: [{ name: "gpt-4", provider: "OpenAI" }],
|
||||
},
|
||||
api_key: {
|
||||
type: "str",
|
||||
password: true,
|
||||
_input_type: "SecretStrInput",
|
||||
load_from_db: true,
|
||||
value: "OPENAI_API_KEY",
|
||||
},
|
||||
};
|
||||
|
||||
const result = clearCredentialFields(template);
|
||||
|
||||
expect(result.model.value).toEqual([]);
|
||||
expect(result.model._connection_mode).toBe(true);
|
||||
});
|
||||
|
||||
it("should clear password fields value and disable load_from_db", () => {
|
||||
const template = {
|
||||
model: { type: "model", value: [{ name: "gpt-4" }] },
|
||||
api_key: {
|
||||
type: "str",
|
||||
password: true,
|
||||
_input_type: "SecretStrInput",
|
||||
load_from_db: true,
|
||||
value: "OPENAI_API_KEY",
|
||||
},
|
||||
temperature: {
|
||||
type: "float",
|
||||
value: 0.7,
|
||||
password: false,
|
||||
},
|
||||
};
|
||||
|
||||
const result = clearCredentialFields(template);
|
||||
|
||||
expect(result.api_key.value).toBe("");
|
||||
expect(result.api_key.load_from_db).toBe(false);
|
||||
// Non-secret fields should not be affected
|
||||
expect(result.temperature.value).toBe(0.7);
|
||||
});
|
||||
});
|
||||
|
||||
describe("exiting connection mode clears _connection_mode flag", () => {
|
||||
const exitConnectionMode = (
|
||||
template: Record<string, any>,
|
||||
): Record<string, any> => {
|
||||
const updated = { ...template };
|
||||
if (updated.model?._connection_mode) {
|
||||
updated.model = { ...updated.model, _connection_mode: false };
|
||||
}
|
||||
return updated;
|
||||
};
|
||||
|
||||
it("should set _connection_mode to false when selecting a model", () => {
|
||||
const template = {
|
||||
model: { type: "model", value: [], _connection_mode: true },
|
||||
};
|
||||
|
||||
const result = exitConnectionMode(template);
|
||||
|
||||
expect(result.model._connection_mode).toBe(false);
|
||||
});
|
||||
|
||||
it("should not change template when _connection_mode is already false", () => {
|
||||
const template = {
|
||||
model: {
|
||||
type: "model",
|
||||
value: [{ name: "gpt-4" }],
|
||||
_connection_mode: false,
|
||||
},
|
||||
};
|
||||
|
||||
const result = exitConnectionMode(template);
|
||||
|
||||
expect(result.model._connection_mode).toBe(false);
|
||||
expect(result.model.value).toEqual([{ name: "gpt-4" }]);
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -1,5 +1,4 @@
|
||||
import { useCallback } from "react";
|
||||
import { mutateTemplate } from "@/CustomNodes/helpers/mutate-template";
|
||||
import useFlowStore from "@/stores/flowStore";
|
||||
import { useTypesStore } from "@/stores/typesStore";
|
||||
import { scapedJSONStringfy } from "@/utils/reactflowUtils";
|
||||
@ -7,120 +6,111 @@ import { groupByFamily } from "@/utils/utils";
|
||||
|
||||
interface UseModelConnectionLogicProps {
|
||||
nodeId: string;
|
||||
nodeClass: any;
|
||||
handleNodeClass: (nodeClass: any) => void;
|
||||
postTemplateValue: any;
|
||||
setErrorData: (error: any) => void;
|
||||
handleOnNewValue: (newValue: any) => void;
|
||||
closePopover: () => void;
|
||||
clearSelection: () => void;
|
||||
}
|
||||
|
||||
export function useModelConnectionLogic({
|
||||
nodeId,
|
||||
nodeClass,
|
||||
handleNodeClass,
|
||||
postTemplateValue,
|
||||
setErrorData,
|
||||
handleOnNewValue,
|
||||
closePopover,
|
||||
clearSelection,
|
||||
}: UseModelConnectionLogicProps) {
|
||||
const handleExternalOptions = useCallback(
|
||||
async (optionValue: string) => {
|
||||
(optionValue: string) => {
|
||||
closePopover();
|
||||
clearSelection();
|
||||
|
||||
// Pass the optionValue ("connect_other_models") as both the field value and to mutateTemplate
|
||||
// This way the backend knows we're in connection mode
|
||||
handleOnNewValue({ value: optionValue });
|
||||
if (optionValue !== "connect_other_models") {
|
||||
return;
|
||||
}
|
||||
|
||||
await mutateTemplate(
|
||||
optionValue,
|
||||
nodeId!,
|
||||
nodeClass!,
|
||||
handleNodeClass!,
|
||||
postTemplateValue,
|
||||
setErrorData,
|
||||
"model",
|
||||
() => {
|
||||
// Enable connection mode for connect_other_models AFTER mutation completes
|
||||
try {
|
||||
if (optionValue === "connect_other_models") {
|
||||
const store = useFlowStore.getState();
|
||||
const node = store.getNode(nodeId!);
|
||||
const templateField = node?.data?.node?.template?.["model"];
|
||||
if (!templateField) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const store = useFlowStore.getState();
|
||||
const node = store.getNode(nodeId!);
|
||||
const templateField = node?.data?.node?.template?.["model"];
|
||||
if (!templateField) {
|
||||
return;
|
||||
}
|
||||
|
||||
const inputTypes: string[] =
|
||||
(Array.isArray(templateField.input_types)
|
||||
? templateField.input_types
|
||||
: []) || [];
|
||||
const effectiveInputTypes =
|
||||
inputTypes.length > 0 ? inputTypes : ["LanguageModel"];
|
||||
const inputTypes: string[] =
|
||||
(Array.isArray(templateField.input_types)
|
||||
? templateField.input_types
|
||||
: []) || [];
|
||||
const effectiveInputTypes =
|
||||
inputTypes.length > 0 ? inputTypes : ["LanguageModel"];
|
||||
|
||||
const tooltipTitle: string =
|
||||
(inputTypes && inputTypes.length > 0
|
||||
? inputTypes.join("\n")
|
||||
: templateField.type) || "";
|
||||
const tooltipTitle: string =
|
||||
(inputTypes && inputTypes.length > 0
|
||||
? inputTypes.join("\n")
|
||||
: templateField.type) || "";
|
||||
|
||||
const myId = scapedJSONStringfy({
|
||||
inputTypes: effectiveInputTypes,
|
||||
type: templateField.type,
|
||||
id: nodeId,
|
||||
fieldName: "model",
|
||||
proxy: templateField.proxy,
|
||||
});
|
||||
const typesData = useTypesStore.getState().data;
|
||||
const grouped = groupByFamily(
|
||||
typesData,
|
||||
(effectiveInputTypes && effectiveInputTypes.length > 0
|
||||
? effectiveInputTypes.join("\n")
|
||||
: tooltipTitle) || "",
|
||||
true,
|
||||
store.nodes,
|
||||
);
|
||||
|
||||
const typesData = useTypesStore.getState().data;
|
||||
const grouped = groupByFamily(
|
||||
typesData,
|
||||
(effectiveInputTypes && effectiveInputTypes.length > 0
|
||||
? effectiveInputTypes.join("\n")
|
||||
: tooltipTitle) || "",
|
||||
true,
|
||||
store.nodes,
|
||||
);
|
||||
// Build a pseudo source so compatible target handles (left side) glow
|
||||
const pseudoSourceHandle = scapedJSONStringfy({
|
||||
fieldName: "model",
|
||||
id: nodeId,
|
||||
inputTypes: effectiveInputTypes,
|
||||
type: "str",
|
||||
});
|
||||
|
||||
// Build a pseudo source so compatible target handles (left side) glow
|
||||
const pseudoSourceHandle = scapedJSONStringfy({
|
||||
fieldName: "model",
|
||||
id: nodeId,
|
||||
inputTypes: effectiveInputTypes,
|
||||
type: "str",
|
||||
});
|
||||
const filterObj = {
|
||||
source: undefined,
|
||||
sourceHandle: undefined,
|
||||
target: nodeId,
|
||||
targetHandle: pseudoSourceHandle,
|
||||
type: "LanguageModel",
|
||||
color: "datatype-fuchsia",
|
||||
} as any;
|
||||
|
||||
const filterObj = {
|
||||
source: undefined,
|
||||
sourceHandle: undefined,
|
||||
target: nodeId,
|
||||
targetHandle: pseudoSourceHandle,
|
||||
type: "LanguageModel",
|
||||
color: "datatype-fuchsia",
|
||||
} as any;
|
||||
|
||||
// Show compatible handles glow
|
||||
store.setFilterEdge(grouped);
|
||||
store.setFilterType(filterObj);
|
||||
// Mark this node as being in connection mode, clear the model value
|
||||
// and provider-specific credential fields so the backend cannot
|
||||
// execute with stale config when no external model is connected.
|
||||
store.setNode(
|
||||
nodeId,
|
||||
(prevNode) => {
|
||||
const template = { ...prevNode.data.node.template };
|
||||
if (template.model) {
|
||||
template.model = {
|
||||
...template.model,
|
||||
value: [],
|
||||
_connection_mode: true,
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn("Error setting up connection mode:", error);
|
||||
}
|
||||
},
|
||||
);
|
||||
for (const [key, field] of Object.entries(template)) {
|
||||
const f = field as any;
|
||||
if (f?.password || f?._input_type === "SecretStrInput") {
|
||||
template[key] = { ...f, value: "", load_from_db: false };
|
||||
}
|
||||
}
|
||||
return {
|
||||
...prevNode,
|
||||
data: {
|
||||
...prevNode.data,
|
||||
_connectionMode: true,
|
||||
node: { ...prevNode.data.node, template },
|
||||
},
|
||||
};
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
// Show compatible handles glow
|
||||
store.setFilterEdge(grouped);
|
||||
store.setFilterType(filterObj);
|
||||
} catch (error) {
|
||||
console.warn("Error setting up connection mode:", error);
|
||||
}
|
||||
},
|
||||
[
|
||||
nodeId,
|
||||
nodeClass,
|
||||
handleNodeClass,
|
||||
postTemplateValue,
|
||||
setErrorData,
|
||||
handleOnNewValue,
|
||||
closePopover,
|
||||
clearSelection,
|
||||
],
|
||||
[nodeId, closePopover, clearSelection],
|
||||
);
|
||||
|
||||
return { handleExternalOptions };
|
||||
|
||||
@ -7,7 +7,10 @@ import { usePostTemplateValue } from "@/controllers/API/queries/nodes/use-post-t
|
||||
import { useRefreshModelInputs } from "@/hooks/use-refresh-model-inputs";
|
||||
import ModelProviderModal from "@/modals/modelProviderModal";
|
||||
import useAlertStore from "@/stores/alertStore";
|
||||
import useFlowStore from "@/stores/flowStore";
|
||||
import type { APIClassType } from "@/types/api";
|
||||
import type { NodeDataType } from "@/types/flow";
|
||||
import { useModelConnectionLogic } from "./hooks/useModelConnectionLogic";
|
||||
import ForwardedIconComponent from "../../../../common/genericIconComponent";
|
||||
import { Button } from "../../../../ui/button";
|
||||
import { Command } from "../../../../ui/command";
|
||||
@ -48,6 +51,31 @@ export default function ModelInputComponent({
|
||||
useState(false);
|
||||
const [isRefreshingAfterClose, setIsRefreshingAfterClose] = useState(false);
|
||||
const [refreshOptions, setRefreshOptions] = useState(false);
|
||||
|
||||
// Connection mode: local state for reactivity, persisted in node data for reload
|
||||
const [isConnectionMode, setIsConnectionMode] = useState(() => {
|
||||
if (!nodeId) return false;
|
||||
const node = useFlowStore.getState().nodes.find((n) => n.id === nodeId);
|
||||
return (node?.data as any)?._connectionMode === true;
|
||||
});
|
||||
|
||||
const setConnectionMode = useCallback(
|
||||
(enabled: boolean) => {
|
||||
setIsConnectionMode(enabled);
|
||||
if (!nodeId) return;
|
||||
const store = useFlowStore.getState();
|
||||
store.setNode(
|
||||
nodeId,
|
||||
(node) => ({
|
||||
...node,
|
||||
data: { ...node.data, _connectionMode: enabled },
|
||||
}),
|
||||
false,
|
||||
);
|
||||
},
|
||||
[nodeId],
|
||||
);
|
||||
|
||||
const { refreshAllModelInputs } = useRefreshModelInputs();
|
||||
|
||||
// Ref to track if we've already processed the empty options state
|
||||
@ -60,6 +88,19 @@ export default function ModelInputComponent({
|
||||
node: (nodeClass as APIClassType) || null,
|
||||
});
|
||||
|
||||
const { handleExternalOptions } = useModelConnectionLogic({
|
||||
nodeId: nodeId || "",
|
||||
closePopover: () => setOpen(false),
|
||||
clearSelection: () => {
|
||||
// Only set the _connectionMode flag on the node data.
|
||||
// Don't call handleOnNewValue — it triggers a backend round-trip
|
||||
// that tries to resolve __default_language_model__ and fails.
|
||||
// The credential fields are cleared by useModelConnectionLogic
|
||||
// directly in the node template via setNode.
|
||||
setConnectionMode(true);
|
||||
},
|
||||
});
|
||||
|
||||
const modelType =
|
||||
nodeClass?.template?.model?.model_type === "language"
|
||||
? "llm"
|
||||
@ -111,9 +152,15 @@ export default function ModelInputComponent({
|
||||
|
||||
// Derive the currently selected model from the value prop
|
||||
const selectedModel = useMemo(() => {
|
||||
// If we're in connection mode, we don't have a normal selected model
|
||||
if (value === "connect_other_models") {
|
||||
return null;
|
||||
// If we're in connection mode, show the connection option as selected
|
||||
if (isConnectionMode) {
|
||||
return {
|
||||
name:
|
||||
externalOptions?.fields?.data?.node?.display_name ||
|
||||
"Connect other models",
|
||||
icon: externalOptions?.fields?.data?.node?.icon || "CornerDownLeft",
|
||||
provider: "",
|
||||
} as SelectedModel;
|
||||
}
|
||||
|
||||
const currentName = value?.[0]?.name;
|
||||
@ -133,11 +180,15 @@ export default function ModelInputComponent({
|
||||
// Or keep displaying the stale one? Original logic selected first available.
|
||||
(flatOptions.length > 0 ? flatOptions[0] : null)
|
||||
);
|
||||
}, [value, flatOptions]);
|
||||
}, [value, flatOptions, isConnectionMode, externalOptions]);
|
||||
|
||||
useEffect(() => {
|
||||
// Only proceed if we have options and haven't selected a value
|
||||
if (flatOptions.length > 0 && (!value || value.length === 0)) {
|
||||
// Only proceed if we have options, haven't selected a value, and NOT in connection mode
|
||||
if (
|
||||
flatOptions.length > 0 &&
|
||||
(!value || value.length === 0) &&
|
||||
!isConnectionMode
|
||||
) {
|
||||
// Check ref to avoid infinite loops
|
||||
if (!hasProcessedEmptyRef.current) {
|
||||
const firstOption = flatOptions[0];
|
||||
@ -162,6 +213,37 @@ export default function ModelInputComponent({
|
||||
*/
|
||||
const handleModelSelect = useCallback(
|
||||
(modelName: string) => {
|
||||
setConnectionMode(false);
|
||||
// Clear the _connection_mode flag from the model field template
|
||||
// so the backend resumes normal update_build_config behavior.
|
||||
if (nodeId) {
|
||||
const store = useFlowStore.getState();
|
||||
const node = store.getNode(nodeId);
|
||||
const nodeData = node?.data as NodeDataType | undefined;
|
||||
if (nodeData?.node?.template?.model?._connection_mode) {
|
||||
store.setNode(
|
||||
nodeId,
|
||||
(prev) => ({
|
||||
...prev,
|
||||
data: {
|
||||
...prev.data,
|
||||
_connectionMode: false,
|
||||
node: {
|
||||
...(prev.data as NodeDataType).node,
|
||||
template: {
|
||||
...(prev.data as NodeDataType).node.template,
|
||||
model: {
|
||||
...(prev.data as NodeDataType).node.template.model,
|
||||
_connection_mode: false,
|
||||
},
|
||||
},
|
||||
},
|
||||
} as NodeDataType,
|
||||
}),
|
||||
false,
|
||||
);
|
||||
}
|
||||
}
|
||||
const selectedOption = flatOptions.find(
|
||||
(option) => option.name === modelName,
|
||||
);
|
||||
@ -295,6 +377,17 @@ export default function ModelInputComponent({
|
||||
"refresh-model-list",
|
||||
)}
|
||||
{renderManageProvidersButton()}
|
||||
{externalOptions?.fields?.data?.node && (
|
||||
<div className="border-t bg-background">
|
||||
{renderFooterButton(
|
||||
externalOptions.fields.data.node.display_name ||
|
||||
"Connect other models",
|
||||
externalOptions.fields.data.node.icon || "CornerDownLeft",
|
||||
() => handleExternalOptions("connect_other_models"),
|
||||
"connect-other-models",
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</Command>
|
||||
</PopoverContentInput>
|
||||
);
|
||||
|
||||
@ -69,6 +69,8 @@ export type NodeDataType = {
|
||||
selected_output_type?: string;
|
||||
buildStatus?: BuildStatus;
|
||||
selected_output?: string;
|
||||
/** Transient flag: true while "Connect other models" mode is active */
|
||||
_connectionMode?: boolean;
|
||||
};
|
||||
|
||||
export type EdgeType = Edge<EdgeDataType, "default">;
|
||||
|
||||
@ -33,9 +33,13 @@ test("user must be able outdated message on error", async ({ page }) => {
|
||||
|
||||
await page.getByTestId("list-card").first().click();
|
||||
|
||||
await page
|
||||
.getByTestId("popover-anchor-input-api_key")
|
||||
.fill("this is a test to crash");
|
||||
// The api_key field may show a global variable badge instead of an input
|
||||
// when OPENAI_API_KEY is in the environment. Either way, we just need to
|
||||
// trigger a run to get the outdated components error message.
|
||||
const apiKeyInput = page.getByTestId("popover-anchor-input-api_key");
|
||||
if (await apiKeyInput.isVisible({ timeout: 3000 }).catch(() => false)) {
|
||||
await apiKeyInput.fill("this is a test to crash");
|
||||
}
|
||||
|
||||
await page.getByTestId("button_run_chat output").click();
|
||||
|
||||
|
||||
@ -313,6 +313,11 @@ def handle_model_input_update(
|
||||
if field_name in provider_mapped_fields:
|
||||
return build_config
|
||||
|
||||
# If the model field is in connection mode (user chose "Connect other models"),
|
||||
# skip auto-selection and provider re-population so credentials stay cleared.
|
||||
if build_config.get(model_field_name, {}).get("_connection_mode"):
|
||||
return build_config
|
||||
|
||||
# When the user changes the model selection, we need to reset/hide fields that may no longer apply
|
||||
if field_name == model_field_name:
|
||||
options = build_config[model_field_name].get("options", [])
|
||||
|
||||
@ -206,7 +206,22 @@ class ParameterHandler:
|
||||
params = self._handle_other_direct_types(field_name, field, val, params)
|
||||
|
||||
if field.get("load_from_db"):
|
||||
load_from_db_fields.append(field_name)
|
||||
# Skip load_from_db if the field itself has an incoming edge
|
||||
has_incoming_edge = self.vertex.get_incoming_edge_by_target_param(field_name) is not None
|
||||
# Skip credential fields when the model field has an incoming edge,
|
||||
# because the connected model component provides its own credentials
|
||||
is_secret = field.get("_input_type") == "SecretStrInput" or field.get("password")
|
||||
model_has_edge = (
|
||||
is_secret
|
||||
and "model" in self.template_dict
|
||||
and self.vertex.get_incoming_edge_by_target_param("model") is not None
|
||||
)
|
||||
# Skip credential fields when the node is in "Connect other models" mode
|
||||
# (user chose to wire an external model instead of the built-in provider)
|
||||
model_field = self.template_dict.get("model", {})
|
||||
in_connection_mode = is_secret and model_field.get("_connection_mode", False)
|
||||
if not has_incoming_edge and not model_has_edge and not in_connection_mode:
|
||||
load_from_db_fields.append(field_name)
|
||||
|
||||
return params, load_from_db_fields
|
||||
|
||||
|
||||
254
src/lfx/tests/unit/graph/vertex/test_param_handler_model_edge.py
Normal file
254
src/lfx/tests/unit/graph/vertex/test_param_handler_model_edge.py
Normal file
@ -0,0 +1,254 @@
|
||||
"""Tests for skipping load_from_db credential fields when model has incoming edge.
|
||||
|
||||
When a model component is connected via wire to a node's model input,
|
||||
the connected component provides its own credentials. The node's own
|
||||
secret fields (like api_key) should NOT be resolved from the database,
|
||||
preventing errors like 'OPENAI_API_KEY variable not found'.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from lfx.graph.vertex.param_handler import ParameterHandler
|
||||
|
||||
|
||||
class TestSkipLoadFromDbWhenModelHasEdge:
|
||||
"""Tests for _process_direct_type_field skipping load_from_db for secrets when model has edge."""
|
||||
|
||||
def _create_handler_with_template(self, template: dict, *, model_has_edge: bool = False) -> ParameterHandler:
|
||||
"""Create a ParameterHandler with the given template and optional model edge."""
|
||||
mock_vertex = MagicMock()
|
||||
mock_vertex.data = {"node": {"template": template}}
|
||||
|
||||
def mock_get_incoming_edge(field_name: str):
|
||||
if field_name == "model" and model_has_edge:
|
||||
return "source-node-id"
|
||||
return None
|
||||
|
||||
mock_vertex.get_incoming_edge_by_target_param = MagicMock(side_effect=mock_get_incoming_edge)
|
||||
return ParameterHandler(mock_vertex, storage_service=None)
|
||||
|
||||
def test_should_add_secret_to_load_from_db_when_no_model_edge(self):
|
||||
"""Secret fields should be resolved normally when model has no wire."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {"type": "model", "value": None, "show": True},
|
||||
"api_key": {
|
||||
"type": "str",
|
||||
"_input_type": "SecretStrInput",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "OPENAI_API_KEY",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=False)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "api_key" in load_from_db_fields
|
||||
|
||||
def test_should_skip_secret_load_from_db_when_model_has_edge(self):
|
||||
"""Secret fields should NOT be resolved when model has a connected wire."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {"type": "model", "value": None, "show": True},
|
||||
"api_key": {
|
||||
"type": "str",
|
||||
"_input_type": "SecretStrInput",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "OPENAI_API_KEY",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=True)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "api_key" not in load_from_db_fields
|
||||
|
||||
def test_should_skip_password_field_when_model_has_edge(self):
|
||||
"""Any password field should be skipped when model has edge, not just SecretStrInput."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {"type": "model", "value": None, "show": True},
|
||||
"custom_token": {
|
||||
"type": "str",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "MY_CUSTOM_TOKEN",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=True)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("custom_token", template["custom_token"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "custom_token" not in load_from_db_fields
|
||||
|
||||
def test_should_not_skip_non_secret_field_when_model_has_edge(self):
|
||||
"""Non-secret fields with load_from_db should still be resolved even with model edge."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {"type": "model", "value": None, "show": True},
|
||||
"system_prompt": {
|
||||
"type": "str",
|
||||
"load_from_db": True,
|
||||
"value": "MY_PROMPT_VAR",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=True)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("system_prompt", template["system_prompt"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "system_prompt" in load_from_db_fields
|
||||
|
||||
def test_should_skip_field_with_direct_incoming_edge(self):
|
||||
"""Fields with their own incoming edge should be skipped regardless of model edge."""
|
||||
# Arrange
|
||||
template = {
|
||||
"some_field": {
|
||||
"type": "str",
|
||||
"load_from_db": True,
|
||||
"value": "SOME_VAR",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
mock_vertex = MagicMock()
|
||||
mock_vertex.data = {"node": {"template": template}}
|
||||
mock_vertex.get_incoming_edge_by_target_param = MagicMock(return_value="source-node")
|
||||
handler = ParameterHandler(mock_vertex, storage_service=None)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("some_field", template["some_field"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "some_field" not in load_from_db_fields
|
||||
|
||||
def test_should_not_skip_when_no_model_field_in_template(self):
|
||||
"""Components without a model field should resolve secrets normally."""
|
||||
# Arrange
|
||||
template = {
|
||||
"api_key": {
|
||||
"type": "str",
|
||||
"_input_type": "SecretStrInput",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "OPENAI_API_KEY",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
mock_vertex = MagicMock()
|
||||
mock_vertex.data = {"node": {"template": template}}
|
||||
mock_vertex.get_incoming_edge_by_target_param = MagicMock(return_value=None)
|
||||
handler = ParameterHandler(mock_vertex, storage_service=None)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "api_key" in load_from_db_fields
|
||||
|
||||
def test_should_not_affect_fields_without_load_from_db(self):
|
||||
"""Fields without load_from_db should never be added regardless of edges."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {"type": "model", "value": None, "show": True},
|
||||
"description": {
|
||||
"type": "str",
|
||||
"value": "Some description",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=True)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("description", template["description"], params, load_from_db_fields)
|
||||
|
||||
# Assert
|
||||
assert "description" not in load_from_db_fields
|
||||
|
||||
def test_should_skip_secret_load_from_db_when_connection_mode_active(self):
|
||||
"""Secret fields must NOT be resolved via load_from_db in connection mode.
|
||||
|
||||
When the model field has _connection_mode: true (user chose
|
||||
'Connect other models'), credential fields should be skipped.
|
||||
"""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {
|
||||
"type": "model",
|
||||
"value": [],
|
||||
"show": True,
|
||||
"_connection_mode": True,
|
||||
},
|
||||
"api_key": {
|
||||
"type": "str",
|
||||
"_input_type": "SecretStrInput",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=False)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields)
|
||||
|
||||
# Assert — api_key must NOT be in load_from_db_fields
|
||||
assert "api_key" not in load_from_db_fields
|
||||
|
||||
def test_should_allow_secret_load_from_db_when_connection_mode_false(self):
|
||||
"""Secret fields should be resolved normally when _connection_mode is false."""
|
||||
# Arrange
|
||||
template = {
|
||||
"model": {
|
||||
"type": "model",
|
||||
"value": [{"name": "gpt-4", "provider": "OpenAI"}],
|
||||
"show": True,
|
||||
"_connection_mode": False,
|
||||
},
|
||||
"api_key": {
|
||||
"type": "str",
|
||||
"_input_type": "SecretStrInput",
|
||||
"password": True,
|
||||
"load_from_db": True,
|
||||
"value": "OPENAI_API_KEY",
|
||||
"show": True,
|
||||
},
|
||||
}
|
||||
handler = self._create_handler_with_template(template, model_has_edge=False)
|
||||
|
||||
# Act
|
||||
load_from_db_fields: list[str] = []
|
||||
params: dict = {}
|
||||
handler._process_direct_type_field("api_key", template["api_key"], params, load_from_db_fields)
|
||||
|
||||
# Assert — api_key SHOULD be in load_from_db_fields
|
||||
assert "api_key" in load_from_db_fields
|
||||
@ -49,6 +49,8 @@ def mock_vertex() -> Mock:
|
||||
}
|
||||
vertex.id = "test-vertex-id"
|
||||
vertex.display_name = "Test Vertex"
|
||||
# Default: no incoming edges for any field
|
||||
vertex.get_incoming_edge_by_target_param = Mock(return_value=None)
|
||||
return vertex
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user