LongArticleTaskServer/app/domains/llm_tasks/decode_material/fetch_decode_results.py

import json
from typing import List, Dict

from app.core.database import DatabaseManager
from app.core.observability import LogService

from app.infra.shared import run_tasks_with_asyncio_task_group

from ._const import DecodeMaterialConst
from ._mapper import MaterialDecodeTaskMapper
from ._utils import MaterialDecodeUtils


class FetchMaterialDecodeResults(DecodeMaterialConst):
    def __init__(self, pool: DatabaseManager, log_service: LogService):
        self.pool = pool
        self.log_service = log_service
        self.mapper = MaterialDecodeTaskMapper(self.pool)
        self.tool = MaterialDecodeUtils()

    async def _process_batch(self, tasks: List[Dict]):
        source_ids = [t["source_id"] for t in tasks]
        results = await self.tool.query_decode_results_batch(source_ids)

        for task in tasks:
            source_id = task["source_id"]
            result = results.get(source_id)

            if not result:
                await self.mapper.update_task_status_by_source_id(
                    source_id=source_id,
                    new_status=self.TaskStatus.FAILED,
                    remark="素材解构任务在结果查询中未返回",
                )
                await self.log_service.log(
                    contents={
                        "task": "fetch_material_decode_results",
                        "source_id": source_id,
                        "status": "fail",
                        "message": "source_id not in query response",
                    }
                )
                continue

            status = result.get("status")
            if status == "API_ERROR":
                continue
            elif status == self.QueryStatus.SUCCESS:
                data_content = result.get("dataContent") or "{}"
                html = result.get("html")
                await self.mapper.set_decode_result(
                    source_id=source_id,
                    result=json.dumps(
                        {"dataContent": data_content, "html": html},
                        ensure_ascii=False,
                    ),
                    remark="素材解构结果获取成功",
                )
            elif status in (self.QueryStatus.PENDING, self.QueryStatus.RUNNING):
                pass
            elif status == self.QueryStatus.FAILED:
                await self.mapper.update_task_status_by_source_id(
                    source_id=source_id,
                    new_status=self.TaskStatus.FAILED,
                    remark=f"素材解构任务失败: {result.get('errorMessage', '')}",
                )
            else:
                await self.log_service.log(
                    contents={
                        "task": "fetch_material_decode_results",
                        "source_id": source_id,
                        "status": "unknown",
                        "message": f"unexpected query status: {status}",
                        "data": result,
                    }
                )

    async def deal(self):
        pending_tasks = await self.mapper.fetch_pending_tasks()
        if not pending_tasks:
            await self.log_service.log(
                contents={
                    "task": "fetch_material_decode_results",
                    "message": "No more material tasks to fetch",
                }
            )
            return

        batches = [
            pending_tasks[i : i + self.SUBMIT_BATCH]
            for i in range(0, len(pending_tasks), self.SUBMIT_BATCH)
        ]
        await run_tasks_with_asyncio_task_group(
            task_list=batches,
            handler=self._process_batch,
            description="批量查询素材解构结果",
            unit="batch",
        )

        await self.log_service.log(
            contents={
                "task": "fetch_material_decode_results",
                "message": f"Processed {len(pending_tasks)} pending material tasks in {len(batches)} batches",
            }
        )


__all__ = ["FetchMaterialDecodeResults"]