1
1
import { getBaseClasses , getCredentialData , getCredentialParam , ICommonObject , INode , INodeData , INodeParams } from '../../../src'
2
2
import { RedisCache as LangchainRedisCache } from 'langchain/cache/ioredis'
3
3
import { Redis } from 'ioredis'
4
+ import { Generation , ChatGeneration , StoredGeneration , mapStoredMessageToChatMessage } from 'langchain/schema'
5
+ import hash from 'object-hash'
4
6
5
7
class RedisCache implements INode {
6
8
label : string
@@ -30,10 +32,21 @@ class RedisCache implements INode {
30
32
optional : true ,
31
33
credentialNames : [ 'redisCacheApi' ]
32
34
}
33
- this . inputs = [ ]
35
+ this . inputs = [
36
+ {
37
+ label : 'Time to Live (ms)' ,
38
+ name : 'ttl' ,
39
+ type : 'number' ,
40
+ step : 1 ,
41
+ optional : true ,
42
+ additionalParams : true
43
+ }
44
+ ]
34
45
}
35
46
36
47
async init ( nodeData : INodeData , _ : string , options : ICommonObject ) : Promise < any > {
48
+ const ttl = nodeData . inputs ?. ttl as string
49
+
37
50
const credentialData = await getCredentialData ( nodeData . credential ?? '' , options )
38
51
const username = getCredentialParam ( 'redisCacheUser' , credentialData , nodeData )
39
52
const password = getCredentialParam ( 'redisCachePwd' , credentialData , nodeData )
@@ -46,8 +59,60 @@ class RedisCache implements INode {
46
59
username,
47
60
password
48
61
} )
49
- return new LangchainRedisCache ( client )
62
+
63
+ const redisClient = new LangchainRedisCache ( client )
64
+
65
+ redisClient . lookup = async ( prompt : string , llmKey : string ) => {
66
+ let idx = 0
67
+ let key = getCacheKey ( prompt , llmKey , String ( idx ) )
68
+ let value = await client . get ( key )
69
+ const generations : Generation [ ] = [ ]
70
+
71
+ while ( value ) {
72
+ const storedGeneration = JSON . parse ( value )
73
+ generations . push ( deserializeStoredGeneration ( storedGeneration ) )
74
+ idx += 1
75
+ key = getCacheKey ( prompt , llmKey , String ( idx ) )
76
+ value = await client . get ( key )
77
+ }
78
+
79
+ return generations . length > 0 ? generations : null
80
+ }
81
+
82
+ redisClient . update = async ( prompt : string , llmKey : string , value : Generation [ ] ) => {
83
+ for ( let i = 0 ; i < value . length ; i += 1 ) {
84
+ const key = getCacheKey ( prompt , llmKey , String ( i ) )
85
+ if ( ttl !== undefined ) {
86
+ await client . set ( key , JSON . stringify ( serializeGeneration ( value [ i ] ) ) , 'EX' , parseInt ( ttl , 10 ) )
87
+ } else {
88
+ await client . set ( key , JSON . stringify ( serializeGeneration ( value [ i ] ) ) )
89
+ }
90
+ }
91
+ }
92
+
93
+ return redisClient
94
+ }
95
+ }
96
+
97
+ const getCacheKey = ( ...strings : string [ ] ) : string => hash ( strings . join ( '_' ) )
98
+ const deserializeStoredGeneration = ( storedGeneration : StoredGeneration ) => {
99
+ if ( storedGeneration . message !== undefined ) {
100
+ return {
101
+ text : storedGeneration . text ,
102
+ message : mapStoredMessageToChatMessage ( storedGeneration . message )
103
+ }
104
+ } else {
105
+ return { text : storedGeneration . text }
106
+ }
107
+ }
108
+ const serializeGeneration = ( generation : Generation ) => {
109
+ const serializedValue : StoredGeneration = {
110
+ text : generation . text
111
+ }
112
+ if ( ( generation as ChatGeneration ) . message !== undefined ) {
113
+ serializedValue . message = ( generation as ChatGeneration ) . message . toDict ( )
50
114
}
115
+ return serializedValue
51
116
}
52
117
53
118
module . exports = { nodeClass : RedisCache }
0 commit comments