From 12a05ce99e52c1911ee558db77defd565cadd80b Mon Sep 17 00:00:00 2001
From: Collin Dutter <collin@griptape.ai>
Date: Tue, 15 Oct 2024 12:08:00 -0700
Subject: [PATCH] Use generics for event listener types

---
 CHANGELOG.md                                      |  5 +++--
 .../misc/src/events_no_publish.py                 |  2 +-
 griptape/events/event_listener.py                 | 15 +++++++++------
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d2f1645dd6..9ffbe17cc1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -36,15 +36,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Structures not flushing events when not listening for `FinishStructureRunEvent`.
+- `EventListener.event_types` and the argument to `BaseEventListenerDriver.handler` being out of sync.
 
-## [0.33.1] - 2024-10-11
+## \[0.33.1\] - 2024-10-11
 
 ### Fixed
 
 - Pinned `cohere` at `~5.11.0` to resolve slow dependency resolution.
 - Missing `exa-py` from `all` extra.
 
-## [0.33.0] - 2024-10-09
+## \[0.33.0\] - 2024-10-09
 
 ## Added
 
diff --git a/docs/griptape-framework/misc/src/events_no_publish.py b/docs/griptape-framework/misc/src/events_no_publish.py
index a598659232..77267b9bc4 100644
--- a/docs/griptape-framework/misc/src/events_no_publish.py
+++ b/docs/griptape-framework/misc/src/events_no_publish.py
@@ -32,7 +32,7 @@ def handler_maybe_drop_events(event: FinishStructureRunEvent) -> Optional[BaseEv
 EventBus.add_event_listeners(
     [
         EventListener(
-            handler_maybe_drop_events,  # pyright: ignore[reportArgumentType]
+            handler_maybe_drop_events,
             event_types=[FinishStructureRunEvent],
             # By default, GriptapeCloudEventListenerDriver uses the api key provided
             # in the GT_CLOUD_API_KEY environment variable.
diff --git a/griptape/events/event_listener.py b/griptape/events/event_listener.py
index f77980aa06..dbc27e46f7 100644
--- a/griptape/events/event_listener.py
+++ b/griptape/events/event_listener.py
@@ -1,17 +1,20 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Callable, Generic, Optional, TypeVar
 
 from attrs import define, field
 
+from .base_event import BaseEvent
+
 if TYPE_CHECKING:
     from griptape.drivers import BaseEventListenerDriver
 
-    from .base_event import BaseEvent
+
+T = TypeVar("T", bound=BaseEvent)
 
 
 @define
-class EventListener:
+class EventListener(Generic[T]):
     """An event listener that listens for events and handles them.
 
     Args:
@@ -23,8 +26,8 @@ class EventListener:
         event_listener_driver: The driver that will be used to publish events.
     """
 
-    handler: Optional[Callable[[BaseEvent], Optional[BaseEvent | dict]]] = field(default=None)
-    event_types: Optional[list[type[BaseEvent]]] = field(default=None, kw_only=True)
+    handler: 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)
 
     _last_event_listeners: Optional[list[EventListener]] = field(default=None)
@@ -43,7 +46,7 @@ def __exit__(self, type, value, traceback) -> None:  # noqa: ANN001, A002
 
         self._last_event_listeners = None
 
-    def publish_event(self, event: BaseEvent, *, flush: bool = False) -> None:
+    def publish_event(self, event: T, *, flush: bool = False) -> None:
         event_types = self.event_types
 
         if event_types is None or type(event) in event_types: