dify/api/core/rag/entities/event.py

from collections.abc import Mapping
from enum import Enum
from typing import Any, Optional

from pydantic import BaseModel, Field


class DatasourceStreamEvent(Enum):
    """
    Datasource Stream event
    """

    PROCESSING = "datasource_processing"
    COMPLETED = "datasource_completed"
    ERROR = "datasource_error"


class BaseDatasourceEvent(BaseModel):
    pass


class DatasourceErrorEvent(BaseDatasourceEvent):
    event: str = DatasourceStreamEvent.ERROR.value
    error: str = Field(..., description="error message")


class DatasourceCompletedEvent(BaseDatasourceEvent):
    event: str = DatasourceStreamEvent.COMPLETED.value
    data: Mapping[str, Any] | list = Field(..., description="result")
    total: Optional[int] = Field(default=0, description="total")
    completed: Optional[int] = Field(default=0, description="completed")
    time_consuming: Optional[float] = Field(default=0.0, description="time consuming")


class DatasourceProcessingEvent(BaseDatasourceEvent):
    event: str = DatasourceStreamEvent.PROCESSING.value
    total: Optional[int] = Field(..., description="total")
    completed: Optional[int] = Field(..., description="completed")