Skip to content

Commit 34702a9

Browse files
HenryHengZJHenryHengZJ
authored
Merge pull request #1224 from vinodkiran/FEATURE/mongodb
MongoDB Atlas Integration - Chat Memory and Vector Store
2 parents 998922b + e251bd5 commit 34702a9

File tree

10 files changed

+417
-2
lines changed

10 files changed

+417
-2
lines changed

‎packages/components/credentials/MongoDBUrlApi.credential.ts

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
import { INodeParams, INodeCredential } from '../src/Interface'
2+
3+
class MongoDBUrlApi implements INodeCredential {
4+
label: string
5+
name: string
6+
version: number
7+
description: string
8+
inputs: INodeParams[]
9+
10+
constructor() {
11+
this.label = 'MongoDB ATLAS'
12+
this.name = 'mongoDBUrlApi'
13+
this.version = 1.0
14+
this.inputs = [
15+
{
16+
label: 'ATLAS Connection URL',
17+
name: 'mongoDBConnectUrl',
18+
type: 'string',
19+
placeholder: 'mongodb+srv://myDatabaseUser:D1fficultP%[email protected]/?retryWrites=true&w=majority'
20+
}
21+
]
22+
}
23+
}
24+
25+
module.exports = { credClass: MongoDBUrlApi }

‎packages/components/nodes/memory/MongoDBMemory/MongoDBMemory.ts

+146
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
import { getBaseClasses, getCredentialData, getCredentialParam, ICommonObject, INode, INodeData, INodeParams } from '../../../src'
2+
import { MongoDBChatMessageHistory } from 'langchain/stores/message/mongodb'
3+
import { BufferMemory, BufferMemoryInput } from 'langchain/memory'
4+
import { BaseMessage, mapStoredMessageToChatMessage } from 'langchain/schema'
5+
import { MongoClient } from 'mongodb'
6+
7+
class MongoDB_Memory implements INode {
8+
label: string
9+
name: string
10+
version: number
11+
description: string
12+
type: string
13+
icon: string
14+
category: string
15+
baseClasses: string[]
16+
credential: INodeParams
17+
inputs: INodeParams[]
18+
19+
constructor() {
20+
this.label = 'MongoDB Atlas Chat Memory'
21+
this.name = 'MongoDBAtlasChatMemory'
22+
this.version = 1.0
23+
this.type = 'MongoDBAtlasChatMemory'
24+
this.icon = 'mongodb.png'
25+
this.category = 'Memory'
26+
this.description = 'Stores the conversation in MongoDB Atlas'
27+
this.baseClasses = [this.type, ...getBaseClasses(BufferMemory)]
28+
this.credential = {
29+
label: 'Connect Credential',
30+
name: 'credential',
31+
type: 'credential',
32+
credentialNames: ['mongoDBUrlApi']
33+
}
34+
this.inputs = [
35+
{
36+
label: 'Database',
37+
name: 'databaseName',
38+
placeholder: '<DB_NAME>',
39+
type: 'string'
40+
},
41+
{
42+
label: 'Collection Name',
43+
name: 'collectionName',
44+
placeholder: '<COLLECTION_NAME>',
45+
type: 'string'
46+
},
47+
{
48+
label: 'Session Id',
49+
name: 'sessionId',
50+
type: 'string',
51+
description: 'If not specified, the first CHAT_MESSAGE_ID will be used as sessionId',
52+
default: '',
53+
additionalParams: true,
54+
optional: true
55+
},
56+
{
57+
label: 'Memory Key',
58+
name: 'memoryKey',
59+
type: 'string',
60+
default: 'chat_history',
61+
additionalParams: true
62+
}
63+
]
64+
}
65+
66+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
67+
return initializeMongoDB(nodeData, options)
68+
}
69+
70+
async clearSessionMemory(nodeData: INodeData, options: ICommonObject): Promise<void> {
71+
const mongodbMemory = await initializeMongoDB(nodeData, options)
72+
const sessionId = nodeData.inputs?.sessionId as string
73+
const chatId = options?.chatId as string
74+
options.logger.info(`Clearing MongoDB memory session ${sessionId ? sessionId : chatId}`)
75+
await mongodbMemory.clear()
76+
options.logger.info(`Successfully cleared MongoDB memory session ${sessionId ? sessionId : chatId}`)
77+
}
78+
}
79+
80+
const initializeMongoDB = async (nodeData: INodeData, options: ICommonObject): Promise<BufferMemory> => {
81+
const databaseName = nodeData.inputs?.databaseName as string
82+
const collectionName = nodeData.inputs?.collectionName as string
83+
const sessionId = nodeData.inputs?.sessionId as string
84+
const memoryKey = nodeData.inputs?.memoryKey as string
85+
const chatId = options?.chatId as string
86+
87+
let isSessionIdUsingChatMessageId = false
88+
if (!sessionId && chatId) isSessionIdUsingChatMessageId = true
89+
90+
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
91+
let mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)
92+
93+
const client = new MongoClient(mongoDBConnectUrl)
94+
await client.connect()
95+
const collection = client.db(databaseName).collection(collectionName)
96+
97+
const mongoDBChatMessageHistory = new MongoDBChatMessageHistory({
98+
collection,
99+
sessionId: sessionId ? sessionId : chatId
100+
})
101+
102+
mongoDBChatMessageHistory.getMessages = async (): Promise<BaseMessage[]> => {
103+
const document = await collection.findOne({
104+
sessionId: (mongoDBChatMessageHistory as any).sessionId
105+
})
106+
const messages = document?.messages || []
107+
return messages.map(mapStoredMessageToChatMessage)
108+
}
109+
110+
mongoDBChatMessageHistory.addMessage = async (message: BaseMessage): Promise<void> => {
111+
const messages = [message].map((msg) => msg.toDict())
112+
await collection.updateOne(
113+
{ sessionId: (mongoDBChatMessageHistory as any).sessionId },
114+
{
115+
$push: { messages: { $each: messages } }
116+
},
117+
{ upsert: true }
118+
)
119+
}
120+
121+
mongoDBChatMessageHistory.clear = async (): Promise<void> => {
122+
await collection.deleteOne({ sessionId: (mongoDBChatMessageHistory as any).sessionId })
123+
}
124+
125+
return new BufferMemoryExtended({
126+
memoryKey,
127+
chatHistory: mongoDBChatMessageHistory,
128+
returnMessages: true,
129+
isSessionIdUsingChatMessageId
130+
})
131+
}
132+
133+
interface BufferMemoryExtendedInput {
134+
isSessionIdUsingChatMessageId: boolean
135+
}
136+
137+
class BufferMemoryExtended extends BufferMemory {
138+
isSessionIdUsingChatMessageId? = false
139+
140+
constructor(fields: BufferMemoryInput & Partial<BufferMemoryExtendedInput>) {
141+
super(fields)
142+
this.isSessionIdUsingChatMessageId = fields.isSessionIdUsingChatMessageId
143+
}
144+
}
145+
146+
module.exports = { nodeClass: MongoDB_Memory }

‎packages/components/nodes/memory/MongoDBMemory/mongodb.png

3.65 KB
Loading

‎packages/components/nodes/vectorstores/Elasticsearch/Elasticsearch_Upsert.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ class ElasicsearchUpsert_VectorStores extends ElasticSearchBase implements INode
5050
delete d.metadata.loc
5151
})
5252
// end of workaround
53-
return super.init(nodeData, _, options, flattenDocs)
53+
return super.init(nodeData, _, options, finalDocs)
5454
}
5555
}
5656

‎packages/components/nodes/vectorstores/MongoDB/MongoDBSearchBase.ts

+145
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
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 { MongoDBAtlasVectorSearch } from 'langchain/vectorstores/mongodb_atlas'
15+
import { Collection, MongoClient } from 'mongodb'
16+
17+
export abstract class MongoDBSearchBase {
18+
label: string
19+
name: string
20+
version: number
21+
description: string
22+
type: string
23+
icon: string
24+
category: string
25+
baseClasses: string[]
26+
inputs: INodeParams[]
27+
credential: INodeParams
28+
outputs: INodeOutputsValue[]
29+
mongoClient: MongoClient
30+
31+
protected constructor() {
32+
this.type = 'MongoDB Atlas'
33+
this.icon = 'mongodb.png'
34+
this.category = 'Vector Stores'
35+
this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever']
36+
this.credential = {
37+
label: 'Connect Credential',
38+
name: 'credential',
39+
type: 'credential',
40+
credentialNames: ['mongoDBUrlApi']
41+
}
42+
this.inputs = [
43+
{
44+
label: 'Embeddings',
45+
name: 'embeddings',
46+
type: 'Embeddings'
47+
},
48+
{
49+
label: 'Database',
50+
name: 'databaseName',
51+
placeholder: '<DB_NAME>',
52+
type: 'string'
53+
},
54+
{
55+
label: 'Collection Name',
56+
name: 'collectionName',
57+
placeholder: '<COLLECTION_NAME>',
58+
type: 'string'
59+
},
60+
{
61+
label: 'Index Name',
62+
name: 'indexName',
63+
placeholder: '<VECTOR_INDEX_NAME>',
64+
type: 'string'
65+
},
66+
{
67+
label: 'Content Field',
68+
name: 'textKey',
69+
description: 'Name of the field (column) that contains the actual content',
70+
type: 'string',
71+
default: 'text',
72+
additionalParams: true,
73+
optional: true
74+
},
75+
{
76+
label: 'Embedded Field',
77+
name: 'embeddingKey',
78+
description: 'Name of the field (column) that contains the Embedding',
79+
type: 'string',
80+
default: 'embedding',
81+
additionalParams: true,
82+
optional: true
83+
},
84+
{
85+
label: 'Top K',
86+
name: 'topK',
87+
description: 'Number of top results to fetch. Default to 4',
88+
placeholder: '4',
89+
type: 'number',
90+
additionalParams: true,
91+
optional: true
92+
}
93+
]
94+
this.outputs = [
95+
{
96+
label: 'MongoDB Retriever',
97+
name: 'retriever',
98+
baseClasses: this.baseClasses
99+
},
100+
{
101+
label: 'MongoDB Vector Store',
102+
name: 'vectorStore',
103+
baseClasses: [this.type, ...getBaseClasses(MongoDBAtlasVectorSearch)]
104+
}
105+
]
106+
}
107+
108+
abstract constructVectorStore(
109+
embeddings: Embeddings,
110+
collection: Collection,
111+
indexName: string,
112+
textKey: string,
113+
embeddingKey: string,
114+
docs: Document<Record<string, any>>[] | undefined
115+
): Promise<VectorStore>
116+
117+
async init(nodeData: INodeData, _: string, options: ICommonObject, docs: Document<Record<string, any>>[] | undefined): Promise<any> {
118+
const credentialData = await getCredentialData(nodeData.credential ?? '', options)
119+
const databaseName = nodeData.inputs?.databaseName as string
120+
const collectionName = nodeData.inputs?.collectionName as string
121+
const indexName = nodeData.inputs?.indexName as string
122+
let textKey = nodeData.inputs?.textKey as string
123+
let embeddingKey = nodeData.inputs?.embeddingKey as string
124+
const embeddings = nodeData.inputs?.embeddings as Embeddings
125+
const topK = nodeData.inputs?.topK as string
126+
const k = topK ? parseFloat(topK) : 4
127+
const output = nodeData.outputs?.output as string
128+
129+
let mongoDBConnectUrl = getCredentialParam('mongoDBConnectUrl', credentialData, nodeData)
130+
131+
this.mongoClient = new MongoClient(mongoDBConnectUrl)
132+
const collection = this.mongoClient.db(databaseName).collection(collectionName)
133+
if (!textKey || textKey === '') textKey = 'text'
134+
if (!embeddingKey || embeddingKey === '') embeddingKey = 'embedding'
135+
const vectorStore = await this.constructVectorStore(embeddings, collection, indexName, textKey, embeddingKey, docs)
136+
137+
if (output === 'retriever') {
138+
return vectorStore.asRetriever(k)
139+
} else if (output === 'vectorStore') {
140+
;(vectorStore as any).k = k
141+
return vectorStore
142+
}
143+
return vectorStore
144+
}
145+
}

‎packages/components/nodes/vectorstores/MongoDB/MongoDB_Existing.ts

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { Collection } from 'mongodb'
2+
import { MongoDBAtlasVectorSearch } from 'langchain/vectorstores/mongodb_atlas'
3+
import { Embeddings } from 'langchain/embeddings/base'
4+
import { VectorStore } from 'langchain/vectorstores/base'
5+
import { Document } from 'langchain/document'
6+
import { MongoDBSearchBase } from './MongoDBSearchBase'
7+
import { ICommonObject, INode, INodeData } from '../../../src/Interface'
8+
9+
class MongoDBExisting_VectorStores extends MongoDBSearchBase implements INode {
10+
constructor() {
11+
super()
12+
this.label = 'MongoDB Atlas Load Existing Index'
13+
this.name = 'MongoDBIndex'
14+
this.version = 1.0
15+
this.description = 'Load existing data from MongoDB Atlas (i.e: Document has been upserted)'
16+
}
17+
18+
async init(nodeData: INodeData, _: string, options: ICommonObject): Promise<any> {
19+
return super.init(nodeData, _, options, undefined)
20+
}
21+
22+
async constructVectorStore(
23+
embeddings: Embeddings,
24+
collection: Collection,
25+
indexName: string,
26+
textKey: string,
27+
embeddingKey: string,
28+
_: Document<Record<string, any>>[] | undefined
29+
): Promise<VectorStore> {
30+
return new MongoDBAtlasVectorSearch(embeddings, {
31+
collection: collection,
32+
indexName: indexName,
33+
textKey: textKey,
34+
embeddingKey: embeddingKey
35+
})
36+
}
37+
}
38+
39+
module.exports = { nodeClass: MongoDBExisting_VectorStores }

0 commit comments

Comments
 (0)