重构数据库访问,替换为统一的数据库实例引用

This commit is contained in:
晴猫
2025-03-12 22:27:59 +09:00
parent 49082267bb
commit 8be087dcad
19 changed files with 138 additions and 284 deletions

View File

@@ -14,7 +14,7 @@ root_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
sys.path.append(root_path)
# 现在可以导入src模块
from src.common.database import Database
from src.common.database import db
# 加载根目录下的env.edv文件
env_path = os.path.join(root_path, ".env.prod")
@@ -24,18 +24,6 @@ load_dotenv(env_path)
class KnowledgeLibrary:
def __init__(self):
# 初始化数据库连接
if Database._instance is None:
Database.initialize(
uri=os.getenv("MONGODB_URI"),
host=os.getenv("MONGODB_HOST", "127.0.0.1"),
port=int(os.getenv("MONGODB_PORT", "27017")),
db_name=os.getenv("DATABASE_NAME", "MegBot"),
username=os.getenv("MONGODB_USERNAME"),
password=os.getenv("MONGODB_PASSWORD"),
auth_source=os.getenv("MONGODB_AUTH_SOURCE"),
)
self.db = Database.get_instance()
self.raw_info_dir = "data/raw_info"
self._ensure_dirs()
self.api_key = os.getenv("SILICONFLOW_KEY")
@@ -176,7 +164,7 @@ class KnowledgeLibrary:
try:
current_hash = self.calculate_file_hash(file_path)
processed_record = self.db.processed_files.find_one({"file_path": file_path})
processed_record = db.processed_files.find_one({"file_path": file_path})
if processed_record:
if processed_record.get("hash") == current_hash:
@@ -197,14 +185,14 @@ class KnowledgeLibrary:
"split_length": knowledge_length,
"created_at": datetime.now()
}
self.db.knowledges.insert_one(knowledge)
db.knowledges.insert_one(knowledge)
result["chunks_processed"] += 1
split_by = processed_record.get("split_by", []) if processed_record else []
if knowledge_length not in split_by:
split_by.append(knowledge_length)
self.db.knowledges.processed_files.update_one(
db.knowledges.processed_files.update_one(
{"file_path": file_path},
{
"$set": {
@@ -322,7 +310,7 @@ class KnowledgeLibrary:
{"$project": {"content": 1, "similarity": 1, "file_path": 1}}
]
results = list(self.db.knowledges.aggregate(pipeline))
results = list(db.knowledges.aggregate(pipeline))
return results
# 创建单例实例
@@ -346,7 +334,7 @@ if __name__ == "__main__":
elif choice == '2':
confirm = input("确定要删除所有知识吗?这个操作不可撤销!(y/n): ").strip().lower()
if confirm == 'y':
knowledge_library.db.knowledges.delete_many({})
db.knowledges.delete_many({})
console.print("[green]已清空所有知识![/green]")
continue
elif choice == '1':