Maas-Agent/apps/knowledge/views/sample_center.py

# coding=utf-8
"""
样本中心知识库对接视图
提供样本中心知识库查询、批量入库等功能
"""
import uuid_utils.compat as uuid
from django.utils.translation import gettext as _
from rest_framework.request import Request
from rest_framework.views import APIView

from common.auth import TokenAuth
from common.exception.app_exception import AppApiException
from common.result import result
from knowledge.services.sample_center_client import get_sample_center_client


class SampleCenterView(APIView):
    """样本中心知识库对接"""
    authentication_classes = [TokenAuth]

    class ListKnowledgeBases(APIView):
        """查询样本中心知识库列表"""
        authentication_classes = [TokenAuth]

        def get(self, request: Request):
            page = int(request.query_params.get('page', 1))
            page_size = int(request.query_params.get('page_size', 20))
            base_url = request.query_params.get('base_url', '')
            app_id = request.query_params.get('app_id', '')
            app_secret = request.query_params.get('app_secret', '')

            if not all([base_url, app_id, app_secret]):
                raise AppApiException(400, _('base_url, app_id, app_secret are required'))

            client = get_sample_center_client(
                base_url=base_url,
                app_id=app_id,
                app_secret=app_secret,
            )
            data = client.list_knowledge_bases(page=page, page_size=page_size)
            return result.success(data)

    class GetKnowledgeBase(APIView):
        """查询样本中心知识库详情"""
        authentication_classes = [TokenAuth]

        def get(self, request: Request):
            kb_id = request.query_params.get('kb_id', '')
            base_url = request.query_params.get('base_url', '')
            app_id = request.query_params.get('app_id', '')
            app_secret = request.query_params.get('app_secret', '')

            if not all([kb_id, base_url, app_id, app_secret]):
                raise AppApiException(400, _('kb_id, base_url, app_id, app_secret are required'))

            client = get_sample_center_client(
                base_url=base_url,
                app_id=app_id,
                app_secret=app_secret,
            )
            data = client.get_knowledge_base(kb_id)
            return result.success(data)

    class BatchImport(APIView):
        """提交批量入库任务（直接转发到样本中心）"""
        authentication_classes = [TokenAuth]

        def post(self, request: Request):
            kb_id = request.data.get('kb_id', '')
            base_url = request.data.get('base_url', '')
            app_id = request.data.get('app_id', '')
            app_secret = request.data.get('app_secret', '')
            parents = request.data.get('parents', [])
            children = request.data.get('children', [])
            callback_url = request.data.get('callback_url', '')

            if not all([kb_id, base_url, app_id, app_secret, parents]):
                raise AppApiException(400, _('kb_id, base_url, app_id, app_secret, parents are required'))

            task_no = f"IMP{uuid.uuid4().hex[:16].upper()}"

            client = get_sample_center_client(
                base_url=base_url,
                app_id=app_id,
                app_secret=app_secret,
            )
            data = client.batch_import(
                kb_id=kb_id,
                task_no=task_no,
                parents=parents,
                children=children,
                callback_url=callback_url,
            )
            return result.success(data)

    class GetImportTask(APIView):
        """查询批量入库任务状态"""
        authentication_classes = [TokenAuth]

        def get(self, request: Request):
            task_id = request.query_params.get('task_id', '')
            base_url = request.query_params.get('base_url', '')
            app_id = request.query_params.get('app_id', '')
            app_secret = request.query_params.get('app_secret', '')

            if not all([task_id, base_url, app_id, app_secret]):
                raise AppApiException(400, _('task_id, base_url, app_id, app_secret are required'))

            client = get_sample_center_client(
                base_url=base_url,
                app_id=app_id,
                app_secret=app_secret,
            )
            data = client.get_import_task(task_id)
            return result.success(data)

    class SyncDocuments(APIView):
        """将本地知识库文档段落推送到样本中心"""
        authentication_classes = [TokenAuth]

        def post(self, request: Request, workspace_id: str, knowledge_id: str):
            from knowledge.models import Knowledge, Document, Paragraph, DocumentTag, Tag

            document_ids = request.data.get('document_ids', [])
            if not document_ids:
                raise AppApiException(400, _('document_ids is required'))

            # 获取知识库信息
            try:
                knowledge = Knowledge.objects.get(id=knowledge_id, workspace_id=workspace_id)
            except Knowledge.DoesNotExist:
                raise AppApiException(404, _('Knowledge base not found'))

            # 获取样本中心配置
            sample_center_config = knowledge.meta.get('sample_center', {})
            base_url = sample_center_config.get('base_url', '')
            app_id = sample_center_config.get('app_id', '')
            app_secret = sample_center_config.get('app_secret', '')
            kb_id = sample_center_config.get('kb_id', '')

            if not all([base_url, app_id, app_secret, kb_id]):
                raise AppApiException(400, _('Knowledge base is not linked to a sample center'))

            # 查询文档
            documents = Document.objects.filter(
                id__in=document_ids,
                knowledge_id=knowledge_id
            )
            if not documents.exists():
                raise AppApiException(400, _('No valid documents found'))

            # 查询段落
            paragraphs = Paragraph.objects.filter(
                document_id__in=document_ids,
                is_active=True
            ).order_by('document_id', 'position')

            if not paragraphs.exists():
                raise AppApiException(400, _('No paragraphs found in selected documents'))

            # 查询文档标签
            doc_tags = {}
            tag_mappings = DocumentTag.objects.filter(
                document_id__in=document_ids
            ).select_related('tag')
            for mapping in tag_mappings:
                doc_id = str(mapping.document_id)
                if doc_id not in doc_tags:
                    doc_tags[doc_id] = []
                doc_tags[doc_id].append(f"{mapping.tag.key}:{mapping.tag.value}")

            # 构建文档名称映射
            doc_name_map = {str(doc.id): doc.name for doc in documents}

            # 将段落转换为样本中心的 parents 格式
            # 每个段落作为一个 parent 条目
            parents = []
            for idx, paragraph in enumerate(paragraphs):
                doc_id = str(paragraph.document_id)
                parent_item = {
                    'index': idx,
                    'parent_id': kb_id,
                    'hierarchy': doc_name_map.get(doc_id, ''),
                    'text': paragraph.content,
                    'metadata': {
                        'source_knowledge_id': str(knowledge_id),
                        'source_document_id': doc_id,
                        'source_paragraph_id': str(paragraph.id),
                        'title': paragraph.title or '',
                    },
                    'doc_id': doc_id,
                    'tag_list': doc_tags.get(doc_id, []),
                }
                parents.append(parent_item)

            # 生成任务号
            task_no = f"IMP{uuid.uuid4().hex[:16].upper()}"

            # 调用样本中心批量入库接口
            client = get_sample_center_client(
                base_url=base_url,
                app_id=app_id,
                app_secret=app_secret,
            )
            data = client.batch_import(
                kb_id=kb_id,
                task_no=task_no,
                parents=parents,
            )

            return result.success({
                'task_id': data.get('task_id', ''),
                'task_no': task_no,
                'status': data.get('status', 'pending'),
                'total_paragraphs': len(parents),
                'total_documents': documents.count(),
            })