From 7f89efee4b4201aa166c8ddba19e58336a936031 Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Sat, 5 Aug 2023 23:01:56 +0800 Subject: [PATCH] fix kb_doc_api.py --- server/knowledge_base/kb_doc_api.py | 4 ++-- server/knowledge_base/knowledge_base.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/server/knowledge_base/kb_doc_api.py b/server/knowledge_base/kb_doc_api.py index 71da9463..9330e699 100644 --- a/server/knowledge_base/kb_doc_api.py +++ b/server/knowledge_base/kb_doc_api.py @@ -54,7 +54,7 @@ async def upload_doc(file: UploadFile = File(description="上传文件"), kb_file = KnowledgeFile(filename=file.filename, knowledge_base_name=knowledge_base_name) kb = KnowledgeBase.load(knowledge_base_name=knowledge_base_name) - kb.add_file(kb_file) + kb.add_doc(kb_file) return BaseResponse(code=200, msg=f"成功上传文件 {file.filename}") @@ -115,7 +115,7 @@ async def recreate_vector_store(knowledge_base_name: str): knowledge_base_name=kb_name) print(f"processing {kb_file.filepath} to vector store.") kb = KnowledgeBase.load(knowledge_base_name=kb_name) - kb.add_file(kb_file) + kb.add_doc(kb_file) yield json.dumps({ "total": len(docs), "finished": i + 1, diff --git a/server/knowledge_base/knowledge_base.py b/server/knowledge_base/knowledge_base.py index f9dc05ea..deab8939 100644 --- a/server/knowledge_base/knowledge_base.py +++ b/server/knowledge_base/knowledge_base.py @@ -149,7 +149,7 @@ def list_docs_from_db(kb_name): conn.close() return kbs -def add_file_to_db(kb_file: KnowledgeFile): +def add_doc_to_db(kb_file: KnowledgeFile): conn = sqlite3.connect(DB_ROOT_PATH) c = conn.cursor() # Create table @@ -163,6 +163,7 @@ def add_file_to_db(kb_file: KnowledgeFile): file_version INTEGER, create_time DATETIME) ''') # Insert a row of data + # TODO: 同名文件添加至知识库时，file_version增加 c.execute(f"""INSERT INTO knowledge_files (file_name, file_ext, kb_name, document_loader_name, text_splitter_name, file_version, create_time) VALUES @@ -219,7 +220,7 @@ class KnowledgeBase: vector_store = FAISS.from_documents(docs, embeddings) # docs 为Document列表 torch_gc() vector_store.save_local(vs_path) - add_file_to_db(kb_file) + add_doc_to_db(kb_file) refresh_vs_cache(self.kb_name) elif self.vs_type in ["milvus"]: # TODO: 向milvus库中增加文件