From 07ef701c2b4666a423ed5fcf2713bb4dd74e1064 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Thu, 19 Oct 2023 12:39:53 +0100
Subject: [PATCH] add im-mem cache embeddings

---
 .../InMemoryCache/InMemoryEmbeddingCache.ts   | 131 ++++++++++++++++++
 1 file changed, 131 insertions(+)
 create mode 100644 packages/components/nodes/cache/InMemoryCache/InMemoryEmbeddingCache.ts

diff --git a/packages/components/nodes/cache/InMemoryCache/InMemoryEmbeddingCache.ts b/packages/components/nodes/cache/InMemoryCache/InMemoryEmbeddingCache.ts
new file mode 100644
index 00000000..fad72482
--- /dev/null
+++ b/packages/components/nodes/cache/InMemoryCache/InMemoryEmbeddingCache.ts
@@ -0,0 +1,131 @@
+import { getBaseClasses, ICommonObject, INode, INodeData, INodeParams } from '../../../src'
+import { CacheBackedEmbeddings } from 'langchain/embeddings/cache_backed'
+import { Embeddings } from 'langchain/embeddings/base'
+import { BaseStore } from 'langchain/schema/storage'
+
+class InMemoryEmbeddingCache 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 = 'InMemory Embedding Cache'
+        this.name = 'inMemoryEmbeddingCache'
+        this.version = 1.0
+        this.type = 'InMemoryEmbeddingCache'
+        this.description = 'Cache generated Embeddings in memory to avoid needing to recompute them.'
+        this.icon = 'inmemorycache.png'
+        this.category = 'Cache'
+        this.baseClasses = [this.type, ...getBaseClasses(CacheBackedEmbeddings)]
+        this.inputs = [
+            {
+                label: 'Embeddings',
+                name: 'embeddings',
+                type: 'Embeddings'
+            },
+            {
+                label: 'Namespace',
+                name: 'namespace',
+                type: 'string',
+                optional: true,
+                additionalParams: true
+            }
+        ]
+    }
+
+    async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
+        const namespace = nodeData.inputs?.namespace as string
+        const underlyingEmbeddings = nodeData.inputs?.embeddings as Embeddings
+        const memoryMap = options.cachePool.getEmbeddingCache(options.chatflowid) ?? {}
+        const inMemCache = new InMemoryEmbeddingCacheExtended(memoryMap)
+
+        inMemCache.mget = async (keys: string[]) => {
+            const memory = options.cachePool.getEmbeddingCache(options.chatflowid) ?? inMemCache.store
+            return keys.map((key) => memory[key])
+        }
+
+        inMemCache.mset = async (keyValuePairs: [string, any][]): Promise<void> => {
+            for (const [key, value] of keyValuePairs) {
+                inMemCache.store[key] = value
+            }
+            options.cachePool.addEmbeddingCache(options.chatflowid, inMemCache.store)
+        }
+
+        inMemCache.mdelete = async (keys: string[]): Promise<void> => {
+            for (const key of keys) {
+                delete inMemCache.store[key]
+            }
+            options.cachePool.addEmbeddingCache(options.chatflowid, inMemCache.store)
+        }
+
+        return CacheBackedEmbeddings.fromBytesStore(underlyingEmbeddings, inMemCache, {
+            namespace: namespace
+        })
+    }
+}
+
+class InMemoryEmbeddingCacheExtended<T = any> extends BaseStore<string, T> {
+    lc_namespace = ['langchain', 'storage', 'in_memory']
+
+    store: Record<string, T> = {}
+
+    constructor(map: Record<string, T>) {
+        super()
+        this.store = map
+    }
+
+    /**
+     * Retrieves the values associated with the given keys from the store.
+     * @param keys Keys to retrieve values for.
+     * @returns Array of values associated with the given keys.
+     */
+    async mget(keys: string[]) {
+        return keys.map((key) => this.store[key])
+    }
+
+    /**
+     * Sets the values for the given keys in the store.
+     * @param keyValuePairs Array of key-value pairs to set in the store.
+     * @returns Promise that resolves when all key-value pairs have been set.
+     */
+    async mset(keyValuePairs: [string, T][]): Promise<void> {
+        for (const [key, value] of keyValuePairs) {
+            this.store[key] = value
+        }
+    }
+
+    /**
+     * Deletes the given keys and their associated values from the store.
+     * @param keys Keys to delete from the store.
+     * @returns Promise that resolves when all keys have been deleted.
+     */
+    async mdelete(keys: string[]): Promise<void> {
+        for (const key of keys) {
+            delete this.store[key]
+        }
+    }
+
+    /**
+     * Asynchronous generator that yields keys from the store. If a prefix is
+     * provided, it only yields keys that start with the prefix.
+     * @param prefix Optional prefix to filter keys.
+     * @returns AsyncGenerator that yields keys from the store.
+     */
+    async *yieldKeys(prefix?: string | undefined): AsyncGenerator<string> {
+        const keys = Object.keys(this.store)
+        for (const key of keys) {
+            if (prefix === undefined || key.startsWith(prefix)) {
+                yield key
+            }
+        }
+    }
+}
+
+module.exports = { nodeClass: InMemoryEmbeddingCache }