ai-flow-python/app/repository/record_interface.py

import uuid
import logging
import time
from abc import ABC, abstractmethod
from typing import Any, Optional

from app.repository.task_record import TaskRecord

logger = logging.getLogger(__name__)


class RecordInterface(ABC):
    """任务记录存储的抽象接口"""

    def create_record(self, request: dict[str, Any]) -> str:
        record_id = str(uuid.uuid4())
        record = TaskRecord(
            task_id=record_id,
            status="submitted",
            message="",
            request=request,
        )
        self._on_save(record)
        logger.info("[%s] Record created, status=submitted", record_id)
        return record_id

    def set_processing(self, task_id: str):
        logger.info("Status -> processing")
        self._on_update_status(task_id, "processing", "", None)

    def set_failed(self, task_id: str, message: str):
        logger.warning("Status -> failed, message=%s", message)
        self._on_update_status(task_id, "failed", message or "", None)
        self._on_publish(task_id)

    def set_success(self, task_id: str, outputs: dict[str, Any]):
        logger.info("Status -> successes, output_keys=%s", list(outputs.keys()) if outputs else [])
        self._on_update_status(task_id, "successes", "", outputs)
        self._on_publish(task_id)

    def get_record(self, task_id: str) -> Optional[TaskRecord]:
        return self._on_get(task_id)

    def get_all_records(self) -> dict[str, TaskRecord]:
        return self._on_get_all()

    def delete_records(self, max_age_hours: int):
        self._on_delete(max_age_hours)

    @abstractmethod
    def _on_save(self, record: TaskRecord): ...

    @abstractmethod
    def _on_get(self, task_id: str) -> Optional[TaskRecord]: ...

    @abstractmethod
    def _on_get_all(self) -> dict[str, TaskRecord]: ...

    @abstractmethod
    def _on_delete(self, max_age_hours: int): ...

    @abstractmethod
    def _on_update_status(
        self, task_id: str, status: str, message: str, data: Optional[dict]
    ): ...

    def _on_publish(self, task_id: str):
        """任务到达终态时发布事件，子类可覆盖实现"""