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:
Hans 2025-04-14 23:26:03 +08:00 committed by GitHub
parent 654bd48849
commit d3510d1054
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 160 additions and 5 deletions

View 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

View 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

View 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 }

View File

@ -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
}

View File

@ -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 () => {