From ac68bad8d49af73e7e2543fef012498a282e9784 Mon Sep 17 00:00:00 2001
From: Nicholas Albion <nalbion@yahoo.com>
Date: Thu, 5 Oct 2023 23:50:32 +1100
Subject: [PATCH] fix for python3.9

---
 pilot/helpers/cli.py   | 4 ++--
 pilot/test_main_e2e.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/pilot/helpers/cli.py b/pilot/helpers/cli.py
index dcba947..c2b1003 100644
--- a/pilot/helpers/cli.py
+++ b/pilot/helpers/cli.py
@@ -5,7 +5,7 @@ import threading
 import queue
 import time
 import platform
-from typing import Dict
+from typing import Dict, Union
 
 from logger.logger import logger
 from utils.style import yellow, green, red, yellow_bold, white_bold
@@ -322,7 +322,7 @@ def execute_command_and_check_cli_response(command, timeout, convo):
 
 
 def run_command_until_success(convo, command,
-                              timeout: int | None,
+                              timeout: Union[int, None],
                               process_name: str | None = None,
                               additional_message=None,
                               force=False,
diff --git a/pilot/test_main_e2e.py b/pilot/test_main_e2e.py
index 10086e6..a56dc62 100644
--- a/pilot/test_main_e2e.py
+++ b/pilot/test_main_e2e.py
@@ -26,9 +26,9 @@ def test_init():
 @pytest.mark.uses_tokens
 @pytest.mark.skip(reason="Uses lots of tokens")
 @pytest.mark.parametrize("endpoint, model", [
-    # ("OPENAI", "gpt-4"),
-    # ("OPENROUTER", "openai/gpt-3.5-turbo"),
-    # ("OPENROUTER", "meta-llama/codellama-34b-instruct"),
+    ("OPENAI", "gpt-4"),
+    ("OPENROUTER", "openai/gpt-3.5-turbo"),
+    ("OPENROUTER", "meta-llama/codellama-34b-instruct"),
     ("OPENROUTER", "google/palm-2-chat-bison"),
     ("OPENROUTER", "google/palm-2-codechat-bison"),
     # TODO: See https://github.com/1rgs/jsonformer-claude/blob/main/jsonformer_claude/main.py