chat-ai-api/app/db/ai_config.py

from pymongo import MongoClient
from bson import ObjectId
from datetime import datetime
from dotenv import load_dotenv
import os

load_dotenv()

MONGO_URI = os.getenv("ARK_LOGS_MONGO_URI")

client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
db = client["arklogs"]
ai_config_col = db["ai_config"]


def _ensure_index():
    try:
        ai_config_col.create_index([("appName", 1)], unique=True)
        ai_config_col.create_index([("status", 1), ("deleted", 1)])
    except Exception:
        pass


def get_config_by_app_name(app_name: str) -> dict | None:
    """根据appName查询配置，只返回启用且未删除的记录"""
    try:
        _ensure_index()
        doc = ai_config_col.find_one({
            "appName": app_name,
            "status": 1,
            "deleted": 0
        })
        if doc:
            doc["_id"] = str(doc["_id"])
        return doc
    except Exception as e:
        print(f"MongoDB 查询 ai_config 失败: {e}")
        return None


def create_config(data: dict) -> str | None:
    """创建新的配置"""
    try:
        _ensure_index()
        data["status"] = 1
        data["deleted"] = 0
        data["createdAt"] = datetime.now()
        data["updatedAt"] = datetime.now()
        result = ai_config_col.insert_one(data)
        return str(result.inserted_id)
    except Exception as e:
        print(f"MongoDB 创建 ai_config 失败: {e}")
        return None


def update_config(app_name: str, data: dict) -> int:
    """更新配置"""
    try:
        _ensure_index()
        data["updatedAt"] = datetime.now()
        result = ai_config_col.update_one(
            {"appName": app_name},
            {"$set": data}
        )
        return result.matched_count
    except Exception as e:
        print(f"MongoDB 更新 ai_config 失败: {e}")
        return 0


def delete_config(app_name: str) -> int:
    """软删除配置"""
    try:
        result = ai_config_col.update_one(
            {"appName": app_name},
            {"$set": {"deleted": 1, "updatedAt": datetime.now()}}
        )
        return result.matched_count
    except Exception as e:
        print(f"MongoDB 删除 ai_config 失败: {e}")
        return 0


def get_all_configs() -> list:
    """获取所有启用的配置"""
    try:
        _ensure_index()
        docs = ai_config_col.find({
            "status": 1,
            "deleted": 0
        })
        return [{"_id": str(doc["_id"]), **{k: v for k, v in doc.items() if k != "_id"}} for doc in docs]
    except Exception as e:
        print(f"MongoDB 查询所有 ai_config 失败: {e}")
        return []