昇腾社区首页
中文
注册

类功能

功能描述

知识库管理的入口类,提供文档管理的功能。包括添加文档、删除文档,获取知识库中所有文档等功能。

函数原型

from mx_rag.knowledge import KnowledgeDB
KnowledgeDB(knowledge_store, chunk_store, vector_store, knowledge_name, white_paths, max_file_count, user_id, lock)

参数说明

参数名

数据类型

可选/必选

说明

knowledge_store

KnowledgeStore

必选

为知识库管理提供数据保存功能,保存已上传成功的文档名称。数据类型请参见KnowledgeStore类

chunk_store

Docstore

必选

为文档切片对象列表提供数据库存储功能。数据类型请参见Docstore类

vector_store

VectorStore

必选

向量数据库存储对象。数据类型请参见VectorStore

knowledge_name

str

必选

知识库名称,用户可根据知识库主题自定义,长度取值范围:[1, 1024]。

white_paths

List[str]

必选

上传文档的白名单路径列表。列表长度取值范围:[1, 1024],路径长度取值范围:[1, 1024],路径不能为软链接且不允许存在".."。

上传的文档路径需要在白名单路径下才能上传成功。

max_file_count

int

可选

上传文档时,允许的最大处理数量,取值范围:[1, 8000],不建议设置过大,默认值为“1000”

user_id

str

必选

用户ID,区分不同知识库,需要满足正则表达式^[a-zA-Z0-9_-]{6,64}$。

lock

multiprocessing.synchronize.Lock或_thread.LockType

可选

如果用户需要多进程或者多线程调用此接口需要申请锁。默认值为None。

可选值:

  • None:表示不使用锁,此时该接口不支持并发。
  • multiprocessing.Lock():表示进程锁,此时该接口支持多进程调用。
  • threading.Lock():表示线程锁。此时该接口支持多线程调用。

chunk_store和vector_store必须保障数据一致性(例如:关系数据库文件和向量数据库文件同时生成)。

调用示例

import pathlib
from paddle.base import libpaddle
from mx_rag.embedding.local import TextEmbedding
from mx_rag.knowledge import KnowledgeStore, KnowledgeDB
from mx_rag.storage.document_store import SQLiteDocstore
from mx_rag.storage.vectorstore import MindFAISS
# 设置向量检索使用的NPU卡
dev = 0
# 加载embedding模型
embed_func = TextEmbedding("/path/to/model", dev_id=dev)
# 初始化向量数据库
vector_store = MindFAISS(x_dim=1024, devs=[dev], 
                         load_local_index="./faiss.index", auto_save=True)
# 初始化文档chunk 关系数据库
chunk_store = SQLiteDocstore(db_path="./sql.db")
# 初始化知识管理关系数据库
knowledge_store = KnowledgeStore(db_path="./sql.db")
#添加知识库及管理员
knowledge_store.add_knowledge(knowledge_name="test", user_id='Default', role='admin')
# 初始化知识管理
knowledge_db = KnowledgeDB(knowledge_store=knowledge_store, chunk_store=chunk_store, vector_store=vector_store,
                           knowledge_name="test", user_id="Default", white_paths=["/home/"])
file_path = pathlib.Path("./gaokao.txt")
knowledge_db.add_file(file=file_path,
                      texts=["test1", "test2"],
                      embed_func={"dense": embed_func.embed_documents},
                      metadatas=[{"source": "./gaokao.txt"}, {"source": "./gaokao.txt"}])
documents =[document.document_name for document in knowledge_db.get_all_documents()]
print(documents)
print(knowledge_db.check_document_exist(doc_name=file_path.name))

knowledge_db.delete_file(doc_name=file_path.name)
knowledge_db.delete_all()