LongArticleTaskServer/app/infra/internal/aigc_decode_server.py

from typing import Dict, List

from app.infra.shared import AsyncHttpClient


class AigcDecodeServer:
    base_url: str = "https://aigc-api.aiddit.com"

    @staticmethod
    def _sanitize_images(posts: List[Dict]) -> List[Dict]:
        for post in posts:
            images = post.get("images")
            if images:
                post["images"] = [
                    img for img in images if img and img.startswith(("http://", "https://"))
                ]
        return posts

    async def submit_decode(
        self, config_id: int, posts: List[Dict], skip_completed: bool = False
    ) -> Dict:
        """批量提交帖子解构
        POST /aigc/api/task/decode
        """
        url = f"{self.base_url}/aigc/api/task/decode"
        headers = {"Content-Type": "application/json"}
        payload = {
            "params": {
                "configId": config_id,
                "skipCompleted": skip_completed,
                "posts": self._sanitize_images(posts),
            }
        }
        async with AsyncHttpClient() as client:
            return await client.post(url, json=payload, headers=headers)

    async def query_decode_results(
        self, config_id: int, channel_content_ids: List[str]
    ) -> Dict:
        """批量查询解构结果
        POST /aigc/api/task/decode/result
        """
        url = f"{self.base_url}/aigc/api/task/decode/result"
        headers = {"Content-Type": "application/json"}
        payload = {
            "params": {"configId": config_id, "channelContentIds": channel_content_ids}
        }
        async with AsyncHttpClient() as client:
            return await client.post(url, json=payload, headers=headers)

    async def cancel_decode_tasks(
        self, config_id: int, channel_content_ids: List[str]
    ) -> Dict:
        """取消待执行解构任务
        POST /aigc/api/task/decode/cancel
        """
        url = f"{self.base_url}/aigc/api/task/decode/cancel"
        headers = {"Content-Type": "application/json"}
        payload = {
            "params": {"configId": config_id, "channelContentIds": channel_content_ids}
        }
        async with AsyncHttpClient() as client:
            return await client.post(url, json=payload, headers=headers)


__all__ = ["AigcDecodeServer"]