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

from pymongo import MongoClient
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"]
# 豆包大模型的对话日志
chat_logs = db["chat_logs"]
# 兴趣圈集合
circle_prompts = db["circle_prompt"]

# 带下划线的表示私有方法（Private）
def _ensure_index():
    try:
        chat_logs.create_index([("username", 1), ("asked_at", -1)])
    except Exception:
        pass


def save_chat_log(
    username: str,
    question: str,
    stream_mode: bool,
    raw_response: str = None,
    status: str = "success",
    error: str = None,
):
    """
    保存聊天原始响应日志到 MongoDB

    Args:
        username: 提问人
        question: 提问的问题
        stream_mode: 回答方式（流式或非流式）
        raw_response: API 原始响应的 repr 字符串
        status: 响应状态 success | error
        error: 异常时的错误信息
    """
    try:
        _ensure_index()
        chat_logs.insert_one({
            "username": username,
            "question": question,
            "stream_mode": stream_mode,
            "raw_response": raw_response,
            "status": status,
            "error": error,
            "asked_at": datetime.now(),
        })
    except Exception as e:
        print(f"MongoDB 日志写入失败: {e}")


_DEFAULT_PROMPT_CONFIG = {
    "name": "兴趣圈",
    "role": "活跃用户",
    "style": "自然亲切，有活人感",
    "keywords": [],
    "forbidden": [],
}


def get_circle_prompt(app_name: str) -> dict:
    try:
        doc = circle_prompts.find_one({"appName": app_name})
        return doc if doc else _DEFAULT_PROMPT_CONFIG
    except Exception:
        return _DEFAULT_PROMPT_CONFIG