Feature/nim container (#3966)

* add nim container setup

* check if image or container exist before pulling

* update NIM dialog

* update chat nvidia api key

* update nim container version

* update nim container version

packages/components/credentials/NvdiaNIMApi.credential.ts

@@ -1,6 +1,6 @@
 import { INodeParams, INodeCredential } from '../src/Interface'
 
-class NvdiaNIMApi implements INodeCredential {
+class NvidiaNIMApi implements INodeCredential {
     label: string
     name: string
     version: number
@@ -9,16 +9,16 @@ class NvdiaNIMApi implements INodeCredential {
 
     constructor() {
         this.label = 'Nvdia NIM API Key'
-        this.name = 'nvdiaNIMApi'
+        this.name = 'nvidiaNIMApi'
         this.version = 1.0
         this.inputs = [
             {
-                label: 'Nvdia NIM API Key',
+                label: 'Nvidia NIM API Key',
-                name: 'nvdiaNIMApiKey',
+                name: 'nvidiaNIMApiKey',
                 type: 'password'
             }
         ]
     }
 }
 
-module.exports = { credClass: NvdiaNIMApi }
+module.exports = { credClass: NvidiaNIMApi }

packages/components/nodes/chatmodels/ChatNvdiaNIM/ChatNvdiaNIM.ts

@@ -16,10 +16,10 @@ class ChatNvdiaNIM_ChatModels implements INode {
     inputs: INodeParams[]
 
     constructor() {
-        this.label = 'Chat Nvdia NIM'
+        this.label = 'Chat Nvidia NIM'
-        this.name = 'chatNvdiaNIM'
+        this.name = 'chatNvidiaNIM'
         this.version = 1.0
-        this.type = 'ChatNvdiaNIM'
+        this.type = 'ChatNvidiaNIM'
         this.icon = 'nvdia.svg'
         this.category = 'Chat Models'
         this.description = 'Wrapper around Nvdia NIM Inference API'
@@ -28,7 +28,7 @@ class ChatNvdiaNIM_ChatModels implements INode {
             label: 'Connect Credential',
             name: 'credential',
             type: 'credential',
-            credentialNames: ['nvdiaNIMApi'],
+            credentialNames: ['nvidiaNIMApi'],
             optional: true
         }
         this.inputs = [
@@ -44,6 +44,13 @@ class ChatNvdiaNIM_ChatModels implements INode {
                 type: 'string',
                 placeholder: 'microsoft/phi-3-mini-4k-instruct'
             },
+            {
+                label: 'Base Path',
+                name: 'basePath',
+                type: 'string',
+                description: 'Specify the URL of the deployed NIM Inference API',
+                placeholder: 'https://integrate.api.nvidia.com/v1'
+            },
             {
                 label: 'Temperature',
                 name: 'temperature',
@@ -52,13 +59,6 @@ class ChatNvdiaNIM_ChatModels implements INode {
                 default: 0.9,
                 optional: true
             },
-            {
-                label: 'Base Path',
-                name: 'basePath',
-                type: 'string',
-                description: 'Specify the URL of the deployed NIM Inference API',
-                placeholder: 'https://integrate.api.nvidia.com/v1'
-            },
             {
                 label: 'Streaming',
                 name: 'streaming',
@@ -131,12 +131,12 @@ class ChatNvdiaNIM_ChatModels implements INode {
         const cache = nodeData.inputs?.cache as BaseCache
 
         const credentialData = await getCredentialData(nodeData.credential ?? '', options)
-        const nvdiaNIMApiKey = getCredentialParam('nvdiaNIMApiKey', credentialData, nodeData)
+        const nvidiaNIMApiKey = getCredentialParam('nvidiaNIMApiKey', credentialData, nodeData)
 
         const obj: ChatOpenAIFields & { nvdiaNIMApiKey?: string } = {
             temperature: parseFloat(temperature),
             modelName,
-            openAIApiKey: nvdiaNIMApiKey,
+            openAIApiKey: nvidiaNIMApiKey ?? 'sk-',
             streaming: streaming ?? true
         }
 

packages/server/package.json

@@ -97,6 +97,7 @@
         "multer": "^1.4.5-lts.1",
         "multer-s3": "^3.0.1",
         "mysql2": "^3.11.3",
+        "nim-container-manager": "^1.0.4",
         "openai": "^4.82.0",
         "pg": "^8.11.1",
         "posthog-node": "^3.5.0",

packages/server/src/controllers/nvidia-nim/index.ts

@@ -0,0 +1,106 @@
+import axios from 'axios'
+import { Request, Response, NextFunction } from 'express'
+
+const { NimContainerManager } = require('nim-container-manager')
+
+const getToken = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const headers = {
+            'Content-Type': 'application/json',
+            Accept: 'application/json'
+        }
+        const data = {
+            client_id: 'Flowise',
+            pdi: '0x1234567890abcdeg',
+            access_policy_name: 'nim-dev'
+        }
+        const response = await axios.post('https://nts.ngc.nvidia.com/v1/token', data, { headers })
+        const responseJson = response.data
+        return res.json(responseJson)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const preload = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        await NimContainerManager.preload()
+        return res.send('Preloaded NIM')
+    } catch (error) {
+        next(error)
+    }
+}
+
+const downloadInstaller = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        await NimContainerManager.downloadInstaller()
+        return res.send('NIM Installer completed successfully!')
+    } catch (error) {
+        next(error)
+    }
+}
+
+const pullImage = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const apiKey = req.body.apiKey
+        await NimContainerManager.pullImage(imageTag, apiKey)
+        return res.send(`Pulling image ${imageTag}`)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const startContainer = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const apiKey = req.body.apiKey
+        await NimContainerManager.startContainer(imageTag, apiKey)
+        return res.send(`Starting container ${imageTag}`)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const getImage = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const images = await NimContainerManager.userImageLibrary()
+        const image = images.find((img: any) => img.tag === imageTag)
+        if (!image) {
+            return res.status(404).send(`Image ${imageTag} not found`)
+        }
+        return res.json(image)
+    } catch (error) {
+        next(error)
+    }
+}
+
+const getContainer = async (req: Request, res: Response, next: NextFunction) => {
+    try {
+        const imageTag = req.body.imageTag
+        const images = await NimContainerManager.userImageLibrary()
+        const image = images.find((img: any) => img.tag === imageTag)
+        if (!image) {
+            return res.status(404).send(`Image ${imageTag} not found`)
+        }
+        if (!image.container) {
+            return res.status(404).send(`Container of ${imageTag} not found`)
+        }
+        const container = image.container
+        container.image = image.name
+        return res.json(container)
+    } catch (error) {
+        next(error)
+    }
+}
+
+export default {
+    preload,
+    getToken,
+    downloadInstaller,
+    pullImage,
+    startContainer,
+    getImage,
+    getContainer
+}

packages/server/src/routes/index.ts

@@ -42,6 +42,7 @@ import variablesRouter from './variables'
 import vectorRouter from './vectors'
 import verifyRouter from './verify'
 import versionRouter from './versions'
+import nvidiaNimRouter from './nvidia-nim'
 
 const router = express.Router()
 
@@ -88,5 +89,6 @@ router.use('/vector', vectorRouter)
 router.use('/verify', verifyRouter)
 router.use('/version', versionRouter)
 router.use('/upsert-history', upsertHistoryRouter)
+router.use('/nvidia-nim', nvidiaNimRouter)
 
 export default router

packages/server/src/routes/nvidia-nim/index.ts

@@ -0,0 +1,14 @@
+import express from 'express'
+import nimController from '../../controllers/nvidia-nim'
+const router = express.Router()
+
+// READ
+router.get('/preload', nimController.preload)
+router.get('/get-token', nimController.getToken)
+router.get('/download-installer', nimController.downloadInstaller)
+router.post('/pull-image', nimController.pullImage)
+router.post('/start-container', nimController.startContainer)
+router.post('/get-image', nimController.getImage)
+router.post('/get-container', nimController.getContainer)
+
+export default router

packages/server/src/utils/constants.ts

@@ -17,7 +17,8 @@ export const WHITELIST_URLS = [
     '/api/v1/ping',
     '/api/v1/version',
     '/api/v1/attachments',
-    '/api/v1/metrics'
+    '/api/v1/metrics',
+    '/api/v1/nvidia-nim'
 ]
 
 export const OMIT_QUEUE_JOB_DATA = ['componentNodes', 'appDataSource', 'sseStreamer', 'telemetry', 'cachePool']

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

@@ -0,0 +1,328 @@
+import { useState, useEffect } from 'react'
+import { createPortal } from 'react-dom'
+import axios from 'axios'
+import PropTypes from 'prop-types'
+import {
+    Dialog,
+    DialogTitle,
+    DialogContent,
+    DialogActions,
+    Button,
+    CircularProgress,
+    Stepper,
+    Step,
+    StepLabel,
+    Select,
+    MenuItem,
+    FormControl,
+    InputLabel
+} from '@mui/material'
+
+const NvidiaNIMDialog = ({ open, onClose, onComplete }) => {
+    const portalElement = document.getElementById('portal')
+
+    const modelOptions = {
+        'nv-mistralai/mistral-nemo-12b-instruct:latest': {
+            label: 'Mistral Nemo 12B Instruct',
+            licenseUrl: 'https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nv-mistralai/containers/mistral-nemo-12b-instruct'
+        },
+        'meta/llama-3.1-8b-instruct-rtx:latest': {
+            label: 'Llama 3.1 8B Instruct',
+            licenseUrl: 'https://catalog.ngc.nvidia.com/orgs/nim/teams/meta/containers/llama-3.1-8b-instruct'
+        }
+    }
+
+    const [activeStep, setActiveStep] = useState(0)
+    const [loading, setLoading] = useState(false)
+    const [imageTag, setImageTag] = useState('')
+    const [pollInterval, setPollInterval] = useState(null)
+
+    const steps = ['Download Installer', 'Pull Image', 'Start Container']
+
+    const handleDownloadInstaller = async () => {
+        try {
+            setLoading(true)
+            await axios.get('/api/v1/nvidia-nim/download-installer')
+            setLoading(false)
+        } catch (err) {
+            let errorData = err.message
+            if (typeof err === 'string') {
+                errorData = err
+            } else if (err.response?.data) {
+                errorData = err.response.data.message
+            }
+            alert('Failed to download installer: ' + errorData)
+            setLoading(false)
+        }
+    }
+
+    const preload = async () => {
+        try {
+            setLoading(true)
+            await axios.get('/api/v1/nvidia-nim/preload')
+            setLoading(false)
+            setActiveStep(1)
+        } catch (err) {
+            let errorData = err.message
+            if (typeof err === 'string') {
+                errorData = err
+            } else if (err.response?.data) {
+                errorData = err.response.data.message
+            }
+            alert('Failed to preload: ' + errorData)
+            setLoading(false)
+        }
+    }
+
+    const handlePullImage = async () => {
+        try {
+            setLoading(true)
+            try {
+                const imageResponse = await axios.post('/api/v1/nvidia-nim/get-image', { imageTag })
+                if (imageResponse.data && imageResponse.data.tag === imageTag) {
+                    setLoading(false)
+                    setActiveStep(2)
+                    return
+                }
+            } catch (err) {
+                // Continue if image not found
+                if (err.response?.status !== 404) {
+                    throw err
+                }
+            }
+
+            // Get token first
+            const tokenResponse = await axios.get('/api/v1/nvidia-nim/get-token')
+            const apiKey = tokenResponse.data.access_token
+
+            // Pull image
+            await axios.post('/api/v1/nvidia-nim/pull-image', {
+                imageTag,
+                apiKey
+            })
+
+            // Start polling for image status
+            const interval = setInterval(async () => {
+                try {
+                    const imageResponse = await axios.post('/api/v1/nvidia-nim/get-image', { imageTag })
+                    if (imageResponse.data) {
+                        clearInterval(interval)
+                        setLoading(false)
+                        setActiveStep(2)
+                    }
+                } catch (err) {
+                    // Continue polling if image not found
+                    if (err.response?.status !== 404) {
+                        clearInterval(interval)
+                        alert('Failed to check image status: ' + err.message)
+                        setLoading(false)
+                    }
+                }
+            }, 5000)
+
+            setPollInterval(interval)
+        } catch (err) {
+            let errorData = err.message
+            if (typeof err === 'string') {
+                errorData = err
+            } else if (err.response?.data) {
+                errorData = err.response.data.message
+            }
+            alert('Failed to pull image: ' + errorData)
+            setLoading(false)
+        }
+    }
+
+    const handleStartContainer = async () => {
+        try {
+            setLoading(true)
+            try {
+                const containerResponse = await axios.post('/api/v1/nvidia-nim/get-container', { imageTag })
+                if (containerResponse.data && containerResponse.data && containerResponse.data.status === 'running') {
+                    // wait additional 10 seconds for container to be ready
+                    await new Promise((resolve) => setTimeout(resolve, 10000))
+                    setLoading(false)
+                    onComplete(containerResponse.data)
+                    onClose()
+                    return
+                }
+            } catch (err) {
+                // Continue if container not found
+                if (err.response?.status !== 404) {
+                    throw err
+                }
+            }
+
+            const tokenResponse = await axios.get('/api/v1/nvidia-nim/get-token')
+            const apiKey = tokenResponse.data.access_token
+
+            await axios.post('/api/v1/nvidia-nim/start-container', {
+                imageTag,
+                apiKey
+            })
+
+            // Start polling for container status
+            const interval = setInterval(async () => {
+                try {
+                    const containerResponse = await axios.post('/api/v1/nvidia-nim/get-container', { imageTag })
+                    if (containerResponse.data) {
+                        clearInterval(interval)
+                        setLoading(false)
+                        onComplete(containerResponse.data)
+                        onClose()
+                    }
+                } catch (err) {
+                    // Continue polling if container not found
+                    if (err.response?.status !== 404) {
+                        clearInterval(interval)
+                        alert('Failed to check container status: ' + err.message)
+                        setLoading(false)
+                    }
+                }
+            }, 5000)
+
+            setPollInterval(interval)
+        } catch (err) {
+            let errorData = err.message
+            if (typeof err === 'string') {
+                errorData = err
+            } else if (err.response?.data) {
+                errorData = err.response.data.message
+            }
+            alert('Failed to start container: ' + errorData)
+            setLoading(false)
+        }
+    }
+
+    const handleNext = () => {
+        if (activeStep === 1 && !imageTag) {
+            alert('Please enter an image tag')
+            return
+        }
+
+        switch (activeStep) {
+            case 0:
+                preload()
+                break
+            case 1:
+                handlePullImage()
+                break
+            case 2:
+                handleStartContainer()
+                break
+            default:
+                setActiveStep((prev) => prev + 1)
+        }
+    }
+
+    // Cleanup polling on unmount
+    useEffect(() => {
+        return () => {
+            if (pollInterval) {
+                clearInterval(pollInterval)
+            }
+        }
+    }, [pollInterval])
+
+    // clear state on close
+    useEffect(() => {
+        if (!open) {
+            setActiveStep(0)
+            setLoading(false)
+            setImageTag('')
+        }
+    }, [open])
+
+    const component = open ? (
+        <Dialog open={open}>
+            <DialogTitle>NIM Setup</DialogTitle>
+            <DialogContent>
+                <Stepper activeStep={activeStep}>
+                    {steps.map((label) => (
+                        <Step key={label}>
+                            <StepLabel>{label}</StepLabel>
+                        </Step>
+                    ))}
+                </Stepper>
+
+                {activeStep === 0 && (
+                    <div style={{ marginTop: 20 }}>
+                        <p style={{ marginBottom: 20 }}>
+                            Would you like to download the NIM installer? Click Next if it has been installed
+                        </p>
+                        {loading && <CircularProgress />}
+                    </div>
+                )}
+
+                {activeStep === 1 && (
+                    <div>
+                        <FormControl fullWidth sx={{ mt: 2 }}>
+                            <InputLabel>Model</InputLabel>
+                            <Select label='Model' value={imageTag} onChange={(e) => setImageTag(e.target.value)}>
+                                {Object.entries(modelOptions).map(([value, { label }]) => (
+                                    <MenuItem key={value} value={value}>
+                                        {label}
+                                    </MenuItem>
+                                ))}
+                            </Select>
+                        </FormControl>
+                        {imageTag && (
+                            <Button
+                                variant='text'
+                                size='small'
+                                sx={{ mt: 1 }}
+                                onClick={() => window.open(modelOptions[imageTag].licenseUrl, '_blank')}
+                            >
+                                View License
+                            </Button>
+                        )}
+                        {loading && (
+                            <div>
+                                <div style={{ marginBottom: 20 }} />
+                                <CircularProgress />
+                                <p>Pulling image...</p>
+                            </div>
+                        )}
+                    </div>
+                )}
+
+                {activeStep === 2 && (
+                    <div>
+                        {loading ? (
+                            <>
+                                <div style={{ marginBottom: 20 }} />
+                                <CircularProgress />
+                                <p>Starting container...</p>
+                            </>
+                        ) : (
+                            <p>Image is ready! Click Next to start the container.</p>
+                        )}
+                    </div>
+                )}
+            </DialogContent>
+            <DialogActions>
+                <Button onClick={onClose} variant='outline'>
+                    Cancel
+                </Button>
+                {activeStep === 0 && (
+                    <Button onClick={handleNext} variant='outline' color='secondary'>
+                        Next
+                    </Button>
+                )}
+                <Button onClick={activeStep === 0 ? handleDownloadInstaller : handleNext} disabled={loading}>
+                    {activeStep === 0 ? 'Download' : 'Next'}
+                </Button>
+            </DialogActions>
+        </Dialog>
+    ) : null
+
+    return createPortal(component, portalElement)
+}
+
+NvidiaNIMDialog.propTypes = {
+    open: PropTypes.bool,
+    onClose: PropTypes.func,
+    onComplete: PropTypes.func
+}
+
+export default NvidiaNIMDialog

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

@@ -37,6 +37,7 @@ import PromptLangsmithHubDialog from '@/ui-component/dialog/PromptLangsmithHubDi
 import ManageScrapedLinksDialog from '@/ui-component/dialog/ManageScrapedLinksDialog'
 import CredentialInputHandler from './CredentialInputHandler'
 import InputHintDialog from '@/ui-component/dialog/InputHintDialog'
+import NvidiaNIMDialog from '@/ui-component/dialog/NvidiaNIMDialog'
 
 // utils
 import { getInputVariables, getCustomConditionOutputs, isValidConnection, getAvailableNodesForVariable } from '@/utils/genericHelper'
@@ -95,6 +96,7 @@ const NodeInputHandler = ({
     const [inputHintDialogProps, setInputHintDialogProps] = useState({})
     const [showConditionDialog, setShowConditionDialog] = useState(false)
     const [conditionDialogProps, setConditionDialogProps] = useState({})
+    const [isNvidiaNIMDialogOpen, setIsNvidiaNIMDialogOpen] = useState(false)
     const [tabValue, setTabValue] = useState(0)
 
     const onInputHintDialogClicked = (hint) => {
@@ -443,6 +445,13 @@ const NodeInputHandler = ({
         setAsyncOptionEditDialog('')
     }
 
+    const handleNvidiaNIMDialogComplete = (containerData) => {
+        if (containerData) {
+            data.inputs['basePath'] = containerData.baseUrl
+            data.inputs['modelName'] = containerData.image
+        }
+    }
+
     useEffect(() => {
         if (ref.current && ref.current.offsetTop && ref.current.clientHeight) {
             setPosition(ref.current.offsetTop + ref.current.clientHeight / 2)
@@ -528,6 +537,22 @@ const NodeInputHandler = ({
                                     ></PromptLangsmithHubDialog>
                                 </>
                             )}
+                        {data.name === 'chatNvidiaNIM' && inputParam.name === 'modelName' && (
+                            <>
+                                <Button
+                                    style={{
+                                        display: 'flex',
+                                        flexDirection: 'row',
+                                        width: '100%'
+                                    }}
+                                    sx={{ borderRadius: '12px', width: '100%', mb: 2, mt: -1 }}
+                                    variant='outlined'
+                                    onClick={() => setIsNvidiaNIMDialogOpen(true)}
+                                >
+                                    Setup NIM Locally
+                                </Button>
+                            </>
+                        )}
                         <div style={{ display: 'flex', flexDirection: 'row' }}>
                             <Typography>
                                 {inputParam.label}
@@ -893,6 +918,11 @@ const NodeInputHandler = ({
                 dialogProps={inputHintDialogProps}
                 onCancel={() => setShowInputHintDialog(false)}
             ></InputHintDialog>
+            <NvidiaNIMDialog
+                open={isNvidiaNIMDialogOpen}
+                onClose={() => setIsNvidiaNIMDialogOpen(false)}
+                onComplete={handleNvidiaNIMDialogComplete}
+            ></NvidiaNIMDialog>
         </div>
     )
 }