ai-agent-admin/backend-fastapi/ai_platform/knowledge/services/knowledge_service.py

"""
知识库服务

知识库 CRUD 操作

数据权限：
- 使用 get_list_with_data_scope() 自动应用数据权限
- 支持本人、本部门、本部门及下级、全部等数据范围
"""
import logging
from typing import Optional, List, Tuple

from sqlalchemy import select, func, or_, and_
from sqlalchemy.ext.asyncio import AsyncSession

from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
from ai_platform.knowledge.schemas.knowledge_base_schema import KnowledgeBaseCreate, KnowledgeBaseUpdate
from app.data_scope_utils import get_data_scope_filter, apply_data_scope_to_conditions

logger = logging.getLogger(__name__)

# 资源类型（用于数据权限配置）
RESOURCE_TYPE = "knowledge_base"
RESOURCE_DISPLAY_NAME = "知识库管理"


class KnowledgeService:
    """知识库服务"""

    def __init__(self, db: AsyncSession):
        self._db = db

    async def get_list(
        self,
        page: int = 1,
        page_size: int = 20,
        name: Optional[str] = None,
        status: Optional[str] = None,
        application_id: Optional[str] = None,
    ) -> Tuple[List[KnowledgeBase], int]:
        """获取知识库列表"""
        query = select(KnowledgeBase).where(KnowledgeBase.is_deleted == False)

        if application_id:
            query = query.where(or_(
                KnowledgeBase.application_id == application_id,
                and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
            ))
        if name:
            query = query.where(KnowledgeBase.name.ilike(f"%{name}%"))
        if status:
            query = query.where(KnowledgeBase.status == status)

        # 总数
        count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
        total = count_result.scalar() or 0

        # 分页
        offset = (page - 1) * page_size
        query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
        query = query.offset(offset).limit(page_size)

        result = await self._db.execute(query)
        items = result.scalars().all()

        return items, total

    async def get_list_with_data_scope(
        self,
        page: int = 1,
        page_size: int = 20,
        name: Optional[str] = None,
        status: Optional[str] = None,
        application_id: Optional[str] = None,
    ) -> Tuple[List[KnowledgeBase], int]:
        """
        获取知识库列表（带数据权限过滤）
        
        自动从上下文获取当前用户信息，应用数据权限过滤
        """
        conditions = [KnowledgeBase.is_deleted == False]

        if application_id:
            conditions.append(or_(
                KnowledgeBase.application_id == application_id,
                and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
            ))
        if name:
            conditions.append(KnowledgeBase.name.ilike(f"%{name}%"))
        if status:
            conditions.append(KnowledgeBase.status == status)

        # 获取数据权限过滤条件并应用
        data_scope_filter = await get_data_scope_filter(self._db, RESOURCE_TYPE)
        scope_conditions = apply_data_scope_to_conditions(KnowledgeBase, data_scope_filter)
        conditions.extend(scope_conditions)

        # 总数
        query = select(KnowledgeBase).where(and_(*conditions))
        count_result = await self._db.execute(select(func.count()).select_from(query.subquery()))
        total = count_result.scalar() or 0

        # 分页
        offset = (page - 1) * page_size
        query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
        query = query.offset(offset).limit(page_size)

        result = await self._db.execute(query)
        items = result.scalars().all()

        return items, total

    async def get_by_id(self, kb_id: str) -> Optional[KnowledgeBase]:
        """获取知识库详情"""
        result = await self._db.execute(
            select(KnowledgeBase).where(
                KnowledgeBase.id == kb_id,
                KnowledgeBase.is_deleted == False
            )
        )
        return result.scalar_one_or_none()

    async def get_by_code(self, code: str) -> Optional[KnowledgeBase]:
        """根据编码获取知识库"""
        result = await self._db.execute(
            select(KnowledgeBase).where(
                KnowledgeBase.code == code,
                KnowledgeBase.is_deleted == False
            )
        )
        return result.scalar_one_or_none()

    async def create(self, data: KnowledgeBaseCreate) -> KnowledgeBase:
        """创建知识库"""
        # 检查编码唯一性
        existing = await self.get_by_code(data.code)
        if existing:
            raise ValueError(f'知识库编码 {data.code} 已存在')

        kb_data = data.model_dump()

        # 自动检测 embedding 模型的真实维度
        if data.embedding_model_id:
            try:
                from ai_platform.knowledge.services.embedding_service import EmbeddingService
                embedding_service = EmbeddingService(self._db)
                real_dim = await embedding_service.get_embedding_dimensions(data.embedding_model_id)
                kb_data['embedding_dimensions'] = real_dim
                logger.info(f'自动检测 embedding 维度: {real_dim}')
            except Exception as e:
                logger.warning(f'自动检测 embedding 维度失败，使用默认值: {e}')

        kb = KnowledgeBase(**kb_data)
        
        # 自动填充创建人和部门
        from utils.context import get_current_user_info_from_context
        user_info = get_current_user_info_from_context()
        if user_info:
            if not kb.sys_creator_id:
                kb.sys_creator_id = user_info.get('user_id')
            if not kb.sys_dept_id and user_info.get('dept_id'):
                kb.sys_dept_id = user_info.get('dept_id')
        
        self._db.add(kb)
        await self._db.commit()
        await self._db.refresh(kb)
        return kb

    async def update(self, kb_id: str, data: KnowledgeBaseUpdate) -> Optional[KnowledgeBase]:
        """更新知识库"""
        kb = await self.get_by_id(kb_id)
        if not kb:
            return None

        update_data = data.model_dump(exclude_unset=True)

        # 如果更换了 embedding 模型，自动重新检测维度
        new_model_id = update_data.get('embedding_model_id')
        if new_model_id and new_model_id != kb.embedding_model_id:
            try:
                from ai_platform.knowledge.services.embedding_service import EmbeddingService
                embedding_service = EmbeddingService(self._db)
                real_dim = await embedding_service.get_embedding_dimensions(new_model_id)
                update_data['embedding_dimensions'] = real_dim
                logger.info(f'更换模型后自动检测 embedding 维度: {real_dim}')
            except Exception as e:
                logger.warning(f'自动检测 embedding 维度失败: {e}')

        for key, value in update_data.items():
            setattr(kb, key, value)

        await self._db.commit()
        await self._db.refresh(kb)
        return kb

    async def delete(self, kb_id: str) -> bool:
        """删除知识库（软删除 + 清理 Qdrant collection）"""
        kb = await self.get_by_id(kb_id)
        if not kb:
            return False

        kb.is_deleted = True

        # 同时软删除所有文档和分段
        doc_result = await self._db.execute(
            select(KnowledgeDocument).where(
                KnowledgeDocument.knowledge_base_id == kb_id,
                KnowledgeDocument.is_deleted == False
            )
        )
        docs = doc_result.scalars().all()
        for doc in docs:
            doc.is_deleted = True

        # 软删除分段
        from sqlalchemy import update
        await self._db.execute(
            update(KnowledgeSegment).where(
                KnowledgeSegment.knowledge_base_id == kb_id
            ).values(is_deleted=True)
        )

        await self._db.commit()

        # 删除 Qdrant 中对应的 collection
        try:
            from ai_platform.knowledge.vector_store import get_vector_store
            vector_store = get_vector_store()
            await vector_store.delete_collection(kb_id)
        except Exception as e:
            logger.warning(f'删除 Qdrant collection 失败: {e}')

        return True

    async def get_simple_list(self, application_id: Optional[str] = None) -> List[dict]:
        """获取知识库简单列表（用于下拉选择）"""
        query = select(
            KnowledgeBase.id,
            KnowledgeBase.name,
            KnowledgeBase.code,
            KnowledgeBase.document_count,
            KnowledgeBase.segment_count,
        ).where(
            KnowledgeBase.is_deleted == False,
            KnowledgeBase.status == 'active',
        )

        if application_id:
            query = query.where(or_(
                KnowledgeBase.application_id == application_id,
                and_(KnowledgeBase.application_id.is_(None), KnowledgeBase.is_global == True)
            ))

        query = query.order_by(KnowledgeBase.sort.desc(), KnowledgeBase.sys_create_datetime.desc())
        result = await self._db.execute(query)
        rows = result.all()

        return [
            {
                'id': row.id,
                'name': row.name,
                'code': row.code,
                'document_count': row.document_count or 0,
                'segment_count': row.segment_count or 0,
            }
            for row in rows
        ]