Addition of Cache Option for Chat Models

This commit is contained in:
vinodkiran
2023-10-08 11:44:43 +05:30
parent dbd655580d
commit 679eac1d0b
10 changed files with 115 additions and 16 deletions
@@ -1,6 +1,7 @@
import { INode, INodeData, INodeParams } from '../../../src/Interface'
import { getBaseClasses } from '../../../src/utils'
import { NIBittensorChatModel, BittensorInput } from 'langchain/experimental/chat_models/bittensor'
import { BaseCache } from 'langchain/schema'
class Bittensor_ChatModels implements INode {
label: string
@@ -16,13 +17,19 @@ class Bittensor_ChatModels implements INode {
constructor() {
this.label = 'NIBittensorChat'
this.name = 'NIBittensorChatModel'
this.version = 1.0
this.version = 2.0
this.type = 'BittensorChat'
this.icon = 'logo.png'
this.category = 'Chat Models'
this.description = 'Wrapper around Bittensor subnet 1 large language models'
this.baseClasses = [this.type, ...getBaseClasses(NIBittensorChatModel)]
this.inputs = [
{
label: 'Cache',
name: 'cache',
type: 'LLMCache',
optional: true
},
{
label: 'System prompt',
name: 'system_prompt',
@@ -35,9 +42,13 @@ class Bittensor_ChatModels implements INode {
async init(nodeData: INodeData, _: string): Promise<any> {
const system_prompt = nodeData.inputs?.system_prompt as string
const cache = nodeData.inputs?.llmCache as BaseCache
const obj: Partial<BittensorInput> = {
systemPrompt: system_prompt
}
if (cache) obj.cache = cache
const model = new NIBittensorChatModel(obj)
return model
}