支持使用 MongoDB Atlas 作为向量数据库

packages/vector-store-service/src/index.ts

@@ -26,6 +26,10 @@ export interface Config extends ChatLunaPlugin.Config {
     milvusUrl: string
     milvusUsername: string
     milvusPassword: string
+
+    mongodbUrl: string
+    mongodbDbName: string
+    mongodbCollectionName: string
 }
 
 export const Config: Schema<Config> = Schema.intersect([
@@ -35,7 +39,8 @@ export const Config: Schema<Config> = Schema.intersect([
                 Schema.const('faiss').description('Faiss'),
                 Schema.const('redis').description('Redis'),
                 Schema.const('milvus').description('Milvus'),
-                Schema.const('luna-vdb').description('lunavdb')
+                Schema.const('luna-vdb').description('lunavdb'),
+                Schema.const('mongodb').description('MongoDB Atlas')
             ])
         )
             .default(['luna-vdb'])
@@ -52,6 +57,14 @@ export const Config: Schema<Config> = Schema.intersect([
             .default('http://127.0.0.1:19530'),
         milvusUsername: Schema.string().default(''),
         milvusPassword: Schema.string().role('secret').default('')
+    }),
+
+    Schema.object({
+        mongodbUrl: Schema.string()
+            .role('url')
+            .default('mongodb://localhost:27017'),
+        mongodbDbName: Schema.string().default('chatluna'),
+        mongodbCollectionName: Schema.string().default('chatluna_collection')
     })
 ]).i18n({
     'zh-CN': require('./locales/zh-CN.schema.yml'),
@@ -69,6 +82,7 @@ export const usage = `
 
 要查看如何配置 Milvus 数据库，看[这里](https://js.langchain.com/docs/integrations/vectorstores/milvus/)
 
+要查看如何配置 MongoDB 数据库，看[这里](https://js.langchain.com/docs/integrations/vectorstores/mongodb_atlas/)
 
 目前配置 Faiss 数据库安装后可能会导致 koishi 环境不安全，如果安装完成后进行某些操作完成后出现了问题（如，升级 node 版本），开发者不对此负直接责任。
 `

packages/vector-store-service/src/vectorstore/mongodb.ts

@@ -0,0 +1,119 @@
+import { MongoDBAtlasVectorSearch } from '@langchain/mongodb'
+import { Context, Logger } from 'koishi'
+import { ChatLunaPlugin } from 'koishi-plugin-chatluna/services/chat'
+import { createLogger } from 'koishi-plugin-chatluna/utils/logger'
+import { Config } from '..'
+import { ChatLunaSaveableVectorStore } from 'koishi-plugin-chatluna/llm-core/model/base'
+import { MongoClient, ObjectId } from 'mongodb'
+
+let logger: Logger
+
+export async function apply(
+    ctx: Context,
+    config: Config,
+    plugin: ChatLunaPlugin
+) {
+    logger = createLogger(ctx, 'chatluna-vector-store-service')
+
+    if (!config.vectorStore.includes('mongodb')) {
+        return
+    }
+
+    await importMongoDB()
+
+    plugin.registerVectorStore('mongodb', async (params) => {
+        const embeddings = params.embeddings
+
+        const client = new MongoClient(config.mongodbUrl)
+        await client.connect()
+
+        ctx.on('dispose', async () => {
+            await client.close()
+            logger.info('MongoDB connection closed')
+        })
+
+        const collection = client
+            .db(config.mongodbDbName)
+            .collection(config.mongodbCollectionName)
+
+        const vectorStore = new MongoDBAtlasVectorSearch(embeddings, {
+            collection,
+            indexName: params.key ?? 'vector_index',
+            textKey: 'text',
+            embeddingKey: 'embedding'
+        })
+
+        const wrapperStore =
+            new ChatLunaSaveableVectorStore<MongoDBAtlasVectorSearch>(
+                vectorStore,
+                {
+                    async deletableFunction(_store, options) {
+                        if (options.deleteAll) {
+                            await collection.deleteMany({})
+                            return
+                        }
+
+                        const ids: string[] = []
+                        if (options.ids) {
+                            ids.push(...options.ids)
+                        }
+
+                        if (options.documents) {
+                            const documentIds = options.documents
+                                ?.map(
+                                    (document) =>
+                                        document.metadata?.raw_id as
+                                            | string
+                                            | undefined
+                                )
+                                .filter((id): id is string => id != null)
+
+                            ids.push(...documentIds)
+                        }
+
+                        if (ids.length > 0) {
+                            await collection.deleteMany({
+                                _id: { $in: ids.map((id) => new ObjectId(id)) }
+                            })
+                        }
+                    },
+                    async addDocumentsFunction(
+                        store,
+                        documents,
+                        options: { ids?: string[] }
+                    ) {
+                        let keys = options?.ids ?? []
+
+                        keys = documents.map((document, i) => {
+                            const id = keys[i] ?? crypto.randomUUID()
+                            document.metadata = {
+                                ...document.metadata,
+                                raw_id: id
+                            }
+                            return id
+                        })
+
+                        await store.addDocuments(documents)
+                    },
+                    async saveableFunction(_store) {
+                        await client.close()
+                        logger.info('MongoDB connection closed during save')
+                    }
+                }
+            )
+
+        return wrapperStore
+    })
+}
+
+async function importMongoDB() {
+    try {
+        const { MongoClient } = await import('mongodb')
+        return { MongoClient }
+    } catch (err) {
+        logger.error(err)
+        throw new Error(
+            'Please install mongodb as a dependency with, e.g. `npm install -S mongodb`'
+        )
+    }
+}