ai_config.py 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. from pymongo import MongoClient
  2. from datetime import datetime
  3. from dotenv import load_dotenv
  4. import os
  5. load_dotenv()
  6. MONGO_URI = os.getenv("ARK_LOGS_MONGO_URI")
  7. client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
  8. db = client["arklogs"]
  9. ai_config_col = db["ai_config"]
  10. def _ensure_index():
  11. try:
  12. ai_config_col.create_index([("appName", 1)], unique=True)
  13. ai_config_col.create_index([("status", 1), ("deleted", 1)])
  14. except Exception:
  15. pass
  16. def get_config_by_app_name(app_name: str) -> dict | None:
  17. """根据appName查询配置,只返回启用且未删除的记录"""
  18. try:
  19. _ensure_index()
  20. doc = ai_config_col.find_one({
  21. "appName": app_name,
  22. "status": 1,
  23. "deleted": 0
  24. })
  25. if doc:
  26. doc["_id"] = str(doc["_id"])
  27. return doc
  28. except Exception as e:
  29. print(f"MongoDB 查询 ai_config 失败: {e}")
  30. return None
  31. def create_config(data: dict) -> str | None:
  32. """创建新的配置"""
  33. try:
  34. _ensure_index()
  35. data["status"] = 1
  36. data["deleted"] = 0
  37. data["createdAt"] = datetime.now()
  38. data["updatedAt"] = datetime.now()
  39. result = ai_config_col.insert_one(data)
  40. return str(result.inserted_id)
  41. except Exception as e:
  42. print(f"MongoDB 创建 ai_config 失败: {e}")
  43. return None
  44. def update_config(app_name: str, data: dict) -> int:
  45. """更新配置"""
  46. try:
  47. _ensure_index()
  48. data["updatedAt"] = datetime.now()
  49. result = ai_config_col.update_one(
  50. {"appName": app_name},
  51. {"$set": data}
  52. )
  53. return result.matched_count
  54. except Exception as e:
  55. print(f"MongoDB 更新 ai_config 失败: {e}")
  56. return 0
  57. def delete_config(app_name: str) -> int:
  58. """软删除配置"""
  59. try:
  60. result = ai_config_col.update_one(
  61. {"appName": app_name},
  62. {"$set": {"deleted": 1, "updatedAt": datetime.now()}}
  63. )
  64. return result.matched_count
  65. except Exception as e:
  66. print(f"MongoDB 删除 ai_config 失败: {e}")
  67. return 0
  68. def get_all_configs() -> list:
  69. """获取所有启用的配置"""
  70. try:
  71. _ensure_index()
  72. docs = ai_config_col.find({
  73. "status": 1,
  74. "deleted": 0
  75. })
  76. return [{"_id": str(doc["_id"]), **{k: v for k, v in doc.items() if k != "_id"}} for doc in docs]
  77. except Exception as e:
  78. print(f"MongoDB 查询所有 ai_config 失败: {e}")
  79. return []