From 1e346002d734b9770c83e9f053e08fde3c9c1528 Mon Sep 17 00:00:00 2001 From: Vincent Emonet Date: Thu, 14 Sep 2023 07:59:20 +0200 Subject: [PATCH] fix type annotations --- src/libre_chat/ui.py | 6 +++--- tests/test_api.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/libre_chat/ui.py b/src/libre_chat/ui.py index 616211d..de24c69 100644 --- 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 index d4235ec..ce926bc 100644 --- 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