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

import threading
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"]
# 豆包大模型的对话日志
chat_logs = db["chat_logs"]
# 聊天历史记录
chat_history_col = db["chat_history"]
# 兴趣圈集合
circle_prompts = db["circle_prompt"]
# 历史人物集合
historical_figures = db["historical_figures"]
# 知识库文档上传记录
knowledge_docs = db["knowledge_docs"]

def _ensure_index():
    try:
        chat_logs.create_index([("user_id", 1), ("asked_at", -1)])
        chat_history_col.create_index([("user_id", 1), ("session_id", 1), ("timestamp", -1)])
    except Exception:
        pass


def _do_save_chat_log(doc: dict):
    try:
        _ensure_index()
        chat_logs.insert_one(doc)
    except Exception as e:
        print(f"MongoDB 日志写入失败: {e}")


def save_chat_log(
    user_id: str,
    question: str,
    stream_mode: bool,
    raw_response: str = None,
    status: str = "success",
    error: str = None,
):
    """
    运维功能:保存聊天原始响应日志到 MongoDB（异步写入，不阻塞调用方）【日志审计】

    Args:
        user_id: 提问人
        question: 提问的问题
        stream_mode: 回答方式（流式或非流式）
        raw_response: API 原始响应的 repr 字符串
        status: 响应状态 success | error
        error: 异常时的错误信息
    """
    doc = {
        "user_id": user_id,
        "question": question,
        "stream_mode": stream_mode,
        "raw_response": raw_response,
        "status": status,
        "error": error,
        "asked_at": datetime.now(),
    }
    threading.Thread(target=_do_save_chat_log, args=(doc,), daemon=True).start()


def save_chat_history(
    user_id: str,
    session_id: str,
    role: str,
    content: str,
    timestamp: datetime,
    response_id: str = None,
    thinking: str = None,
    searching: str = None,
    attachments: list = None,
    input_tokens: int = None,
    output_tokens: int = None,
    total_tokens: int = None,
    app_name:str=None,
):
    """
    业务功能: 保存用户跟chat机器人的聊天记录，用于会话存储
    """

    try:
        _ensure_index()
        doc = {
            "user_id": user_id,
            "session_id": session_id,
            "role": role,
            "content": content,
            "thinking": thinking,
            "searching": searching,
            "response_id": response_id,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": total_tokens,
            "timestamp": timestamp,
            "app_name": app_name,
        }
        if attachments:
            doc["attachments"] = attachments
        chat_history_col.insert_one(doc)
    except Exception as e:
        print(f"MongoDB 聊天历史写入失败: {e}")


def get_chat_history(user_id: str, session_id: str) -> list:
    try:
        docs = chat_history_col.find(
            {"user_id": user_id, "session_id": session_id},
            {"_id": 0}
        ).sort("timestamp", 1)
        return list(docs)
    except Exception as e:
        print(f"MongoDB 聊天历史读取失败: {e}")
        return []


def get_last_response_id(user_id: str, session_id: str) -> str | None:
    try:
        doc = chat_history_col.find_one(
            {"user_id": user_id, "session_id": session_id, "role": "assistant", "response_id": {"$ne": None}},
            {"response_id": 1, "_id": 0},
            sort=[("timestamp", -1)]
        )
        return doc["response_id"] if doc else None
    except Exception as e:
        print(f"MongoDB 查询 response_id 失败: {e}")
        return None


def delete_chat_history(user_id: str, session_id: str) -> int:
    try:
        result = chat_history_col.delete_many({"user_id": user_id, "session_id": session_id})
        return result.deleted_count
    except Exception as e:
        print(f"MongoDB 聊天历史删除失败: {e}")
        return 0


def get_sessions(user_id: str) -> list:
    try:
        pipeline = [
            {"$match": {"user_id": user_id, "role": "user"}},
            {"$sort": {"timestamp": 1}},
            {"$group": {
                "_id": "$session_id",
                "createdAt": {"$first": "$timestamp"},
                "preview": {"$first": "$content"},
            }},
            {"$sort": {"createdAt": -1}},
            {"$project": {
                "_id": 0,
                "sessionId": "$_id",
                "createdAt": 1,
                "preview": {"$substrCP": ["$preview", 0, 20]},
            }},
        ]
        return list(chat_history_col.aggregate(pipeline))
    except Exception as e:
        print(f"MongoDB 会话列表查询失败: {e}")
        return []


_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


def upsert_circle_prompt(data: dict) -> None:
    circle_prompts.update_one(
        {"appName": data["appName"]},
        {"$set": data},
        upsert=True,
    )


# ===================== 历史人物 =====================

def get_all_figures() -> list:
    try:
        docs = historical_figures.find({})
        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 历史人物列表查询失败: {e}")
        return []


def get_figure_by_id(figure_id: str) -> dict | None:
    try:
        doc = historical_figures.find_one({"_id": ObjectId(figure_id)})
        if doc:
            doc["_id"] = str(doc["_id"])
        return doc
    except Exception as e:
        print(f"MongoDB 历史人物查询失败: {e}")
        return None


def insert_figure(data: dict) -> str:
    try:
        result = historical_figures.insert_one(data)
        return str(result.inserted_id)
    except Exception as e:
        print(f"MongoDB 历史人物新增失败: {e}")
        return None


def update_figure(figure_id: str, data: dict) -> int:
    try:
        result = historical_figures.update_one(
            {"_id": ObjectId(figure_id)},
            {"$set": data},
        )
        return result.matched_count
    except Exception as e:
        print(f"MongoDB 历史人物修改失败: {e}")
        return 0


def delete_figure(figure_id: str) -> int:
    try:
        result = historical_figures.delete_one({"_id": ObjectId(figure_id)})
        return result.deleted_count
    except Exception as e:
        print(f"MongoDB 历史人物删除失败: {e}")
        return 0


# ===================== 群聊AI回复历史 =====================

# 群聊AI回复记录集合（只存AI回复，不存群消息）
group_chat_history_col = db["group_chat_history"]


def _ensure_group_chat_index():
    try:
        group_chat_history_col.create_index([("session_id", 1), ("timestamp", -1)])
        group_chat_history_col.create_index([("group_id", 1), ("ai_role", 1), ("timestamp", -1)])
    except Exception:
        pass


def save_group_ai_reply(
    session_id: str,
    group_id: str,
    ai_role: str,
    content: str,
    response_id: str,
    trigger_message: str,
    trigger_user_id: str,
    timestamp: datetime,
    app_name: str = "souyue_im",
):
    """
    业务功能：保存AI在群聊中的回复

    注意：只存储AI回复，不存储群消息（群消息通过腾讯云API获取）

    Args:
        session_id: 会话ID（群ID + 小时窗口）
        group_id: 群ID
        ai_role: AI角色（entrepreneur | femininity）
        content: AI回复内容
        response_id: 大模型response ID
        trigger_message: 触发AI的原始消息
        trigger_user_id: 触发用户ID
        timestamp: 时间戳
        app_name: 应用名称
    """
    try:
        _ensure_group_chat_index()
        doc = {
            "session_id": session_id,
            "group_id": group_id,
            "role": "assistant",
            "user_id": f"ai_{ai_role}",
            "content": content,
            "ai_role": ai_role,
            "response_id": response_id,
            "trigger_message": trigger_message,
            "trigger_user_id": trigger_user_id,
            "timestamp": timestamp,
            "app_name": app_name,
        }
        group_chat_history_col.insert_one(doc)
    except Exception as e:
        print(f"MongoDB 群聊AI回复写入失败: {e}")


def get_last_group_response_id(session_id: str, ai_role: str) -> str | None:
    """
    获取指定AI的最后一次回复的response_id

    Args:
        session_id: 会话ID
        ai_role: AI角色（必须指定，防止串线）

    Returns:
        response_id 或 None
    """
    try:
        doc = group_chat_history_col.find_one(
            {
                "session_id": session_id,
                "ai_role": ai_role,
                "response_id": {"$ne": None}
            },
            {"response_id": 1, "_id": 0},
            sort=[("timestamp", -1)]
        )
        return doc["response_id"] if doc else None
    except Exception as e:
        print(f"MongoDB 查询群聊 response_id 失败: {e}")
        return None


def get_group_ai_reply_count(session_id: str, ai_role: str = None) -> int:
    """
    获取会话中的AI回复次数

    Args:
        session_id: 会话ID
        ai_role: AI角色（可选，不传则统计所有AI）

    Returns:
        回复次数
    """
    try:
        query = {"session_id": session_id}
        if ai_role:
            query["ai_role"] = ai_role
        return group_chat_history_col.count_documents(query)
    except Exception as e:
        print(f"MongoDB 查询AI回复次数失败: {e}")
        return 0



# ==================== 知识库文档上传记录 ====================

def _ensure_knowledge_docs_index():
    try:
        knowledge_docs.create_index([("app_name", 1), ("created_at", -1)])
        knowledge_docs.create_index([("doc_id", 1)], unique=True)
    except Exception:
        pass


def save_knowledge_doc(doc_data: dict) -> str | None:
    """
    保存知识库文档上传记录，返回插入的 ObjectId 字符串。
    doc_data 应包含：doc_id, resource_id, collection_name, project,
                     doc_name, doc_type, app_name
    """
    try:
        _ensure_knowledge_docs_index()
        doc_data["created_at"] = datetime.now()
        result = knowledge_docs.insert_one(doc_data)
        return str(result.inserted_id)
    except Exception as e:
        print(f"MongoDB 知识库文档记录写入失败: {e}")
        return None


def get_knowledge_docs(app_name: str, page: int = 1, page_size: int = 20) -> dict:
    """
    分页查询知识库文档上传记录。
    返回 {"total": int, "page": int, "page_size": int, "list": [...]}
    """
    try:
        _ensure_knowledge_docs_index()
        query = {"app_name": app_name}
        total = knowledge_docs.count_documents(query)
        skip = (page - 1) * page_size
        docs = list(
            knowledge_docs.find(query, {"_id": 1, "doc_id": 1, "resource_id": 1,
                                        "collection_name": 1, "project": 1,
                                        "doc_name": 1, "doc_type": 1,
                                        "app_name": 1, "created_at": 1})
            .sort("created_at", -1)
            .skip(skip)
            .limit(page_size)
        )
        for doc in docs:
            doc["id"] = str(doc.pop("_id"))
            if "created_at" in doc:
                doc["created_at"] = doc["created_at"].strftime("%Y-%m-%d %H:%M:%S")
        return {"total": total, "page": page, "page_size": page_size, "list": docs}
    except Exception as e:
        print(f"MongoDB 知识库文档记录查询失败: {e}")
        return {"total": 0, "page": page, "page_size": page_size, "list": []}


def get_knowledge_doc_by_doc_id(doc_id: str) -> dict | None:
    """
    按 doc_id 查询单条知识库文档记录。
    返回文档 dict（含 id 字符串字段），未找到返回 None。
    """
    try:
        doc = knowledge_docs.find_one({"doc_id": doc_id})
        if doc:
            doc["id"] = str(doc.pop("_id"))
        return doc
    except Exception as e:
        print(f"MongoDB 知识库文档记录查询失败: {e}")
        return None


def delete_knowledge_doc(doc_id: str) -> bool:
    """
    按 doc_id 删除知识库文档记录。
    返回 True 表示成功删除，False 表示记录不存在或删除异常。
    """
    try:
        result = knowledge_docs.delete_one({"doc_id": doc_id})
        return result.deleted_count > 0
    except Exception as e:
        print(f"MongoDB 知识库文档记录删除失败: {e}")
        return False


# ==================== 用户长期记忆 ====================

# facts 使用数组结构 [{key, value, updated_at}]，避免动态字段名带来的 Mongo 注入风险
user_memory_col = db["user_memory"]
user_memory_events_col = db["user_memory_events"]

_ALLOWED_FACT_KEYS = {"行业", "市场区域", "职能背景", "痛点", "目标", "产品", "融资阶段"}
_MAX_FACT_VALUE_LEN = 100
_MAX_EVENT_LEN = 200


def _ensure_user_memory_index():
    try:
        user_memory_col.create_index(
            [("user_id", 1), ("group_id", 1), ("ai_role", 1)], unique=True
        )
        user_memory_events_col.create_index(
            [("user_id", 1), ("group_id", 1), ("timestamp", -1)]
        )
        user_memory_events_col.create_index([("user_id", 1), ("keywords", 1)])
    except Exception:
        pass


def get_user_memory(user_id: str, group_id: str, ai_role: str) -> dict | None:
    """查询用户在指定群和角色下的记忆文档。"""
    try:
        _ensure_user_memory_index()
        return user_memory_col.find_one(
            {"user_id": user_id, "group_id": group_id, "ai_role": ai_role},
            {"_id": 0},
        )
    except Exception as e:
        print(f"MongoDB 查询用户记忆失败: {e}")
        return None


def upsert_user_fact(user_id: str, group_id: str, ai_role: str, key: str, value: str) -> None:
    """
    更新或新增用户事实（数组结构），按 key 匹配更新。
    key 必须在白名单内；value 超长截断。
    interaction_count 不在此处自增，由调用方统一自增一次。
    """
    if key not in _ALLOWED_FACT_KEYS:
        return
    value = str(value)[:_MAX_FACT_VALUE_LEN]
    try:
        _ensure_user_memory_index()
        # 先尝试更新已有 key
        result = user_memory_col.update_one(
            {"user_id": user_id, "group_id": group_id, "ai_role": ai_role,
             "facts": {"$elemMatch": {"key": key}}},
            {"$set": {"facts.$.value": value, "facts.$.updated_at": datetime.now(),
                      "last_active": datetime.now()}},
        )
        if result.matched_count == 0:
            # key 不存在，追加到数组
            user_memory_col.update_one(
                {"user_id": user_id, "group_id": group_id, "ai_role": ai_role},
                {
                    "$push": {"facts": {"key": key, "value": value, "updated_at": datetime.now()}},
                    "$set": {"last_active": datetime.now()},
                    "$setOnInsert": {"created_at": datetime.now(), "interaction_count": 0},
                },
                upsert=True,
            )
    except Exception as e:
        print(f"MongoDB 更新用户事实失败（key={key}）: {e}")


def increment_interaction_count(user_id: str, group_id: str, ai_role: str) -> None:
    """每次 _extract_memory 调用成功后，统一自增一次互动次数。"""
    try:
        user_memory_col.update_one(
            {"user_id": user_id, "group_id": group_id, "ai_role": ai_role},
            {
                "$inc": {"interaction_count": 1},
                "$set": {"last_active": datetime.now()},
                "$setOnInsert": {"created_at": datetime.now(), "facts": []},
            },
            upsert=True,
        )
    except Exception as e:
        print(f"MongoDB 自增互动次数失败: {e}")


def append_memory_event(
    user_id: str, group_id: str, ai_role: str,
    event: str, keywords: list, trigger_msg_id: str | None = None,
) -> None:
    """追加场景记忆事件。"""
    event = str(event)[:_MAX_EVENT_LEN]
    try:
        _ensure_user_memory_index()
        user_memory_events_col.insert_one({
            "user_id": user_id, "group_id": group_id, "ai_role": ai_role,
            "event": event,
            "keywords": [str(k)[:30] for k in keywords[:10]],
            "trigger_message_id": trigger_msg_id,
            "timestamp": datetime.now(),
        })
    except Exception as e:
        print(f"MongoDB 追加记忆事件失败: {e}")


def search_memory_events(
    user_id: str, group_id: str, keywords: list, limit: int = 5,
) -> list:
    """按关键词检索近期场景记忆事件，按时间倒序返回。"""
    if not keywords:
        return []
    try:
        docs = user_memory_events_col.find(
            {"user_id": user_id, "group_id": group_id,
             "keywords": {"$in": [str(k)[:30] for k in keywords[:10]]}},
            {"_id": 0, "event": 1, "timestamp": 1},
        ).sort("timestamp", -1).limit(limit)
        return list(docs)
    except Exception as e:
        print(f"MongoDB 检索记忆事件失败: {e}")
        return []