events

@define
class ActionChunkEvent(BaseChunkEvent):
    partial_input: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
    tag: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
    name: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
    path: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})

    def __str__(self) -> str:
        parts = []

        if self.name:
            parts.append(self.name)
            if self.path:
                parts.append(f".{self.path}")
                if self.tag:
                    parts.append(f" ({self.tag})")

        if self.partial_input:
            if parts:
                parts.append(f"\n{self.partial_input}")
            else:
                parts.append(self.partial_input)

        return "".join(parts)

def __str__(self) -> str:
    parts = []

    if self.name:
        parts.append(self.name)
        if self.path:
            parts.append(f".{self.path}")
            if self.tag:
                parts.append(f" ({self.tag})")

    if self.partial_input:
        if parts:
            parts.append(f"\n{self.partial_input}")
        else:
            parts.append(self.partial_input)

    return "".join(parts)

@define
class AudioChunkEvent(BaseChunkEvent):
    """Stores a chunk of audio data.

    Attributes:
        data: Base64 encoded audio data.
    """

    data: str = field(kw_only=True, metadata={"serializable": True})

    def __str__(self) -> str:
        return self.data

def __str__(self) -> str:
    return self.data

@define
class BaseActionsSubtaskEvent(BaseTaskEvent, ABC):
    subtask_parent_task_id: Optional[str] = field(kw_only=True, metadata={"serializable": True})
    subtask_thought: Optional[str] = field(kw_only=True, metadata={"serializable": True})
    subtask_actions: Optional[list[dict]] = field(kw_only=True, metadata={"serializable": True})

    @classmethod
    def from_task(cls, task: BaseTask) -> BaseActionsSubtaskEvent:
        from griptape.tasks import ActionsSubtask

        if not isinstance(task, ActionsSubtask):
            raise ValueError("Event must be of instance ActionSubtask.")
        return cls(
            task_id=task.id,
            task_parent_ids=task.parent_ids,
            task_child_ids=task.child_ids,
            task_input=task.input,
            task_output=task.output,
            subtask_parent_task_id=task.origin_task.id,
            subtask_thought=task.thought,
            subtask_actions=task.actions_to_dicts(),
        )

@classmethod
def from_task(cls, task: BaseTask) -> BaseActionsSubtaskEvent:
    from griptape.tasks import ActionsSubtask

    if not isinstance(task, ActionsSubtask):
        raise ValueError("Event must be of instance ActionSubtask.")
    return cls(
        task_id=task.id,
        task_parent_ids=task.parent_ids,
        task_child_ids=task.child_ids,
        task_input=task.input,
        task_output=task.output,
        subtask_parent_task_id=task.origin_task.id,
        subtask_thought=task.thought,
        subtask_actions=task.actions_to_dicts(),
    )

class BaseAudioTranscriptionEvent(BaseEvent): ...

@define
class BaseChunkEvent(BaseEvent):
    index: int = field(default=0, metadata={"serializable": True})

    @abstractmethod
    def __str__(self) -> str: ...

@abstractmethod
def __str__(self) -> str: ...

@define
class BaseEvent(SerializableMixin, ABC):
    id: str = field(default=Factory(lambda: uuid.uuid4().hex), kw_only=True, metadata={"serializable": True})
    timestamp: float = field(default=Factory(lambda: time.time()), kw_only=True, metadata={"serializable": True})
    meta: dict[str, Any] = field(factory=dict, kw_only=True, metadata={"serializable": True})

@define
class BasePromptEvent(BaseEvent, ABC):
    model: str = field(kw_only=True, metadata={"serializable": True})

@define
class BaseTaskEvent(BaseEvent, ABC):
    task_id: str = field(kw_only=True, metadata={"serializable": True})
    task_parent_ids: list[str] = field(kw_only=True, metadata={"serializable": True})
    task_child_ids: list[str] = field(kw_only=True, metadata={"serializable": True})

    task_input: BaseArtifact = field(kw_only=True, metadata={"serializable": True})
    task_output: Optional[BaseArtifact] = field(kw_only=True, metadata={"serializable": True})

@define
class BaseTextToSpeechEvent(BaseMediaGenerationEvent, ABC): ...

@define
class EventListener(Generic[T]):
    """An event listener that listens for events and handles them.

    Attributes:
        on_event: The on_event function that will be called when an event is published.
            The on_event function should accept an event and return either the event or a dictionary.
            If the on_event returns None, the event will not be published.
        event_types: A list of event types that the event listener should listen for.
            If not provided, the event listener will listen for all event types.
        event_listener_driver: The driver that will be used to publish events.
    """

    on_event: Optional[Callable[[T], Optional[BaseEvent | dict]]] = field(default=None)
    event_types: Optional[list[type[T]]] = field(default=None, kw_only=True)
    event_listener_driver: Optional[BaseEventListenerDriver] = field(default=None, kw_only=True)

    def __enter__(self) -> Self:
        from griptape.events import EventBus

        EventBus.add_event_listener(self)

        return self

    def __exit__(self, type, value, traceback) -> None:  # noqa: ANN001, A002
        from griptape.events import EventBus

        EventBus.remove_event_listener(self)

    def publish_event(self, event: T, *, flush: bool = False) -> None:
        event_types = self.event_types

        if event_types is None or any(isinstance(event, event_type) for event_type in event_types):
            handled_event = event
            if self.on_event is not None:
                handled_event = self.on_event(event)

            if self.event_listener_driver is not None and handled_event is not None:
                self.event_listener_driver.publish_event(handled_event)

        if self.event_listener_driver is not None and flush:
            self.event_listener_driver.flush_events()

def __enter__(self) -> Self:
    from griptape.events import EventBus

    EventBus.add_event_listener(self)

    return self

def __exit__(self, type, value, traceback) -> None:  # noqa: ANN001, A002
    from griptape.events import EventBus

    EventBus.remove_event_listener(self)

def publish_event(self, event: T, *, flush: bool = False) -> None:
    event_types = self.event_types

    if event_types is None or any(isinstance(event, event_type) for event_type in event_types):
        handled_event = event
        if self.on_event is not None:
            handled_event = self.on_event(event)

        if self.event_listener_driver is not None and handled_event is not None:
            self.event_listener_driver.publish_event(handled_event)

    if self.event_listener_driver is not None and flush:
        self.event_listener_driver.flush_events()

@define
class FinishActionsSubtaskEvent(BaseActionsSubtaskEvent): ...

class FinishAudioTranscriptionEvent(BaseAudioTranscriptionEvent): ...

@define
class FinishImageGenerationEvent(BaseImageGenerationEvent): ...

@define
class FinishImageQueryEvent(BaseImageQueryEvent):
    result: str = field(kw_only=True, metadata={"serializable": True})

@define
class FinishPromptEvent(BasePromptEvent):
    result: str = field(kw_only=True, metadata={"serializable": True})
    input_token_count: Optional[float] = field(kw_only=True, metadata={"serializable": True})
    output_token_count: Optional[float] = field(kw_only=True, metadata={"serializable": True})

@define
class FinishStructureRunEvent(BaseEvent):
    structure_id: Optional[str] = field(kw_only=True, default=None, metadata={"serializable": True})
    output_task_input: Optional[BaseArtifact] = field(kw_only=True, default=None, metadata={"serializable": True})
    output_task_output: Optional[BaseArtifact] = field(kw_only=True, default=None, metadata={"serializable": True})

@define
class FinishTaskEvent(BaseTaskEvent): ...

@define
class FinishTextToSpeechEvent(BaseTextToSpeechEvent): ...

@define
class StartActionsSubtaskEvent(BaseActionsSubtaskEvent): ...

class StartAudioTranscriptionEvent(BaseAudioTranscriptionEvent): ...

@define
class StartImageGenerationEvent(BaseImageGenerationEvent):
    prompts: list[str] = field(kw_only=True, metadata={"serializable": True})
    negative_prompts: Optional[list[str]] = field(default=None, kw_only=True, metadata={"serializable": True})

@define
class StartImageQueryEvent(BaseImageQueryEvent):
    query: str = field(kw_only=True, metadata={"serializable": True})
    images_info: list[str] = field(kw_only=True, metadata={"serializable": True})

@define
class StartPromptEvent(BasePromptEvent):
    prompt_stack: PromptStack = field(kw_only=True, metadata={"serializable": True})

@define
class StartStructureRunEvent(BaseEvent):
    structure_id: Optional[str] = field(kw_only=True, default=None, metadata={"serializable": True})
    input_task_input: BaseArtifact = field(kw_only=True, metadata={"serializable": True})
    input_task_output: Optional[BaseArtifact] = field(kw_only=True, metadata={"serializable": True})

@define
class StartTaskEvent(BaseTaskEvent): ...

@define
class StartTextToSpeechEvent(BaseTextToSpeechEvent):
    prompts: list[str] = field(kw_only=True, metadata={"serializable": True})

@define
class TextChunkEvent(BaseChunkEvent):
    token: str = field(kw_only=True, metadata={"serializable": True})

    def __str__(self) -> str:
        return self.token

def __str__(self) -> str:
    return self.token