From 27e026d9436ddf58b51726d272a8db674f23d1af Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Thu, 20 Apr 2023 00:46:33 +0100
Subject: [PATCH] add baseClasses for outputs

---
 .../components/nodes/chains/LLMChain/LLMChain.ts    |  8 ++++----
 packages/components/src/Interface.ts                |  2 +-
 packages/ui/src/utils/genericHelper.js              | 13 +++++++------
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/packages/components/nodes/chains/LLMChain/LLMChain.ts b/packages/components/nodes/chains/LLMChain/LLMChain.ts
index ae441eae7..1e45cc44b 100644
--- a/packages/components/nodes/chains/LLMChain/LLMChain.ts
+++ b/packages/components/nodes/chains/LLMChain/LLMChain.ts
@@ -43,14 +43,14 @@ class LLMChain_Chains implements INode {
         ]
         this.outputs = [
             {
-                label: this.label,
-                name: this.name,
-                type: this.type
+                label: 'LLM Chain',
+                name: 'llmChain',
+                baseClasses: [this.type, ...getBaseClasses(LLMChain)]
             },
             {
                 label: 'Output Prediction',
                 name: 'outputPrediction',
-                type: 'string'
+                baseClasses: ['string']
             }
         ]
     }
diff --git a/packages/components/src/Interface.ts b/packages/components/src/Interface.ts
index fd26020c9..44c87d746 100644
--- a/packages/components/src/Interface.ts
+++ b/packages/components/src/Interface.ts
@@ -32,7 +32,7 @@ export interface INodeOptionsValue {
 export interface INodeOutputsValue {
     label: string
     name: string
-    type: string
+    baseClasses: string[]
     description?: string
 }
 
diff --git a/packages/ui/src/utils/genericHelper.js b/packages/ui/src/utils/genericHelper.js
index 845f934d0..dd63c1313 100644
--- a/packages/ui/src/utils/genericHelper.js
+++ b/packages/ui/src/utils/genericHelper.js
@@ -61,12 +61,13 @@ export const initNode = (nodeData, newNodeId) => {
                 let baseClasses = ''
                 let type = ''
 
-                if (whitelistTypes.includes(nodeData.outputs[j].type)) {
-                    baseClasses = nodeData.outputs[j].type
-                    type = nodeData.outputs[j].type
-                } else {
-                    baseClasses = nodeData.baseClasses.join('|')
-                    type = nodeData.baseClasses.join(' | ')
+                const outputBaseClasses = nodeData.outputs[j].baseClasses ?? []
+                if (outputBaseClasses.length > 1) {
+                    baseClasses = outputBaseClasses.join('|')
+                    type = outputBaseClasses.join(' | ')
+                } else if (outputBaseClasses.length === 1) {
+                    baseClasses = outputBaseClasses[0]
+                    type = outputBaseClasses[0]
                 }
 
                 const newOutputOption = {