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/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']