use union

mbektas · Jan 5, 2025 · d8457b5 · d8457b5
1 parent 8ea2c16
commit d8457b5
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/notebook_intelligence/extension.py b/notebook_intelligence/extension.py
@@ -5,7 +5,7 @@
 from os import path
 import os
 import sys
-from typing import Callable, Dict
+from typing import Callable, Dict, Union
 from dataclasses import dataclass
 from enum import Enum
 
@@ -239,7 +239,7 @@ def description(self) -> str:
     def schema(self) -> dict:
         raise NotImplemented
 
-    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> ToolPreInvokeResponse | None:
+    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> Union[ToolPreInvokeResponse, None]:
         return None
 
     async def handle_tool_call(self, request: ChatRequest, response: ChatResponse, tool_context: dict, tool_args: dict) -> dict:

diff --git a/notebook_intelligence/handlers.py b/notebook_intelligence/handlers.py
@@ -5,6 +5,7 @@
 from os import path
 import os
 import sys
+from typing import Union
 import uuid
 
 from jupyter_server.base.handlers import APIHandler
@@ -125,7 +126,7 @@ def chat_id(self) -> str:
     def message_id(self) -> str:
         return self.messageId
 
-    def stream(self, data: ResponseStreamData | dict):
+    def stream(self, data: Union[ResponseStreamData | dict]):
         data_type = ResponseStreamDataType.LLMRaw if type(data) is dict else data.data_type
 
         if data_type == ResponseStreamDataType.Markdown:

diff --git a/notebook_intelligence/test_extension.py b/notebook_intelligence/test_extension.py
@@ -1,4 +1,5 @@
 from time import sleep
+from typing import Union
 
 from .extension import AnchorData, ButtonData, ChatCommand, CompletionContext, CompletionContextProvider, ContextItem, HTMLData, MarkdownData, NotebookIntelligenceExtension, Host, ChatParticipant, ChatRequest, ChatResponse, ProgressData, ResponseStreamDataType, Tool, ToolPreInvokeResponse, ContextRequest
 
@@ -89,7 +90,7 @@ def schema(self) -> dict:
             },
         }
 
-    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> ToolPreInvokeResponse | None:
+    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> Union[ToolPreInvokeResponse, None]:
         return ToolPreInvokeResponse(
             message="Converting fahrenheit to celcius",
             confirmationTitle="Confirm conversion",
@@ -138,7 +139,7 @@ def schema(self) -> dict:
             },
         }
 
-    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> ToolPreInvokeResponse | None:
+    def pre_invoke(self, request: ChatRequest, tool_args: dict) -> Union[ToolPreInvokeResponse, None]:
         return ToolPreInvokeResponse(
             message="Converting celcius to kelvin",
             # confirmationTitle="Confirm",