model_script/client/AIGCClient.py

import json
from typing import Tuple

import requests


class AIGCClient(object):
    def __init__(self, token, base_url):
        self.token = token
        self.base_url = base_url

    def update_produce_plan_status(self, produce_plan_id: str, status: int):
        params = {
            "id": produce_plan_id,
            "status": status,
        }
        url = f"{self.base_url}/aigc/produce/plan/updatePlanStatus"
        self.post(url, params)

    def get_content_crawler_plan_by_id(self, crawler_plan_id: str) -> Tuple[str, dict]:
        params = {
            "filterItems": [
                {
                    "itemName": "id",
                    "selectValues": [
                        crawler_plan_id
                    ]
                }
            ],
            "pageNum": 1,
            "pageSize": 100,
            "contentModal": 4
        }
        resp = self.get_content_crawler_plan_list(params=params)
        if resp['code'] == 0 and resp['data']['data']:
            return "", resp['data']['data'][0]
        else:
            return resp['msg'], {}

    def get_content_crawler_plan_list(self, params: dict) -> dict:
        url = f"{self.base_url}/aigc/crawler/plan/list/contentCrawlerPlan"
        return self.post(url, params)

    def crawler_plan_save(self, params: dict):
        url = f"{self.base_url}/aigc/crawler/plan/save"
        return self.post(url, params)

    def post(self, url: str, params: dict) -> dict:
        request_param = {
            "params": params,
            "baseInfo": {
                "token": self.token,
            }
        }
        print(f"invoke aigc platform url: {url} request: {json.dumps(request_param)}")
        response = requests.post(url, json=request_param)
        print(f"invoke aigc platform url: {url} request: {json.dumps(request_param)} response: {response.json()}")
        return response.json()