From d19fdf53736146ef928744f382d1d47df3898e09 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <lucas.edu.oli@hotmail.com>
Date: Mon, 4 Dec 2023 16:35:06 -0300
Subject: [PATCH] Added colors to types

---
 src/frontend/src/CustomNodes/GenericNode/index.tsx | 12 +++++++++---
 src/frontend/src/utils/styleUtils.ts               |  2 ++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index e212b5329..c2a340675 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -464,10 +464,10 @@ export default function GenericNode({
                         data={data}
                         color={
                           nodeColors[
-                            types[data.node?.template[templateField].type!]
+                            data.node?.template[templateField].type!
                           ] ??
                           nodeColors[
-                            data.node?.template[templateField].type!
+                            types[data.node?.template[templateField].type!]
                           ] ??
                           nodeColors.unknown
                         }
@@ -519,7 +519,13 @@ export default function GenericNode({
                     dataType: data.type,
                   })}
                   data={data}
-                  color={nodeColors[types[data.type]] ?? nodeColors.unknown}
+                  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("|")
diff --git a/src/frontend/src/utils/styleUtils.ts b/src/frontend/src/utils/styleUtils.ts
index 795d32988..fae4be1f0 100644
--- a/src/frontend/src/utils/styleUtils.ts
+++ b/src/frontend/src/utils/styleUtils.ts
@@ -174,6 +174,8 @@ export const nodeColors: { [char: string]: string } = {
   prompts: "#4367BF",
   llms: "#6344BE",
   chains: "#FE7500",
+  Document: "#7AAE42",
+  list: "#9AAE42",
   agents: "#903BBE",
   tools: "#FF3434",
   memories: "#F5B85A",