Skip to content

Commit 561ea99

Browse files
HenryHengZJJJK801
HenryHengZJ
authored and
JJK801
committed
Feature/nim container (FlowiseAI#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
1 parent 27d5d10 commit 561ea99

File tree

10 files changed

+598
-19
lines changed

10 files changed

+598
-19
lines changed

‎packages/components/credentials/NvdiaNIMApi.credential.ts

+5-5
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { INodeParams, INodeCredential } from '../src/Interface'
22

3-
class NvdiaNIMApi implements INodeCredential {
3+
class NvidiaNIMApi implements INodeCredential {
44
label: string
55
name: string
66
version: number
@@ -9,16 +9,16 @@ class NvdiaNIMApi implements INodeCredential {
99

1010
constructor() {
1111
this.label = 'Nvdia NIM API Key'
12-
this.name = 'nvdiaNIMApi'
12+
this.name = 'nvidiaNIMApi'
1313
this.version = 1.0
1414
this.inputs = [
1515
{
16-
label: 'Nvdia NIM API Key',
17-
name: 'nvdiaNIMApiKey',
16+
label: 'Nvidia NIM API Key',
17+
name: 'nvidiaNIMApiKey',
1818
type: 'password'
1919
}
2020
]
2121
}
2222
}
2323

24-
module.exports = { credClass: NvdiaNIMApi }
24+
module.exports = { credClass: NvidiaNIMApi }

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

+13-13
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,10 @@ class ChatNvdiaNIM_ChatModels implements INode {
1616
inputs: INodeParams[]
1717

1818
constructor() {
19-
this.label = 'Chat Nvdia NIM'
20-
this.name = 'chatNvdiaNIM'
19+
this.label = 'Chat Nvidia NIM'
20+
this.name = 'chatNvidiaNIM'
2121
this.version = 1.0
22-
this.type = 'ChatNvdiaNIM'
22+
this.type = 'ChatNvidiaNIM'
2323
this.icon = 'nvdia.svg'
2424
this.category = 'Chat Models'
2525
this.description = 'Wrapper around Nvdia NIM Inference API'
@@ -28,7 +28,7 @@ class ChatNvdiaNIM_ChatModels implements INode {
2828
label: 'Connect Credential',
2929
name: 'credential',
3030
type: 'credential',
31-
credentialNames: ['nvdiaNIMApi'],
31+
credentialNames: ['nvidiaNIMApi'],
3232
optional: true
3333
}
3434
this.inputs = [
@@ -44,6 +44,13 @@ class ChatNvdiaNIM_ChatModels implements INode {
4444
type: 'string',
4545
placeholder: 'microsoft/phi-3-mini-4k-instruct'
4646
},
47+
{
48+
label: 'Base Path',
49+
name: 'basePath',
50+
type: 'string',
51+
description: 'Specify the URL of the deployed NIM Inference API',
52+
placeholder: 'https://integrate.api.nvidia.com/v1'
53+
},
4754
{
4855
label: 'Temperature',
4956
name: 'temperature',
@@ -52,13 +59,6 @@ class ChatNvdiaNIM_ChatModels implements INode {
5259
default: 0.9,
5360
optional: true
5461
},
55-
{
56-
label: 'Base Path',
57-
name: 'basePath',
58-
type: 'string',
59-
description: 'Specify the URL of the deployed NIM Inference API',
60-
placeholder: 'https://integrate.api.nvidia.com/v1'
61-
},
6262
{
6363
label: 'Streaming',
6464
name: 'streaming',
@@ -131,12 +131,12 @@ class ChatNvdiaNIM_ChatModels implements INode {
131131
const cache = nodeData.inputs?.cache as BaseCache
132132

133133
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
134-
const nvdiaNIMApiKey = getCredentialParam('nvdiaNIMApiKey', credentialData, nodeData)
134+
const nvidiaNIMApiKey = getCredentialParam('nvidiaNIMApiKey', credentialData, nodeData)
135135

136136
const obj: ChatOpenAIFields & { nvdiaNIMApiKey?: string } = {
137137
temperature: parseFloat(temperature),
138138
modelName,
139-
openAIApiKey: nvdiaNIMApiKey,
139+
openAIApiKey: nvidiaNIMApiKey ?? 'sk-',
140140
streaming: streaming ?? true
141141
}
142142

‎packages/server/package.json

+1
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@
9797
"multer": "^1.4.5-lts.1",
9898
"multer-s3": "^3.0.1",
9999
"mysql2": "^3.11.3",
100+
"nim-container-manager": "^1.0.4",
100101
"openai": "^4.82.0",
101102
"pg": "^8.11.1",
102103
"posthog-node": "^3.5.0",

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

+106
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import axios from 'axios'
2+
import { Request, Response, NextFunction } from 'express'
3+
4+
const { NimContainerManager } = require('nim-container-manager')
5+
6+
const getToken = async (req: Request, res: Response, next: NextFunction) => {
7+
try {
8+
const headers = {
9+
'Content-Type': 'application/json',
10+
Accept: 'application/json'
11+
}
12+
const data = {
13+
client_id: 'Flowise',
14+
pdi: '0x1234567890abcdeg',
15+
access_policy_name: 'nim-dev'
16+
}
17+
const response = await axios.post('https://nts.ngc.nvidia.com/v1/token', data, { headers })
18+
const responseJson = response.data
19+
return res.json(responseJson)
20+
} catch (error) {
21+
next(error)
22+
}
23+
}
24+
25+
const preload = async (req: Request, res: Response, next: NextFunction) => {
26+
try {
27+
await NimContainerManager.preload()
28+
return res.send('Preloaded NIM')
29+
} catch (error) {
30+
next(error)
31+
}
32+
}
33+
34+
const downloadInstaller = async (req: Request, res: Response, next: NextFunction) => {
35+
try {
36+
await NimContainerManager.downloadInstaller()
37+
return res.send('NIM Installer completed successfully!')
38+
} catch (error) {
39+
next(error)
40+
}
41+
}
42+
43+
const pullImage = async (req: Request, res: Response, next: NextFunction) => {
44+
try {
45+
const imageTag = req.body.imageTag
46+
const apiKey = req.body.apiKey
47+
await NimContainerManager.pullImage(imageTag, apiKey)
48+
return res.send(`Pulling image ${imageTag}`)
49+
} catch (error) {
50+
next(error)
51+
}
52+
}
53+
54+
const startContainer = async (req: Request, res: Response, next: NextFunction) => {
55+
try {
56+
const imageTag = req.body.imageTag
57+
const apiKey = req.body.apiKey
58+
await NimContainerManager.startContainer(imageTag, apiKey)
59+
return res.send(`Starting container ${imageTag}`)
60+
} catch (error) {
61+
next(error)
62+
}
63+
}
64+
65+
const getImage = async (req: Request, res: Response, next: NextFunction) => {
66+
try {
67+
const imageTag = req.body.imageTag
68+
const images = await NimContainerManager.userImageLibrary()
69+
const image = images.find((img: any) => img.tag === imageTag)
70+
if (!image) {
71+
return res.status(404).send(`Image ${imageTag} not found`)
72+
}
73+
return res.json(image)
74+
} catch (error) {
75+
next(error)
76+
}
77+
}
78+
79+
const getContainer = async (req: Request, res: Response, next: NextFunction) => {
80+
try {
81+
const imageTag = req.body.imageTag
82+
const images = await NimContainerManager.userImageLibrary()
83+
const image = images.find((img: any) => img.tag === imageTag)
84+
if (!image) {
85+
return res.status(404).send(`Image ${imageTag} not found`)
86+
}
87+
if (!image.container) {
88+
return res.status(404).send(`Container of ${imageTag} not found`)
89+
}
90+
const container = image.container
91+
container.image = image.name
92+
return res.json(container)
93+
} catch (error) {
94+
next(error)
95+
}
96+
}
97+
98+
export default {
99+
preload,
100+
getToken,
101+
downloadInstaller,
102+
pullImage,
103+
startContainer,
104+
getImage,
105+
getContainer
106+
}

‎packages/server/src/routes/index.ts

+2
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import variablesRouter from './variables'
4242
import vectorRouter from './vectors'
4343
import verifyRouter from './verify'
4444
import versionRouter from './versions'
45+
import nvidiaNimRouter from './nvidia-nim'
4546

4647
const router = express.Router()
4748

@@ -88,5 +89,6 @@ router.use('/vector', vectorRouter)
8889
router.use('/verify', verifyRouter)
8990
router.use('/version', versionRouter)
9091
router.use('/upsert-history', upsertHistoryRouter)
92+
router.use('/nvidia-nim', nvidiaNimRouter)
9193

9294
export default router

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

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import express from 'express'
2+
import nimController from '../../controllers/nvidia-nim'
3+
const router = express.Router()
4+
5+
// READ
6+
router.get('/preload', nimController.preload)
7+
router.get('/get-token', nimController.getToken)
8+
router.get('/download-installer', nimController.downloadInstaller)
9+
router.post('/pull-image', nimController.pullImage)
10+
router.post('/start-container', nimController.startContainer)
11+
router.post('/get-image', nimController.getImage)
12+
router.post('/get-container', nimController.getContainer)
13+
14+
export default router

‎packages/server/src/utils/constants.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ export const WHITELIST_URLS = [
1717
'/api/v1/ping',
1818
'/api/v1/version',
1919
'/api/v1/attachments',
20-
'/api/v1/metrics'
20+
'/api/v1/metrics',
21+
'/api/v1/nvidia-nim'
2122
]
2223

2324
export const OMIT_QUEUE_JOB_DATA = ['componentNodes', 'appDataSource', 'sseStreamer', 'telemetry', 'cachePool']

0 commit comments

Comments
 (0)