diff --git a/packages/vector-store-service/src/index.ts b/packages/vector-store-service/src/index.ts index 2f9ae978..291eb807 100644 --- a/packages/vector-store-service/src/index.ts +++ b/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 = Schema.intersect([ @@ -35,7 +39,8 @@ export const Config: Schema = 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 = 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 版本),开发者不对此负直接责任。 ` diff --git a/packages/vector-store-service/src/locales/en-US.schema.yml b/packages/vector-store-service/src/locales/en-US.schema.yml index 940b8d4e..2379e9f7 100644 --- a/packages/vector-store-service/src/locales/en-US.schema.yml +++ b/packages/vector-store-service/src/locales/en-US.schema.yml @@ -7,3 +7,7 @@ $inner: milvusUrl: Milvus URL Address milvusUsername: Milvus Username milvusPassword: Milvus Password + - $desc: MongoDB Database Settings + mongodbUrl: MongoDB URL Address + mongodbDbName: MongoDB Database Name + mongodbCollectionName: MongoDB Collection Name diff --git a/packages/vector-store-service/src/locales/zh-CN.schema.yml b/packages/vector-store-service/src/locales/zh-CN.schema.yml index 3c60191d..0415aaf7 100644 --- a/packages/vector-store-service/src/locales/zh-CN.schema.yml +++ b/packages/vector-store-service/src/locales/zh-CN.schema.yml @@ -7,3 +7,7 @@ $inner: milvusUrl: Milvus url 地址 milvusUsername: Milvus 用户名 milvusPassword: Milvus 密码 + - $desc: MongoDB 数据库设置 + mongodbUrl: MongoDB url 地址 + mongodbDbName: MongoDB 数据库名 + mongodbCollectionName: MongoDB 集合名 diff --git a/packages/vector-store-service/src/vectorstore/mongodb.ts b/packages/vector-store-service/src/vectorstore/mongodb.ts new file mode 100644 index 00000000..1ad605cf --- /dev/null +++ b/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( + 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`' + ) + } +}