fix type annotations

vemonet · Sep 14, 2023 · 1e34600 · 1e34600
1 parent 8474c6a
commit 1e34600
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/src/libre_chat/ui.py b/src/libre_chat/ui.py
@@ -49,7 +49,7 @@ def chat_accordion() -> Tuple[float, int]:
 
     def stream(
         input_text: str, memory: Any, instructions: str, temperature: float, max_new_tokens: int
-    ) -> Iterator[Tuple[Any, str]]:
+    ) -> "Iterator[Tuple[Any, str]]":
         # Create a Queue
         q: "Queue[Any]" = Queue()
         job_done = object()
@@ -159,7 +159,7 @@ def run_chat(
             instructions: str,
             temperature: float,
             max_new_tokens: int,
-        ) -> Iterator[List[List[str]]]:
+        ) -> "Iterator[List[List[str]]]":
             if not message or (message == RETRY_COMMAND and len(chat_history) == 0):
                 yield chat_history
                 return
@@ -193,7 +193,7 @@ def run_retry(
             instructions: str,
             temperature: float,
             max_new_tokens: int,
-        ) -> Iterator[List[List[str]]]:
+        ) -> "Iterator[List[List[str]]]":
             yield from run_chat(
                 RETRY_COMMAND, chat_history, instructions, temperature, max_new_tokens
             )

diff --git a/tests/test_api.py b/tests/test_api.py
@@ -48,7 +48,13 @@ def test_websocket_prompt_conversation() -> None:
                 pass  # wait for the end
 
 
-def test_get_ui() -> None:
-    """Test get UI"""
+def test_get_gradio_ui() -> None:
+    """Test get gradio UI"""
     response = client.get("/")
     assert response.status_code == 200
+
+
+def test_get_alternative_ui() -> None:
+    """Test get alternative UI"""
+    response = client.get("/ui")
+    assert response.status_code == 200