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

"""
知识库管理 - MongoDB 操作层
存储文档元数据、切片(chunks)、租户配置
"""
from pymongo import MongoClient
from typing import List, Optional, Dict, Any
from datetime import datetime
from bson.objectid import ObjectId
import os

# MongoDB 连接
MONGO_URI = os.getenv("ARK_LOGS_MONGO_URI", "mongodb://localhost:27017/")
client = MongoClient(MONGO_URI)
db = client["arklogs"]

# 集合定义
documents_collection = db["kb_documents"]  # 文档元数据
chunks_collection = db["kb_chunks"]  # 文档切片
tenants_collection = db["kb_tenants"]  # 租户配置


# ==================== 租户管理 ====================

def get_tenant_config(tenant_id: str) -> Optional[Dict]:
    """获取租户配置"""
    return tenants_collection.find_one({"tenant_id": tenant_id, "deleted": False})


def upsert_tenant_config(tenant_id: str, config: Dict) -> bool:
    """创建或更新租户配置"""
    config.update({
        "tenant_id": tenant_id,
        "updated_at": datetime.utcnow(),
    })
    result = tenants_collection.update_one(
        {"tenant_id": tenant_id},
        {"$set": config, "$setOnInsert": {"created_at": datetime.utcnow(), "deleted": False}},
        upsert=True
    )
    return result.acknowledged


def get_agent_config(tenant_id: str) -> Optional[Dict]:
    """获取租户的智能客服配置"""
    tenant = get_tenant_config(tenant_id)
    if not tenant:
        return None
    return {
        "enabled": tenant.get("agent_enabled", False),
        "name": tenant.get("agent_name", "智能小助手"),
        "greeting": tenant.get("agent_greeting", "您好！我是智能客服小助手，有什么可以帮您的吗？"),
        "h5_url": f"https://aichat.zhongsou.com/?tenant={tenant_id}",
    }


def update_agent_config(tenant_id: str, enabled: bool, name: str, greeting: str) -> bool:
    """更新智能客服配置"""
    result = tenants_collection.update_one(
        {"tenant_id": tenant_id},
        {
            "$set": {
                "agent_enabled": enabled,
                "agent_name": name,
                "agent_greeting": greeting,
                "updated_at": datetime.utcnow(),
            }
        },
        upsert=True
    )
    return result.acknowledged


# ==================== 文档管理 ====================

def create_document(tenant_id: str, filename: str, file_type: str, file_size: int, oss_path: str) -> str:
    """创建文档记录"""
    doc = {
        "tenant_id": tenant_id,
        "filename": filename,
        "file_type": file_type,
        "file_size": file_size,
        "oss_path": oss_path,
        "status": "pending",  # pending, processing, completed, failed
        "chunk_count": 0,
        "created_at": datetime.utcnow(),
        "updated_at": datetime.utcnow(),
        "deleted": False,
    }
    result = documents_collection.insert_one(doc)
    return str(result.inserted_id)


def get_document_by_id(doc_id: str) -> Optional[Dict]:
    """根据ID获取文档"""
    try:
        doc = documents_collection.find_one({"_id": ObjectId(doc_id), "deleted": False})
        if doc:
            doc["_id"] = str(doc["_id"])
        return doc
    except:
        return None


def list_documents(tenant_id: str, skip: int = 0, limit: int = 100) -> List[Dict]:
    """列出租户的所有文档"""
    cursor = documents_collection.find(
        {"tenant_id": tenant_id, "deleted": False}
    ).sort("created_at", -1).skip(skip).limit(limit)

    docs = []
    for doc in cursor:
        doc["_id"] = str(doc["_id"])
        docs.append(doc)
    return docs


def count_documents(tenant_id: str) -> int:
    """统计租户文档数量"""
    return documents_collection.count_documents({"tenant_id": tenant_id, "deleted": False})


def update_document_status(doc_id: str, status: str, chunk_count: int = 0, error_msg: str = None):
    """更新文档处理状态"""
    update_data = {
        "status": status,
        "updated_at": datetime.utcnow(),
    }
    if chunk_count > 0:
        update_data["chunk_count"] = chunk_count
    if error_msg:
        update_data["error_msg"] = error_msg

    documents_collection.update_one(
        {"_id": ObjectId(doc_id)},
        {"$set": update_data}
    )


def delete_document(doc_id: str) -> bool:
    """软删除文档"""
    result = documents_collection.update_one(
        {"_id": ObjectId(doc_id)},
        {"$set": {"deleted": True, "updated_at": datetime.utcnow()}}
    )
    return result.modified_count > 0


# ==================== 切片管理 ====================

def create_chunks(doc_id: str, tenant_id: str, chunks_data: List[Dict]) -> List[str]:
    """批量创建文档切片"""
    chunks = []
    for idx, chunk_data in enumerate(chunks_data):
        chunk = {
            "doc_id": doc_id,
            "tenant_id": tenant_id,
            "chunk_index": idx,
            "content": chunk_data["content"],
            "vector_id": chunk_data.get("vector_id"),  # Qdrant中的point_id
            "metadata": chunk_data.get("metadata", {}),
            "created_at": datetime.utcnow(),
            "deleted": False,
        }
        chunks.append(chunk)

    result = chunks_collection.insert_many(chunks)
    return [str(id) for id in result.inserted_ids]


def get_chunks_by_doc(doc_id: str) -> List[Dict]:
    """获取文档的所有切片"""
    cursor = chunks_collection.find(
        {"doc_id": doc_id, "deleted": False}
    ).sort("chunk_index", 1)

    chunks = []
    for chunk in cursor:
        chunk["_id"] = str(chunk["_id"])
        chunks.append(chunk)
    return chunks


def delete_chunks_by_doc(doc_id: str) -> int:
    """删除文档的所有切片"""
    result = chunks_collection.update_many(
        {"doc_id": doc_id},
        {"$set": {"deleted": True}}
    )
    return result.modified_count


def search_chunks_by_ids(chunk_ids: List[str]) -> List[Dict]:
    """根据ID列表获取切片（用于向量检索后的内容获取）"""
    chunks = []
    for chunk_id in chunk_ids:
        chunk = chunks_collection.find_one({"vector_id": chunk_id, "deleted": False})
        if chunk:
            chunk["_id"] = str(chunk["_id"])
            chunks.append(chunk)
    return chunks


# ==================== 统计 ====================

def get_tenant_stats(tenant_id: str) -> Dict:
    """获取租户的知识库统计数据"""
    total_docs = count_documents(tenant_id)
    total_chunks = chunks_collection.count_documents({"tenant_id": tenant_id, "deleted": False})

    # 获取最近上传的文档
    recent_docs = list_documents(tenant_id, skip=0, limit=5)

    return {
        "total_documents": total_docs,
        "total_chunks": total_chunks,
        "recent_documents": recent_docs,
    }