70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
import asyncio
|
||
import logging
|
||
import os
|
||
|
||
from Util.DocxUtil import get_docx_content_by_pandoc
|
||
from Util.LightRagUtil import initialize_rag
|
||
|
||
# 是不是清空重新生成
|
||
IS_CLEAR = False
|
||
|
||
# 更详细地控制日志输出
|
||
logger = logging.getLogger('lightrag')
|
||
logger.setLevel(logging.INFO)
|
||
handler = logging.StreamHandler()
|
||
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||
logger.addHandler(handler)
|
||
|
||
|
||
async def main():
|
||
# 清空文件
|
||
if IS_CLEAR:
|
||
# 注释掉或删除以下清理代码
|
||
files_to_delete = [
|
||
"graph_chunk_entity_relation.graphml",
|
||
"kv_store_doc_status.json",
|
||
"kv_store_full_docs.json",
|
||
"kv_store_text_chunks.json",
|
||
"vdb_chunks.json",
|
||
"vdb_entities.json",
|
||
"vdb_relationships.json",
|
||
]
|
||
|
||
# 删除文件
|
||
for file in files_to_delete:
|
||
file_path = os.path.join(WORKING_DIR, file)
|
||
if os.path.exists(file_path):
|
||
os.remove(file_path)
|
||
logger.info(f"删除的文件:: {file_path}")
|
||
|
||
try:
|
||
# 注意：默认设置使用NetworkX
|
||
rag = await initialize_rag(WORKING_DIR)
|
||
# 在docx_file 目录下遍历所有以KEMU开头的文件
|
||
for filename in os.listdir(docx_file):
|
||
if filename.startswith(KEMU):
|
||
file_path = os.path.join(docx_file, filename)
|
||
# 获取docx文件的内容
|
||
if filename.endswith(".docx"):
|
||
continue
|
||
# content = get_docx_content_by_pandoc(file_path)
|
||
elif filename.endswith(".txt"):
|
||
with open(file_path, "r", encoding="utf-8") as f:
|
||
content = f.read()
|
||
await rag.ainsert(content, file_paths=[filename])
|
||
logger.info(f"Inserted content from {filename}")
|
||
except Exception as e:
|
||
logger.error(f"An error occurred: {e}")
|
||
finally:
|
||
await rag.finalize_storages()
|
||
|
||
|
||
# KEMUS=['JiHe','Math','SuShi','Chemistry','ShiJi','ChangChun']
|
||
KEMUS = ['Math2']
|
||
|
||
for KEMU in KEMUS:
|
||
# 组装文件路径
|
||
WORKING_DIR = "./Topic/" + KEMU
|
||
docx_file = 'static/Txt/'
|
||
asyncio.run(main())
|