From 595f1ed7f2ec634d25386b29d37febc1816abfb1 Mon Sep 17 00:00:00 2001
From: Keith Kacsh <keith@kacsh.com>
Date: Sat, 6 Jan 2024 17:16:06 -0700
Subject: [PATCH] Introduce new credential for LocalAI, Pass optional auth to
 LocalAI, New env var

---
 .../credentials/LcoalAIApi.credential.ts      | 23 +++++++++++++++
 .../chatmodels/ChatLocalAI/ChatLocalAI.ts     | 29 +++++++++++++++----
 packages/server/.env.example                  |  2 ++
 3 files changed, 48 insertions(+), 6 deletions(-)
 create mode 100644 packages/components/credentials/LcoalAIApi.credential.ts

diff --git a/packages/components/credentials/LcoalAIApi.credential.ts b/packages/components/credentials/LcoalAIApi.credential.ts
new file mode 100644
index 000000000..624e07fa4
--- /dev/null
+++ b/packages/components/credentials/LcoalAIApi.credential.ts
@@ -0,0 +1,23 @@
+import { INodeParams, INodeCredential } from '../src/Interface'
+
+class LocalAIApi implements INodeCredential {
+    label: string
+    name: string
+    version: number
+    inputs: INodeParams[]
+
+    constructor() {
+        this.label = 'LocalAI API'
+        this.name = 'LocalAIApi'
+        this.version = 1.0
+        this.inputs = [
+            {
+                label: 'LocalAI Api Key',
+                name: 'LocalAIApiKey',
+                type: 'password'
+            }
+        ]
+    }
+}
+
+module.exports = { credClass: LocalAIApi }
diff --git a/packages/components/nodes/chatmodels/ChatLocalAI/ChatLocalAI.ts b/packages/components/nodes/chatmodels/ChatLocalAI/ChatLocalAI.ts
index 18ed409bf..258db1f80 100644
--- a/packages/components/nodes/chatmodels/ChatLocalAI/ChatLocalAI.ts
+++ b/packages/components/nodes/chatmodels/ChatLocalAI/ChatLocalAI.ts
@@ -1,5 +1,5 @@
-import { INode, INodeData, INodeParams } from '../../../src/Interface'
-import { getBaseClasses } from '../../../src/utils'
+import { ICommonObject, INode, INodeData, INodeParams } from '../../../src/Interface'
+import { getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
 import { OpenAIChat } from 'langchain/llms/openai'
 import { OpenAIChatInput } from 'langchain/chat_models/openai'
 import { BaseCache } from 'langchain/schema'
@@ -14,6 +14,7 @@ class ChatLocalAI_ChatModels implements INode {
     category: string
     description: string
     baseClasses: string[]
+    credential: INodeParams
     inputs: INodeParams[]
 
     constructor() {
@@ -25,6 +26,16 @@ class ChatLocalAI_ChatModels implements INode {
         this.category = 'Chat Models'
         this.description = 'Use local LLMs like llama.cpp, gpt4all using LocalAI'
         this.baseClasses = [this.type, 'BaseChatModel', ...getBaseClasses(OpenAIChat)]
+        this.credential = {
+            label: 'Connect Credential',
+            name: 'credential',
+            type: 'credential',
+            credentialNames: ['LocalAIApi'],
+            optional: true
+        }
+
+        const modelOptions = JSON.parse(process.env.LOCALAI_CHAT_MODELS || '[]');
+        
         this.inputs = [
             {
                 label: 'Cache',
@@ -41,8 +52,10 @@ class ChatLocalAI_ChatModels implements INode {
             {
                 label: 'Model Name',
                 name: 'modelName',
-                type: 'string',
-                placeholder: 'gpt4all-lora-quantized.bin'
+                type: 'options',
+                options: modelOptions,
+                default: modelOptions.length > 0 ? modelOptions[0].name : '',
+                optional: true
             },
             {
                 label: 'Temperature',
@@ -79,19 +92,23 @@ class ChatLocalAI_ChatModels implements INode {
         ]
     }
 
-    async init(nodeData: INodeData): Promise<any> {
+    async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
         const temperature = nodeData.inputs?.temperature as string
         const modelName = nodeData.inputs?.modelName as string
         const maxTokens = nodeData.inputs?.maxTokens as string
         const topP = nodeData.inputs?.topP as string
         const timeout = nodeData.inputs?.timeout as string
         const basePath = nodeData.inputs?.basePath as string
+        
+        const credentialData = await getCredentialData(nodeData.credential ?? '', options)
+        const openAIApiKey = getCredentialParam('LocalAIApiKey', credentialData, nodeData)
+
         const cache = nodeData.inputs?.cache as BaseCache
 
         const obj: Partial<OpenAIChatInput> & BaseLLMParams & { openAIApiKey?: string } = {
             temperature: parseFloat(temperature),
             modelName,
-            openAIApiKey: 'sk-'
+            openAIApiKey
         }
 
         if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)
diff --git a/packages/server/.env.example b/packages/server/.env.example
index 6e746a4df..9b7be0ff8 100644
--- a/packages/server/.env.example
+++ b/packages/server/.env.example
@@ -26,3 +26,5 @@ PORT=3000
 # LANGCHAIN_ENDPOINT=https://api.smith.langchain.com
 # LANGCHAIN_API_KEY=your_api_key
 # LANGCHAIN_PROJECT=your_project
+
+# LOCALAI_CHAT_MODELS='[{"label": "model1", "name": "model1"}, {"label": "model2", "name": "model2"}]'