New Feature: Add post postprocessing of response from LLM using custom JS Function (#4079)

* New Feature: Add post postprocessing of response from LLM using custom Javascript functions

* Disable Save when there is no content

* add post processing ui changes, disable streaming

---------

Co-authored-by: Henry <hzj94@hotmail.com>

packages/components/nodes/utilities/CustomFunction/CustomFunction.ts

@@ -88,6 +88,7 @@ class CustomFunction_Utilities implements INode {
             chatflowId: options.chatflowid,
             sessionId: options.sessionId,
             chatId: options.chatId,
+            rawOutput: options.rawOutput || '',
             input
         }
 

packages/server/src/services/chatflows/index.ts

@@ -1,4 +1,4 @@
-import { removeFolderFromStorage } from 'flowise-components'
+import { ICommonObject, removeFolderFromStorage } from 'flowise-components'
 import { StatusCodes } from 'http-status-codes'
 import { ChatflowType, IReactFlowObject } from '../../Interface'
 import { ChatFlow } from '../../database/entities/ChatFlow'
@@ -28,6 +28,15 @@ const checkIfChatflowIsValidForStreaming = async (chatflowId: string): Promise<a
             throw new InternalFlowiseError(StatusCodes.NOT_FOUND, `Chatflow ${chatflowId} not found`)
         }
 
+        /* Check for post-processing settings, if available isStreamValid is always false */
+        let chatflowConfig: ICommonObject = {}
+        if (chatflow.chatbotConfig) {
+            chatflowConfig = JSON.parse(chatflow.chatbotConfig)
+            if (chatflowConfig?.postProcessing?.enabled === true) {
+                return { isStreaming: false }
+            }
+        }
+
         /*** Get Ending Node with Directed Graph  ***/
         const flowData = chatflow.flowData
         const parsedFlowData: IReactFlowObject = JSON.parse(flowData)

packages/server/src/utils/buildChatflow.ts

@@ -579,7 +579,19 @@ export const executeFlow = async ({
         }
         return undefined
     } else {
-        const isStreamValid = await checkIfStreamValid(endingNodes, nodes, streaming)
+        let chatflowConfig: ICommonObject = {}
+        if (chatflow.chatbotConfig) {
+            chatflowConfig = JSON.parse(chatflow.chatbotConfig)
+        }
+
+        let isStreamValid = false
+
+        /* Check for post-processing settings, if available isStreamValid is always false */
+        if (chatflowConfig?.postProcessing?.enabled === true) {
+            isStreamValid = false
+        } else {
+            isStreamValid = await checkIfStreamValid(endingNodes, nodes, streaming)
+        }
 
         /*** Find the last node to execute ***/
         const { endingNodeData, endingNodeInstance } = await initEndingNode({
@@ -637,8 +649,37 @@ export const executeFlow = async ({
         await utilAddChatMessage(userMessage, appDataSource)
 
         let resultText = ''
-        if (result.text) resultText = result.text
+        if (result.text) {
-        else if (result.json) resultText = '```json\n' + JSON.stringify(result.json, null, 2)
+            resultText = result.text
+            /* Check for post-processing settings */
+            if (chatflowConfig?.postProcessing?.enabled === true) {
+                try {
+                    const postProcessingFunction = JSON.parse(chatflowConfig?.postProcessing?.customFunction)
+                    const nodeInstanceFilePath = componentNodes['customFunction'].filePath as string
+                    const nodeModule = await import(nodeInstanceFilePath)
+                    const nodeData = {
+                        inputs: { javascriptFunction: postProcessingFunction },
+                        outputs: { output: 'output' }
+                    }
+                    const options: ICommonObject = {
+                        chatflowid: chatflow.id,
+                        sessionId,
+                        chatId,
+                        input: question,
+                        rawOutput: resultText,
+                        appDataSource,
+                        databaseEntities,
+                        logger
+                    }
+                    const customFuncNodeInstance = new nodeModule.nodeClass()
+                    let moderatedResponse = await customFuncNodeInstance.init(nodeData, question, options)
+                    result.text = moderatedResponse
+                    resultText = result.text
+                } catch (e) {
+                    logger.log('[server]: Post Processing Error:', e)
+                }
+            }
+        } else if (result.json) resultText = '```json\n' + JSON.stringify(result.json, null, 2)
         else resultText = JSON.stringify(result, null, 2)
 
         const apiMessage: Omit<IChatMessage, 'createdDate'> = {

packages/ui/src/ui-component/dialog/ChatflowConfigurationDialog.jsx

@@ -11,6 +11,7 @@ import StarterPrompts from '@/ui-component/extended/StarterPrompts'
 import Leads from '@/ui-component/extended/Leads'
 import FollowUpPrompts from '@/ui-component/extended/FollowUpPrompts'
 import FileUpload from '@/ui-component/extended/FileUpload'
+import PostProcessing from '@/ui-component/extended/PostProcessing'
 
 const CHATFLOW_CONFIGURATION_TABS = [
     {
@@ -44,6 +45,11 @@ const CHATFLOW_CONFIGURATION_TABS = [
     {
         label: 'File Upload',
         id: 'fileUpload'
+    },
+    {
+        label: 'Post Processing',
+        id: 'postProcessing',
+        hideInAgentFlow: true
     }
 ]
 
@@ -76,10 +82,12 @@ function a11yProps(index) {
     }
 }
 
-const ChatflowConfigurationDialog = ({ show, dialogProps, onCancel }) => {
+const ChatflowConfigurationDialog = ({ show, isAgentCanvas, dialogProps, onCancel }) => {
     const portalElement = document.getElementById('portal')
     const [tabValue, setTabValue] = useState(0)
 
+    const filteredTabs = CHATFLOW_CONFIGURATION_TABS.filter((tab) => !isAgentCanvas || !tab.hideInAgentFlow)
+
     const component = show ? (
         <Dialog
             onClose={onCancel}
@@ -108,7 +116,7 @@ const ChatflowConfigurationDialog = ({ show, dialogProps, onCancel }) => {
                     variant='scrollable'
                     scrollButtons='auto'
                 >
-                    {CHATFLOW_CONFIGURATION_TABS.map((item, index) => (
+                    {filteredTabs.map((item, index) => (
                         <Tab
                             sx={{
                                 minHeight: '40px',
@@ -123,7 +131,7 @@ const ChatflowConfigurationDialog = ({ show, dialogProps, onCancel }) => {
                         ></Tab>
                     ))}
                 </Tabs>
-                {CHATFLOW_CONFIGURATION_TABS.map((item, index) => (
+                {filteredTabs.map((item, index) => (
                     <TabPanel key={index} value={tabValue} index={index}>
                         {item.id === 'security' && <Security dialogProps={dialogProps} />}
                         {item.id === 'conversationStarters' ? <StarterPrompts dialogProps={dialogProps} /> : null}
@@ -133,6 +141,7 @@ const ChatflowConfigurationDialog = ({ show, dialogProps, onCancel }) => {
                         {item.id === 'analyseChatflow' ? <AnalyseFlow dialogProps={dialogProps} /> : null}
                         {item.id === 'leads' ? <Leads dialogProps={dialogProps} /> : null}
                         {item.id === 'fileUpload' ? <FileUpload dialogProps={dialogProps} /> : null}
+                        {item.id === 'postProcessing' ? <PostProcessing dialogProps={dialogProps} /> : null}
                     </TabPanel>
                 ))}
             </DialogContent>
@@ -144,6 +153,7 @@ const ChatflowConfigurationDialog = ({ show, dialogProps, onCancel }) => {
 
 ChatflowConfigurationDialog.propTypes = {
     show: PropTypes.bool,
+    isAgentCanvas: PropTypes.bool,
     dialogProps: PropTypes.object,
     onCancel: PropTypes.func
 }

packages/ui/src/ui-component/extended/PostProcessing.jsx

@@ -0,0 +1,228 @@
+import { useDispatch } from 'react-redux'
+import { useState, useEffect } from 'react'
+import PropTypes from 'prop-types'
+import { useSelector } from 'react-redux'
+
+// material-ui
+import { IconButton, Button, Box, Typography } from '@mui/material'
+import { IconArrowsMaximize, IconBulb, IconX } from '@tabler/icons-react'
+import { useTheme } from '@mui/material/styles'
+
+// Project import
+import { StyledButton } from '@/ui-component/button/StyledButton'
+import { SwitchInput } from '@/ui-component/switch/Switch'
+import { CodeEditor } from '@/ui-component/editor/CodeEditor'
+import ExpandTextDialog from '@/ui-component/dialog/ExpandTextDialog'
+
+// store
+import { enqueueSnackbar as enqueueSnackbarAction, closeSnackbar as closeSnackbarAction, SET_CHATFLOW } from '@/store/actions'
+import useNotifier from '@/utils/useNotifier'
+
+// API
+import chatflowsApi from '@/api/chatflows'
+
+const sampleFunction = `return $flow.rawOutput + " This is a post processed response!";`
+
+const PostProcessing = ({ dialogProps }) => {
+    const dispatch = useDispatch()
+
+    useNotifier()
+    const theme = useTheme()
+    const customization = useSelector((state) => state.customization)
+
+    const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
+    const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args))
+
+    const [postProcessingEnabled, setPostProcessingEnabled] = useState(false)
+    const [postProcessingFunction, setPostProcessingFunction] = useState('')
+    const [chatbotConfig, setChatbotConfig] = useState({})
+    const [showExpandDialog, setShowExpandDialog] = useState(false)
+    const [expandDialogProps, setExpandDialogProps] = useState({})
+
+    const handleChange = (value) => {
+        setPostProcessingEnabled(value)
+    }
+
+    const onExpandDialogClicked = (value) => {
+        const dialogProps = {
+            value,
+            inputParam: {
+                label: 'Post Processing Function',
+                name: 'postProcessingFunction',
+                type: 'code',
+                placeholder: sampleFunction,
+                hideCodeExecute: true
+            },
+            languageType: 'js',
+            confirmButtonName: 'Save',
+            cancelButtonName: 'Cancel'
+        }
+        setExpandDialogProps(dialogProps)
+        setShowExpandDialog(true)
+    }
+
+    const onSave = async () => {
+        try {
+            let value = {
+                postProcessing: {
+                    enabled: postProcessingEnabled,
+                    customFunction: JSON.stringify(postProcessingFunction)
+                }
+            }
+            chatbotConfig.postProcessing = value.postProcessing
+            const saveResp = await chatflowsApi.updateChatflow(dialogProps.chatflow.id, {
+                chatbotConfig: JSON.stringify(chatbotConfig)
+            })
+            if (saveResp.data) {
+                enqueueSnackbar({
+                    message: 'Post Processing Settings Saved',
+                    options: {
+                        key: new Date().getTime() + Math.random(),
+                        variant: 'success',
+                        action: (key) => (
+                            <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                                <IconX />
+                            </Button>
+                        )
+                    }
+                })
+                dispatch({ type: SET_CHATFLOW, chatflow: saveResp.data })
+            }
+        } catch (error) {
+            enqueueSnackbar({
+                message: `Failed to save Post Processing Settings: ${
+                    typeof error.response.data === 'object' ? error.response.data.message : error.response.data
+                }`,
+                options: {
+                    key: new Date().getTime() + Math.random(),
+                    variant: 'error',
+                    persist: true,
+                    action: (key) => (
+                        <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                            <IconX />
+                        </Button>
+                    )
+                }
+            })
+        }
+    }
+
+    useEffect(() => {
+        if (dialogProps.chatflow && dialogProps.chatflow.chatbotConfig) {
+            let chatbotConfig = JSON.parse(dialogProps.chatflow.chatbotConfig)
+            setChatbotConfig(chatbotConfig || {})
+            if (chatbotConfig.postProcessing) {
+                setPostProcessingEnabled(chatbotConfig.postProcessing.enabled)
+                if (chatbotConfig.postProcessing.customFunction) {
+                    setPostProcessingFunction(JSON.parse(chatbotConfig.postProcessing.customFunction))
+                }
+            }
+        }
+
+        return () => {}
+    }, [dialogProps])
+
+    return (
+        <>
+            <Box sx={{ width: '100%', display: 'flex', alignItems: 'center', justifyContent: 'space-between', mb: 2 }}>
+                <SwitchInput label='Enable Post Processing' onChange={handleChange} value={postProcessingEnabled} />
+            </Box>
+            <Box sx={{ width: '100%', display: 'flex', flexDirection: 'column', alignItems: 'flex-start', gap: 1 }}>
+                <Box sx={{ width: '100%', display: 'flex', alignItems: 'center' }}>
+                    <Typography>JS Function</Typography>
+                    <Button
+                        sx={{ ml: 2 }}
+                        variant='outlined'
+                        onClick={() => {
+                            setPostProcessingFunction(sampleFunction)
+                        }}
+                    >
+                        See Example
+                    </Button>
+                    <div style={{ flex: 1 }} />
+                    <IconButton
+                        size='small'
+                        sx={{
+                            height: 25,
+                            width: 25
+                        }}
+                        title='Expand'
+                        color='primary'
+                        onClick={() => onExpandDialogClicked(postProcessingFunction)}
+                    >
+                        <IconArrowsMaximize />
+                    </IconButton>
+                </Box>
+
+                <div
+                    style={{
+                        marginTop: '10px',
+                        border: '1px solid',
+                        borderColor: theme.palette.grey['300'],
+                        borderRadius: '6px',
+                        height: '200px',
+                        width: '100%'
+                    }}
+                >
+                    <CodeEditor
+                        value={postProcessingFunction}
+                        height='200px'
+                        theme={customization.isDarkMode ? 'dark' : 'light'}
+                        lang={'js'}
+                        placeholder={sampleFunction}
+                        onValueChange={(code) => setPostProcessingFunction(code)}
+                        basicSetup={{ highlightActiveLine: false, highlightActiveLineGutter: false }}
+                    />
+                </div>
+            </Box>
+            <div
+                style={{
+                    display: 'flex',
+                    flexDirection: 'column',
+                    borderRadius: 10,
+                    background: '#d8f3dc',
+                    padding: 10,
+                    marginTop: 10
+                }}
+            >
+                <div
+                    style={{
+                        display: 'flex',
+                        flexDirection: 'row',
+                        alignItems: 'center',
+                        paddingTop: 10
+                    }}
+                >
+                    <IconBulb size={30} color='#2d6a4f' />
+                    <span style={{ color: '#2d6a4f', marginLeft: 10, fontWeight: 500 }}>
+                        The following variables are available to use in the custom function:{' '}
+                        <pre>$flow.rawOutput, $flow.input, $flow.chatflowId, $flow.sessionId, $flow.chatId</pre>
+                    </span>
+                </div>
+            </div>
+            <StyledButton
+                style={{ marginBottom: 10, marginTop: 10 }}
+                variant='contained'
+                disabled={!postProcessingFunction || postProcessingFunction?.trim().length === 0}
+                onClick={onSave}
+            >
+                Save
+            </StyledButton>
+            <ExpandTextDialog
+                show={showExpandDialog}
+                dialogProps={expandDialogProps}
+                onCancel={() => setShowExpandDialog(false)}
+                onConfirm={(newValue) => {
+                    setPostProcessingFunction(newValue)
+                    setShowExpandDialog(false)
+                }}
+            ></ExpandTextDialog>
+        </>
+    )
+}
+
+PostProcessing.propTypes = {
+    dialogProps: PropTypes.object
+}
+
+export default PostProcessing

packages/ui/src/views/canvas/CanvasHeader.jsx

@@ -475,6 +475,7 @@ const CanvasHeader = ({ chatflow, isAgentCanvas, handleSaveFlow, handleDeleteFlo
                 show={chatflowConfigurationDialogOpen}
                 dialogProps={chatflowConfigurationDialogProps}
                 onCancel={() => setChatflowConfigurationDialogOpen(false)}
+                isAgentCanvas={isAgentCanvas}
             />
         </>
     )