73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from agno.document import Document
|
|
from agno.utils.log import logger
|
|
from dotenv import load_dotenv
|
|
|
|
from agentic_rag import initialize_knowledge_base, get_reader, initialize_mingci_knowledge_base
|
|
|
|
# 加载.env文件
|
|
load_dotenv()
|
|
import os
|
|
|
|
def main():
|
|
print("Hello from agno-agentic-rag!")
|
|
# 从.env加载知识库来源目录并初始化知识库
|
|
mingci_knowledge_source_dir = os.getenv("MINGCI_KNOWLEDGE_SOURCE_DIR")
|
|
if mingci_knowledge_source_dir and os.path.exists(mingci_knowledge_source_dir):
|
|
# 初始化知识库
|
|
knowledge_base = initialize_mingci_knowledge_base()
|
|
|
|
LoadMingCiKnowledgeToDatabase(knowledge_base, mingci_knowledge_source_dir)
|
|
|
|
knowledge_source_dir = os.getenv("KNOWLEDGE_SOURCE_DIR")
|
|
if knowledge_source_dir and os.path.exists(knowledge_source_dir):
|
|
# 初始化知识库
|
|
knowledge_base = initialize_knowledge_base()
|
|
|
|
LoadKnowledgeToDatabase(knowledge_base, knowledge_source_dir)
|
|
|
|
def LoadMingCiKnowledgeToDatabase(knowledge_base, knowledge_source_dir):
|
|
logger.info(f"加载知识库: {knowledge_source_dir}")
|
|
for root, _, files in os.walk(knowledge_source_dir):
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
file_ext = os.path.splitext(file)[1][1:] # 获取文件扩展名
|
|
reader = get_reader(file_ext)
|
|
if reader:
|
|
try:
|
|
docs = []
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
json_obj = json.load(f)
|
|
for item in json_obj:
|
|
file_contents = "{}\n 同义词: {}".format(item["name"], item["synonymous"])
|
|
|
|
docs.append(Document(
|
|
name=item["name"],
|
|
id=item["name"],
|
|
content=file_contents,
|
|
))
|
|
knowledge_base.load_documents(docs, upsert=True)
|
|
except Exception as e:
|
|
logger.warning(f"无法加载文档 {file_path}: {str(e)}")
|
|
|
|
def LoadKnowledgeToDatabase(knowledge_base, knowledge_source_dir):
|
|
logger.info(f"加载知识库: {knowledge_source_dir}")
|
|
for root, _, files in os.walk(knowledge_source_dir):
|
|
for file in files:
|
|
file_path = os.path.join(root, file)
|
|
file_ext = os.path.splitext(file)[1][1:] # 获取文件扩展名
|
|
reader = get_reader(file_ext)
|
|
if reader:
|
|
try:
|
|
filePath = Path(file_path)
|
|
docs: List[Document] = reader.read(filePath)
|
|
knowledge_base.load_documents(docs, upsert=True)
|
|
except Exception as e:
|
|
logger.warning(f"无法加载文档 {file_path}: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|