fix unittest

l4b4r4b4b4 · Aug 8, 2024 · fe924b9 · fe924b9
1 parent 403293c
commit fe924b9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/functionary/prompt_template/llama31_prompt_template.py b/functionary/prompt_template/llama31_prompt_template.py
@@ -256,10 +256,10 @@ def parse_assistant_response(
 
     def initialize_fsm_gen_state(
         self,
-        tool_choice: str | Tool,
+        tool_choice: Union[str, Tool],
         curr_text: str,
-        curr_tokens: List[int] | None,
-        add_code_interpreter: bool | None,
+        curr_tokens: Optional[List[int]],
+        add_code_interpreter: Optional[bool],
     ) -> Dict:
         func_name = None
         # To force a text response ("tool_choice"="none")
@@ -305,10 +305,10 @@ def stream_delta_text(
         self,
         gen_state: Dict,
         delta_text: str,
-        finish_reason: str | None,
+        finish_reason: Optional[str],
         tools_or_functions: List[Dict],
         tool_choice: Any,
-    ) -> Tuple[Dict | List[Dict] | None]:
+    ) -> Tuple[Dict, Optional[Union[Dict, List[Dict]]]]:
         if finish_reason is not None:  # handle if finish
             if gen_state["stage"] in ["parameter", "code-interpreter"]:
                 finish_reason = "tool_calls"

diff --git a/functionary/prompt_template/prompt_template_v2.py b/functionary/prompt_template/prompt_template_v2.py
@@ -292,10 +292,10 @@ def get_chat_template_jinja(self) -> str:
 
     def initialize_fsm_gen_state(
         self,
-        tool_choice: str | Tool,
+        tool_choice: Union[str, Tool],
         curr_text: str,
-        curr_tokens: List[int] | None,
-        add_code_interpreter: bool | None,
+        curr_tokens: Optional[List[int]],
+        add_code_interpreter: Optional[bool],
     ) -> Dict:
         add_all_recipient = False
         func_name = None
@@ -336,10 +336,10 @@ def stream_delta_text(
         self,
         gen_state: Dict,
         delta_text: str,
-        finish_reason: str | None,
+        finish_reason: Optional[str],
         tools_or_functions: List[Dict],
         tool_choice: Any,
-    ) -> Tuple[Dict | List[Dict] | None]:
+    ) -> Tuple[Dict, Optional[Union[Dict, List[Dict]]]]:
         if finish_reason is not None:  # handle if finish
             if gen_state["func_name"] is not None and gen_state["func_name"] != "all":
                 finish_reason = "tool_calls"