|
| 1 | +import { |
| 2 | + getBaseClasses, |
| 3 | + getCredentialData, |
| 4 | + getCredentialParam, |
| 5 | + ICommonObject, |
| 6 | + INodeData, |
| 7 | + INodeOutputsValue, |
| 8 | + INodeParams |
| 9 | +} from '../../../src' |
| 10 | + |
| 11 | +import { Embeddings } from 'langchain/embeddings/base' |
| 12 | +import { VectorStore } from 'langchain/vectorstores/base' |
| 13 | +import { Document } from 'langchain/document' |
| 14 | +import { createClient, SearchOptions } from 'redis' |
| 15 | +import { RedisVectorStore } from 'langchain/vectorstores/redis' |
| 16 | +import { escapeSpecialChars, unEscapeSpecialChars } from './utils' |
| 17 | + |
| 18 | +export abstract class RedisSearchBase { |
| 19 | + label: string |
| 20 | + name: string |
| 21 | + version: number |
| 22 | + description: string |
| 23 | + type: string |
| 24 | + icon: string |
| 25 | + category: string |
| 26 | + baseClasses: string[] |
| 27 | + inputs: INodeParams[] |
| 28 | + credential: INodeParams |
| 29 | + outputs: INodeOutputsValue[] |
| 30 | + redisClient: ReturnType<typeof createClient> |
| 31 | + |
| 32 | + protected constructor() { |
| 33 | + this.type = 'Redis' |
| 34 | + this.icon = 'redis.svg' |
| 35 | + this.category = 'Vector Stores' |
| 36 | + this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever'] |
| 37 | + this.credential = { |
| 38 | + label: 'Connect Credential', |
| 39 | + name: 'credential', |
| 40 | + type: 'credential', |
| 41 | + credentialNames: ['redisCacheUrlApi', 'redisCacheApi'] |
| 42 | + } |
| 43 | + this.inputs = [ |
| 44 | + { |
| 45 | + label: 'Embeddings', |
| 46 | + name: 'embeddings', |
| 47 | + type: 'Embeddings' |
| 48 | + }, |
| 49 | + { |
| 50 | + label: 'Index Name', |
| 51 | + name: 'indexName', |
| 52 | + placeholder: '<VECTOR_INDEX_NAME>', |
| 53 | + type: 'string' |
| 54 | + }, |
| 55 | + { |
| 56 | + label: 'Replace Index?', |
| 57 | + name: 'replaceIndex', |
| 58 | + description: 'Selecting this option will delete the existing index and recreate a new one', |
| 59 | + default: false, |
| 60 | + type: 'boolean' |
| 61 | + }, |
| 62 | + { |
| 63 | + label: 'Content Field', |
| 64 | + name: 'contentKey', |
| 65 | + description: 'Name of the field (column) that contains the actual content', |
| 66 | + type: 'string', |
| 67 | + default: 'content', |
| 68 | + additionalParams: true, |
| 69 | + optional: true |
| 70 | + }, |
| 71 | + { |
| 72 | + label: 'Metadata Field', |
| 73 | + name: 'metadataKey', |
| 74 | + description: 'Name of the field (column) that contains the metadata of the document', |
| 75 | + type: 'string', |
| 76 | + default: 'metadata', |
| 77 | + additionalParams: true, |
| 78 | + optional: true |
| 79 | + }, |
| 80 | + { |
| 81 | + label: 'Vector Field', |
| 82 | + name: 'vectorKey', |
| 83 | + description: 'Name of the field (column) that contains the vector', |
| 84 | + type: 'string', |
| 85 | + default: 'content_vector', |
| 86 | + additionalParams: true, |
| 87 | + optional: true |
| 88 | + }, |
| 89 | + { |
| 90 | + label: 'Top K', |
| 91 | + name: 'topK', |
| 92 | + description: 'Number of top results to fetch. Default to 4', |
| 93 | + placeholder: '4', |
| 94 | + type: 'number', |
| 95 | + additionalParams: true, |
| 96 | + optional: true |
| 97 | + } |
| 98 | + ] |
| 99 | + this.outputs = [ |
| 100 | + { |
| 101 | + label: 'Redis Retriever', |
| 102 | + name: 'retriever', |
| 103 | + baseClasses: this.baseClasses |
| 104 | + }, |
| 105 | + { |
| 106 | + label: 'Redis Vector Store', |
| 107 | + name: 'vectorStore', |
| 108 | + baseClasses: [this.type, ...getBaseClasses(RedisVectorStore)] |
| 109 | + } |
| 110 | + ] |
| 111 | + } |
| 112 | + |
| 113 | + abstract constructVectorStore( |
| 114 | + embeddings: Embeddings, |
| 115 | + indexName: string, |
| 116 | + replaceIndex: boolean, |
| 117 | + docs: Document<Record<string, any>>[] | undefined |
| 118 | + ): Promise<VectorStore> |
| 119 | + |
| 120 | + async init(nodeData: INodeData, _: string, options: ICommonObject, docs: Document<Record<string, any>>[] | undefined): Promise<any> { |
| 121 | + const credentialData = await getCredentialData(nodeData.credential ?? '', options) |
| 122 | + const indexName = nodeData.inputs?.indexName as string |
| 123 | + let contentKey = nodeData.inputs?.contentKey as string |
| 124 | + let metadataKey = nodeData.inputs?.metadataKey as string |
| 125 | + let vectorKey = nodeData.inputs?.vectorKey as string |
| 126 | + const embeddings = nodeData.inputs?.embeddings as Embeddings |
| 127 | + const topK = nodeData.inputs?.topK as string |
| 128 | + const replaceIndex = nodeData.inputs?.replaceIndex as boolean |
| 129 | + const k = topK ? parseFloat(topK) : 4 |
| 130 | + const output = nodeData.outputs?.output as string |
| 131 | + |
| 132 | + let redisUrl = getCredentialParam('redisUrl', credentialData, nodeData) |
| 133 | + if (!redisUrl || redisUrl === '') { |
| 134 | + const username = getCredentialParam('redisCacheUser', credentialData, nodeData) |
| 135 | + const password = getCredentialParam('redisCachePwd', credentialData, nodeData) |
| 136 | + const portStr = getCredentialParam('redisCachePort', credentialData, nodeData) |
| 137 | + const host = getCredentialParam('redisCacheHost', credentialData, nodeData) |
| 138 | + |
| 139 | + redisUrl = 'redis://' + username + ':' + password + '@' + host + ':' + portStr |
| 140 | + } |
| 141 | + |
| 142 | + this.redisClient = createClient({ url: redisUrl }) |
| 143 | + await this.redisClient.connect() |
| 144 | + |
| 145 | + const vectorStore = await this.constructVectorStore(embeddings, indexName, replaceIndex, docs) |
| 146 | + if (!contentKey || contentKey === '') contentKey = 'content' |
| 147 | + if (!metadataKey || metadataKey === '') metadataKey = 'metadata' |
| 148 | + if (!vectorKey || vectorKey === '') vectorKey = 'content_vector' |
| 149 | + |
| 150 | + const buildQuery = (query: number[], k: number, filter?: string[]): [string, SearchOptions] => { |
| 151 | + const vectorScoreField = 'vector_score' |
| 152 | + |
| 153 | + let hybridFields = '*' |
| 154 | + // if a filter is set, modify the hybrid query |
| 155 | + if (filter && filter.length) { |
| 156 | + // `filter` is a list of strings, then it's applied using the OR operator in the metadata key |
| 157 | + hybridFields = `@${metadataKey}:(${filter.map(escapeSpecialChars).join('|')})` |
| 158 | + } |
| 159 | + |
| 160 | + const baseQuery = `${hybridFields} => [KNN ${k} @${vectorKey} $vector AS ${vectorScoreField}]` |
| 161 | + const returnFields = [metadataKey, contentKey, vectorScoreField] |
| 162 | + |
| 163 | + const options: SearchOptions = { |
| 164 | + PARAMS: { |
| 165 | + vector: Buffer.from(new Float32Array(query).buffer) |
| 166 | + }, |
| 167 | + RETURN: returnFields, |
| 168 | + SORTBY: vectorScoreField, |
| 169 | + DIALECT: 2, |
| 170 | + LIMIT: { |
| 171 | + from: 0, |
| 172 | + size: k |
| 173 | + } |
| 174 | + } |
| 175 | + |
| 176 | + return [baseQuery, options] |
| 177 | + } |
| 178 | + |
| 179 | + vectorStore.similaritySearchVectorWithScore = async ( |
| 180 | + query: number[], |
| 181 | + k: number, |
| 182 | + filter?: string[] |
| 183 | + ): Promise<[Document, number][]> => { |
| 184 | + const results = await this.redisClient.ft.search(indexName, ...buildQuery(query, k, filter)) |
| 185 | + const result: [Document, number][] = [] |
| 186 | + |
| 187 | + if (results.total) { |
| 188 | + for (const res of results.documents) { |
| 189 | + if (res.value) { |
| 190 | + const document = res.value |
| 191 | + if (document.vector_score) { |
| 192 | + const metadataString = unEscapeSpecialChars(document[metadataKey] as string) |
| 193 | + result.push([ |
| 194 | + new Document({ |
| 195 | + pageContent: document[contentKey] as string, |
| 196 | + metadata: JSON.parse(metadataString) |
| 197 | + }), |
| 198 | + Number(document.vector_score) |
| 199 | + ]) |
| 200 | + } |
| 201 | + } |
| 202 | + } |
| 203 | + } |
| 204 | + return result |
| 205 | + } |
| 206 | + |
| 207 | + if (output === 'retriever') { |
| 208 | + return vectorStore.asRetriever(k) |
| 209 | + } else if (output === 'vectorStore') { |
| 210 | + ;(vectorStore as any).k = k |
| 211 | + return vectorStore |
| 212 | + } |
| 213 | + return vectorStore |
| 214 | + } |
| 215 | +} |
0 commit comments