From ca514e579f49d830b34beae43fe288ba9ee188ed Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 29 Feb 2024 23:07:10 -0300
Subject: [PATCH] Add branching to GenericNode

---
 .../components/parameterComponent/index.tsx   |   8 ++
 .../src/CustomNodes/GenericNode/index.tsx     | 110 +++++++++++++-----
 src/frontend/src/types/api/index.ts           |   1 +
 src/frontend/src/types/components/index.ts    |   1 +
 src/frontend/src/types/flow/index.ts          |   1 +
 5 files changed, 90 insertions(+), 31 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index ef8ad8d80..272aad5ba 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -15,6 +15,7 @@ import KeypairListComponent from "../../../../components/keypairListComponent";
 import PromptAreaComponent from "../../../../components/promptComponent";
 import TextAreaComponent from "../../../../components/textAreaComponent";
 import ToggleShadComponent from "../../../../components/toggleShadComponent";
+import { Badge } from "../../../../components/ui/badge";
 import { Button } from "../../../../components/ui/button";
 import {
   LANGFLOW_SUPPORTED_TYPES,
@@ -50,6 +51,7 @@ export default function ParameterComponent({
   data,
   tooltipTitle,
   title,
+  conditionPath,
   color,
   type,
   name = "",
@@ -315,6 +317,12 @@ export default function ParameterComponent({
             (info !== "" ? " flex items-center" : "")
           }
         >
+          {" "}
+          {conditionPath && (
+            <Badge variant="gray" size="md" className="mr-2">
+              {conditionPath}
+            </Badge>
+          )}
           {proxy ? (
             <ShadTooltip content={<span>{proxy.id}</span>}>
               <span>{title}</span>
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index c9844a47d..3729ac418 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -664,37 +664,85 @@ export default function GenericNode({
               >
                 {" "}
               </div>
-              {data.node!.base_classes.length > 0 && (
-                <ParameterComponent
-                  key={scapedJSONStringfy({
-                    baseClasses: data.node!.base_classes,
-                    id: data.id,
-                    dataType: data.type,
-                  })}
-                  data={data}
-                  color={
-                    (data.node?.output_types &&
-                    data.node.output_types.length > 0
-                      ? nodeColors[data.node.output_types[0]] ??
-                        nodeColors[types[data.node.output_types[0]]]
-                      : nodeColors[types[data.type]]) ?? nodeColors.unknown
-                  }
-                  title={
-                    data.node?.output_types && data.node.output_types.length > 0
-                      ? data.node.output_types.join(" | ")
-                      : data.type
-                  }
-                  tooltipTitle={data.node?.base_classes.join("\n")}
-                  id={{
-                    baseClasses: data.node!.base_classes,
-                    id: data.id,
-                    dataType: data.type,
-                  }}
-                  type={data.node?.base_classes.join("|")}
-                  left={false}
-                  showNode={showNode}
-                />
-              )}
+              <div>
+                {data.node!.base_classes.length > 0 && (
+                  <ParameterComponent
+                    key={scapedJSONStringfy({
+                      baseClasses: data.node!.base_classes,
+                      id: data.id,
+                      dataType: data.type,
+                      conditionalPath: data.node!.is_conditional ? true : null,
+                    })}
+                    data={data}
+                    color={
+                      (data.node?.output_types &&
+                      data.node.output_types.length > 0
+                        ? nodeColors[data.node.output_types[0]] ??
+                          nodeColors[types[data.node.output_types[0]]]
+                        : nodeColors[types[data.type]]) ?? nodeColors.unknown
+                    }
+                    title={
+                      data.node?.output_types &&
+                      data.node.output_types.length > 0
+                        ? data.node.output_types.join(" | ")
+                        : data.type
+                    }
+                    conditionPath={data.node?.is_conditional ? "True" : null}
+                    tooltipTitle={data.node?.base_classes.join("\n")}
+                    id={{
+                      baseClasses: data.node!.base_classes,
+                      id: data.id,
+                      dataType: data.type,
+                      // First parameter component should be true
+                      // Second should be false
+                      conditionalPath: data.node!.is_conditional ? true : null,
+                    }}
+                    // Type should be base_classes if it's not a conditional node
+                    // else it should be true in the first parameter component
+                    type={data.node?.base_classes.join("|")}
+                    left={false}
+                    showNode={showNode}
+                  />
+                )}
+                {data.node!.is_conditional && (
+                  <ParameterComponent
+                    key={scapedJSONStringfy({
+                      baseClasses: data.node!.base_classes,
+                      id: data.id,
+                      dataType: data.type,
+                      conditionalPath: data.node!.is_conditional ? true : null,
+                    })}
+                    data={data}
+                    color={
+                      (data.node?.output_types &&
+                      data.node.output_types.length > 0
+                        ? nodeColors[data.node.output_types[0]] ??
+                          nodeColors[types[data.node.output_types[0]]]
+                        : nodeColors[types[data.type]]) ?? nodeColors.unknown
+                    }
+                    title={
+                      data.node?.output_types &&
+                      data.node.output_types.length > 0
+                        ? data.node.output_types.join(" | ")
+                        : data.type
+                    }
+                    conditionPath={data.node?.is_conditional ? "False" : null}
+                    tooltipTitle={data.node?.base_classes.join("\n")}
+                    id={{
+                      baseClasses: data.node!.base_classes,
+                      id: data.id,
+                      dataType: data.type,
+                      // condition should be null if it's not a conditional node
+                      // false if it's a conditional node and the condition is false
+                      // so we should check if it's a conditional node.
+                      conditionalPath: data.node!.is_conditional ? false : null,
+                    }}
+                    type={data.node?.base_classes.join("|")}
+                    left={false}
+                    showNode={showNode}
+                  />
+                )}
+              </div>
             </>
           </div>
         )}
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 6663692dc..646941bd9 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -20,6 +20,7 @@ export type APIClassType = {
   icon?: string;
   is_input?: boolean;
   is_output?: boolean;
+  is_conditional?: boolean;
   input_types?: Array<string>;
   output_types?: Array<string>;
   custom_fields?: CustomFieldsType;
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 4f5d1bb60..7236998bd 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -40,6 +40,7 @@ export type DropDownComponentType = {
 export type ParameterComponentType = {
   data: NodeDataType;
   title: string;
+  conditionPath?: string | null;
   id: sourceHandleType | targetHandleType;
   color: string;
   left: boolean;
diff --git a/src/frontend/src/types/flow/index.ts b/src/frontend/src/types/flow/index.ts
index 967d4e424..02861cf35 100644
--- a/src/frontend/src/types/flow/index.ts
+++ b/src/frontend/src/types/flow/index.ts
@@ -52,6 +52,7 @@ export type sourceHandleType = {
   dataType: string;
   id: string;
   baseClasses: string[];
+  conditionalPath?: boolean | null;
 };
 //left side
 export type targetHandleType = {