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)