Underground-Digital/Workflow-Engine

1from flask_restful import Resource, reqparse2 3from controllers.console import api4from controllers.console.datasets.error import WebsiteCrawlError5from controllers.console.wraps import account_initialization_required, setup_required6from libs.login import login_required7from services.website_service import WebsiteService8 9 10class WebsiteCrawlApi(Resource):11    @setup_required12    @login_required13    @account_initialization_required14    def post(self):15        parser = reqparse.RequestParser()16        parser.add_argument(17            "provider", type=str, choices=["firecrawl", "jinareader"], required=True, nullable=True, location="json"18        )19        parser.add_argument("url", type=str, required=True, nullable=True, location="json")20        parser.add_argument("options", type=dict, required=True, nullable=True, location="json")21        args = parser.parse_args()22        WebsiteService.document_create_args_validate(args)23        # crawl url24        try:25            result = WebsiteService.crawl_url(args)26        except Exception as e:27            raise WebsiteCrawlError(str(e))28        return result, 20029 30 31class WebsiteCrawlStatusApi(Resource):32    @setup_required33    @login_required34    @account_initialization_required35    def get(self, job_id: str):36        parser = reqparse.RequestParser()37        parser.add_argument("provider", type=str, choices=["firecrawl", "jinareader"], required=True, location="args")38        args = parser.parse_args()39        # get crawl status40        try:41            result = WebsiteService.get_crawl_status(job_id, args["provider"])42        except Exception as e:43            raise WebsiteCrawlError(str(e))44        return result, 20045 46 47api.add_resource(WebsiteCrawlApi, "/website/crawl")48api.add_resource(WebsiteCrawlStatusApi, "/website/crawl/status/<string:job_id>")49