Support cache system instructs for Google GenAI (#4148)
* Support cache system instructs for Google GenAI * format code * Update FlowiseGoogleAICacheManager.ts --------- Co-authored-by: Henry Heng <henryheng@flowiseai.com>
This commit is contained in:
parent
654bd48849
commit
d3510d1054
44
packages/components/nodes/cache/GoogleGenerativeAIContextCache/FlowiseGoogleAICacheManager.ts
vendored
Normal file
44
packages/components/nodes/cache/GoogleGenerativeAIContextCache/FlowiseGoogleAICacheManager.ts
vendored
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import type { CachedContentBase, CachedContent, Content } from '@google/generative-ai'
|
||||
import { GoogleAICacheManager as GoogleAICacheManagerBase } from '@google/generative-ai/server'
|
||||
import hash from 'object-hash'
|
||||
|
||||
type CacheContentOptions = Omit<CachedContentBase, 'contents'> & { contents?: Content[] }
|
||||
|
||||
export class GoogleAICacheManager extends GoogleAICacheManagerBase {
|
||||
private ttlSeconds: number
|
||||
private cachedContents: Map<string, CachedContent> = new Map()
|
||||
|
||||
setTtlSeconds(ttlSeconds: number) {
|
||||
this.ttlSeconds = ttlSeconds
|
||||
}
|
||||
|
||||
async lookup(options: CacheContentOptions): Promise<CachedContent | undefined> {
|
||||
const { model, tools, contents } = options
|
||||
if (!contents?.length) {
|
||||
return undefined
|
||||
}
|
||||
const hashKey = hash({
|
||||
model,
|
||||
tools,
|
||||
contents
|
||||
})
|
||||
if (this.cachedContents.has(hashKey)) {
|
||||
return this.cachedContents.get(hashKey)
|
||||
}
|
||||
const { cachedContents } = await this.list()
|
||||
const cachedContent = (cachedContents ?? []).find((cache) => cache.displayName === hashKey)
|
||||
if (cachedContent) {
|
||||
this.cachedContents.set(hashKey, cachedContent)
|
||||
return cachedContent
|
||||
}
|
||||
const res = await this.create({
|
||||
...(options as CachedContentBase),
|
||||
displayName: hashKey,
|
||||
ttlSeconds: this.ttlSeconds
|
||||
})
|
||||
this.cachedContents.set(hashKey, res)
|
||||
return res
|
||||
}
|
||||
}
|
||||
|
||||
export default GoogleAICacheManager
|
||||
34
packages/components/nodes/cache/GoogleGenerativeAIContextCache/GoogleGemini.svg
vendored
Normal file
34
packages/components/nodes/cache/GoogleGenerativeAIContextCache/GoogleGemini.svg
vendored
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
<svg width="32" height="32" viewBox="0 0 32 32" fill="none" xmlns="http://www.w3.org/2000/svg">
|
||||
<mask id="mask0_42_15021" style="mask-type:alpha" maskUnits="userSpaceOnUse" x="4" y="4" width="24" height="24">
|
||||
<path d="M16.9976 4.93059C16.9611 4.40651 16.5253 4 16 4C15.4747 4 15.0389 4.40651 15.0024 4.93059L14.951 5.66926C14.6048 10.645 10.645 14.6048 5.66926 14.951L4.93059 15.0024C4.40651 15.0389 4 15.4747 4 16C4 16.5253 4.40651 16.9611 4.93059 16.9976L5.66926 17.049C10.645 17.3952 14.6048 21.355 14.951 26.3307L15.0024 27.0694C15.0389 27.5935 15.4747 28 16 28C16.5253 28 16.9611 27.5935 16.9976 27.0694L17.049 26.3307C17.3952 21.355 21.355 17.3952 26.3307 17.049L27.0694 16.9976C27.5935 16.9611 28 16.5253 28 16C28 15.4747 27.5935 15.0389 27.0694 15.0024L26.3307 14.951C21.355 14.6048 17.3952 10.645 17.049 5.66926L16.9976 4.93059Z" fill="black"/>
|
||||
</mask>
|
||||
<g mask="url(#mask0_42_15021)">
|
||||
<path d="M16.9976 4.93059C16.9611 4.40651 16.5253 4 16 4C15.4747 4 15.0389 4.40651 15.0024 4.93059L14.951 5.66926C14.6048 10.645 10.645 14.6048 5.66926 14.951L4.93059 15.0024C4.40651 15.0389 4 15.4747 4 16C4 16.5253 4.40651 16.9611 4.93059 16.9976L5.66926 17.049C10.645 17.3952 14.6048 21.355 14.951 26.3307L15.0024 27.0694C15.0389 27.5935 15.4747 28 16 28C16.5253 28 16.9611 27.5935 16.9976 27.0694L17.049 26.3307C17.3952 21.355 21.355 17.3952 26.3307 17.049L27.0694 16.9976C27.5935 16.9611 28 16.5253 28 16C28 15.4747 27.5935 15.0389 27.0694 15.0024L26.3307 14.951C21.355 14.6048 17.3952 10.645 17.049 5.66926L16.9976 4.93059Z" fill="white"/>
|
||||
<g filter="url(#filter0_f_42_15021)">
|
||||
<circle cx="10.4616" cy="13.2307" r="8.30769" fill="#77B6E5"/>
|
||||
</g>
|
||||
<g filter="url(#filter1_f_42_15021)">
|
||||
<circle cx="16" cy="22.4615" r="8.30769" fill="#1E90C9"/>
|
||||
</g>
|
||||
<g filter="url(#filter2_f_42_15021)">
|
||||
<ellipse cx="21.5385" cy="10.4615" rx="10.1538" ry="8.30769" fill="#E9E5DF"/>
|
||||
</g>
|
||||
</g>
|
||||
<defs>
|
||||
<filter id="filter0_f_42_15021" x="-7.84613" y="-5.07697" width="36.6154" height="36.6154" filterUnits="userSpaceOnUse" color-interpolation-filters="sRGB">
|
||||
<feFlood flood-opacity="0" result="BackgroundImageFix"/>
|
||||
<feBlend mode="normal" in="SourceGraphic" in2="BackgroundImageFix" result="shape"/>
|
||||
<feGaussianBlur stdDeviation="5" result="effect1_foregroundBlur_42_15021"/>
|
||||
</filter>
|
||||
<filter id="filter1_f_42_15021" x="-0.307678" y="6.15381" width="32.6154" height="32.6154" filterUnits="userSpaceOnUse" color-interpolation-filters="sRGB">
|
||||
<feFlood flood-opacity="0" result="BackgroundImageFix"/>
|
||||
<feBlend mode="normal" in="SourceGraphic" in2="BackgroundImageFix" result="shape"/>
|
||||
<feGaussianBlur stdDeviation="4" result="effect1_foregroundBlur_42_15021"/>
|
||||
</filter>
|
||||
<filter id="filter2_f_42_15021" x="3.38464" y="-5.84619" width="36.3077" height="32.6154" filterUnits="userSpaceOnUse" color-interpolation-filters="sRGB">
|
||||
<feFlood flood-opacity="0" result="BackgroundImageFix"/>
|
||||
<feBlend mode="normal" in="SourceGraphic" in2="BackgroundImageFix" result="shape"/>
|
||||
<feGaussianBlur stdDeviation="4" result="effect1_foregroundBlur_42_15021"/>
|
||||
</filter>
|
||||
</defs>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 3.0 KiB |
53
packages/components/nodes/cache/GoogleGenerativeAIContextCache/GoogleGenerativeAIContextCache.ts
vendored
Normal file
53
packages/components/nodes/cache/GoogleGenerativeAIContextCache/GoogleGenerativeAIContextCache.ts
vendored
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
import { getBaseClasses, getCredentialData, getCredentialParam, ICommonObject, INode, INodeData, INodeParams } from '../../../src'
|
||||
import FlowiseGoogleAICacheManager from './FlowiseGoogleAICacheManager'
|
||||
|
||||
class GoogleGenerativeAIContextCache implements INode {
|
||||
label: string
|
||||
name: string
|
||||
version: number
|
||||
description: string
|
||||
type: string
|
||||
icon: string
|
||||
category: string
|
||||
baseClasses: string[]
|
||||
inputs: INodeParams[]
|
||||
credential: INodeParams
|
||||
|
||||
constructor() {
|
||||
this.label = 'Google GenAI Context Cache'
|
||||
this.name = 'googleGenerativeAIContextCache'
|
||||
this.version = 1.0
|
||||
this.type = 'GoogleAICacheManager'
|
||||
this.description = 'Large context cache for Google Gemini large language models'
|
||||
this.icon = 'GoogleGemini.svg'
|
||||
this.category = 'Cache'
|
||||
this.baseClasses = [this.type, ...getBaseClasses(FlowiseGoogleAICacheManager)]
|
||||
this.inputs = [
|
||||
{
|
||||
label: 'TTL',
|
||||
name: 'ttl',
|
||||
type: 'number',
|
||||
default: 60 * 60 * 24 * 30
|
||||
}
|
||||
]
|
||||
this.credential = {
|
||||
label: 'Connect Credential',
|
||||
name: 'credential',
|
||||
type: 'credential',
|
||||
credentialNames: ['googleGenerativeAI'],
|
||||
optional: false,
|
||||
description: 'Google Generative AI credential.'
|
||||
}
|
||||
}
|
||||
|
||||
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
|
||||
const ttl = nodeData.inputs?.ttl as number
|
||||
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
|
||||
const apiKey = getCredentialParam('googleGenerativeAPIKey', credentialData, nodeData)
|
||||
const manager = new FlowiseGoogleAICacheManager(apiKey)
|
||||
manager.setTtlSeconds(ttl)
|
||||
return manager
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { nodeClass: GoogleGenerativeAIContextCache }
|
||||
|
|
@ -5,6 +5,7 @@ import { ICommonObject, IMultiModalOption, INode, INodeData, INodeOptionsValue,
|
|||
import { convertMultiOptionsToStringArray, getBaseClasses, getCredentialData, getCredentialParam } from '../../../src/utils'
|
||||
import { getModels, MODEL_TYPE } from '../../../src/modelLoader'
|
||||
import { ChatGoogleGenerativeAI, GoogleGenerativeAIChatInput } from './FlowiseChatGoogleGenerativeAI'
|
||||
import type FlowiseGoogleAICacheManager from '../../cache/GoogleGenerativeAIContextCache/FlowiseGoogleAICacheManager'
|
||||
|
||||
class GoogleGenerativeAI_ChatModels implements INode {
|
||||
label: string
|
||||
|
|
@ -42,6 +43,12 @@ class GoogleGenerativeAI_ChatModels implements INode {
|
|||
type: 'BaseCache',
|
||||
optional: true
|
||||
},
|
||||
{
|
||||
label: 'Context Cache',
|
||||
name: 'contextCache',
|
||||
type: 'GoogleAICacheManager',
|
||||
optional: true
|
||||
},
|
||||
{
|
||||
label: 'Model Name',
|
||||
name: 'modelName',
|
||||
|
|
@ -188,6 +195,7 @@ class GoogleGenerativeAI_ChatModels implements INode {
|
|||
const harmCategory = nodeData.inputs?.harmCategory as string
|
||||
const harmBlockThreshold = nodeData.inputs?.harmBlockThreshold as string
|
||||
const cache = nodeData.inputs?.cache as BaseCache
|
||||
const contextCache = nodeData.inputs?.contextCache as FlowiseGoogleAICacheManager
|
||||
const streaming = nodeData.inputs?.streaming as boolean
|
||||
|
||||
const allowImageUploads = nodeData.inputs?.allowImageUploads as boolean
|
||||
|
|
@ -225,6 +233,7 @@ class GoogleGenerativeAI_ChatModels implements INode {
|
|||
|
||||
const model = new ChatGoogleGenerativeAI(nodeData.id, obj)
|
||||
model.setMultiModalOption(multiModalOption)
|
||||
if (contextCache) model.setContextCache(contextCache)
|
||||
|
||||
return model
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ import { StructuredToolInterface } from '@langchain/core/tools'
|
|||
import { isStructuredTool } from '@langchain/core/utils/function_calling'
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema'
|
||||
import { BaseLanguageModelCallOptions } from '@langchain/core/language_models/base'
|
||||
import type FlowiseGoogleAICacheManager from '../../cache/GoogleGenerativeAIContextCache/FlowiseGoogleAICacheManager'
|
||||
|
||||
const DEFAULT_IMAGE_MAX_TOKEN = 8192
|
||||
const DEFAULT_IMAGE_MODEL = 'gemini-1.5-flash-latest'
|
||||
|
|
@ -86,6 +87,8 @@ class LangchainChatGoogleGenerativeAI
|
|||
|
||||
private client: GenerativeModel
|
||||
|
||||
private contextCache?: FlowiseGoogleAICacheManager
|
||||
|
||||
get _isMultimodalModel() {
|
||||
return this.modelName.includes('vision') || this.modelName.startsWith('gemini-1.5')
|
||||
}
|
||||
|
|
@ -147,7 +150,7 @@ class LangchainChatGoogleGenerativeAI
|
|||
this.getClient()
|
||||
}
|
||||
|
||||
getClient(tools?: Tool[]) {
|
||||
async getClient(prompt?: Content[], tools?: Tool[]) {
|
||||
this.client = new GenerativeAI(this.apiKey ?? '').getGenerativeModel({
|
||||
model: this.modelName,
|
||||
tools,
|
||||
|
|
@ -161,6 +164,14 @@ class LangchainChatGoogleGenerativeAI
|
|||
topK: this.topK
|
||||
}
|
||||
})
|
||||
if (this.contextCache) {
|
||||
const cachedContent = await this.contextCache.lookup({
|
||||
contents: prompt ? [{ ...prompt[0], parts: prompt[0].parts.slice(0, 1) }] : [],
|
||||
model: this.modelName,
|
||||
tools
|
||||
})
|
||||
this.client.cachedContent = cachedContent as any
|
||||
}
|
||||
}
|
||||
|
||||
_combineLLMOutput() {
|
||||
|
|
@ -209,6 +220,10 @@ class LangchainChatGoogleGenerativeAI
|
|||
}
|
||||
}
|
||||
|
||||
setContextCache(contextCache: FlowiseGoogleAICacheManager): void {
|
||||
this.contextCache = contextCache
|
||||
}
|
||||
|
||||
async getNumTokens(prompt: BaseMessage[]) {
|
||||
const contents = convertBaseMessagesToContent(prompt, this._isMultimodalModel)
|
||||
const { totalTokens } = await this.client.countTokens({ contents })
|
||||
|
|
@ -226,9 +241,9 @@ class LangchainChatGoogleGenerativeAI
|
|||
this.convertFunctionResponse(prompt)
|
||||
|
||||
if (tools.length > 0) {
|
||||
this.getClient(tools as Tool[])
|
||||
await this.getClient(prompt, tools as Tool[])
|
||||
} else {
|
||||
this.getClient()
|
||||
await this.getClient(prompt)
|
||||
}
|
||||
const res = await this.caller.callWithOptions({ signal: options?.signal }, async () => {
|
||||
let output
|
||||
|
|
@ -296,9 +311,9 @@ class LangchainChatGoogleGenerativeAI
|
|||
|
||||
const tools = options.tools ?? []
|
||||
if (tools.length > 0) {
|
||||
this.getClient(tools as Tool[])
|
||||
await this.getClient(prompt, tools as Tool[])
|
||||
} else {
|
||||
this.getClient()
|
||||
await this.getClient(prompt)
|
||||
}
|
||||
|
||||
const stream = await this.caller.callWithOptions({ signal: options?.signal }, async () => {
|
||||
|
|
|
|||
Loading…
Reference in New Issue