some changes

Author: kateforsberg13

griptape/drivers/assistant/openai_assistant_driver.py

@@ -14,6 +14,8 @@
 if TYPE_CHECKING:
     import openai
     from openai import AssistantEventHandler
+    from openai.types.beta.threads import Text, TextDelta
+    from openai.types.beta.threads.runs import ToolCall, ToolCallDelta
 
 
 @define
@@ -24,15 +26,13 @@ def _create_event_handler_class() -> type[AssistantEventHandler]:  # pyright: ig
         AssistantEventHandler = import_optional_dependency("openai").AssistantEventHandler  # noqa: N806
 
         class EventHandler(AssistantEventHandler):
-            # Pyright can't verify override since base class is in TYPE_CHECKING
             @override
-            def on_text_delta(self, delta, snapshot) -> None:  # pyright: ignore[reportGeneralTypeIssues,reportUndefinedVariable] # noqa: ANN001
+            def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:  # pyright: ignore[reportGeneralTypeIssues,reportUndefinedVariable] Pyright can't verify override since base class is in TYPE_CHECKING
                 if delta.value is not None:
                     EventBus.publish_event(TextChunkEvent(token=delta.value))
 
-            # Pyright can't verify override since base class is in TYPE_CHECKING
             @override
-            def on_tool_call_delta(self, delta, snapshot) -> None:  # pyright: ignore[reportGeneralTypeIssues,reportUndefinedVariable] # noqa: ANN001
+            def on_tool_call_delta(self, delta: ToolCallDelta, snapshot: ToolCall) -> None:  # pyright: ignore[reportGeneralTypeIssues,reportUndefinedVariable] Pyright can't verify override since base class is in TYPE_CHECKING
                 if delta.type == "code_interpreter" and delta.code_interpreter is not None:
                     if delta.code_interpreter.input:
                         EventBus.publish_event(TextChunkEvent(token=delta.code_interpreter.input))