Merge pull request #1027 from Pythagora-io/virtual-ui

add support for virtual UI, useful for automated end-to-end tests

Author: LeonOstrez

core/cli/helpers.py

@@ -17,6 +17,7 @@
 from core.ui.base import UIBase
 from core.ui.console import PlainConsoleUI
 from core.ui.ipc_client import IPCClientUI
+from core.ui.virtual import VirtualUI
 
 
 def parse_llm_endpoint(value: str) -> Optional[tuple[LLMProvider, str]]:
@@ -313,6 +314,8 @@ def init() -> tuple[UIBase, SessionManager, Namespace]:
 
     if config.ui.type == UIAdapter.IPC_CLIENT:
         ui = IPCClientUI(config.ui)
+    elif config.ui.type == UIAdapter.VIRTUAL:
+        ui = VirtualUI(config.ui.inputs)
     else:
         ui = PlainConsoleUI()
 

core/config/__init__.py

@@ -69,6 +69,7 @@ class UIAdapter(str, Enum):
 
     PLAIN = "plain"
     IPC_CLIENT = "ipc-client"
+    VIRTUAL = "virtual"
 
 
 class ProviderConfig(_StrictModel):
@@ -254,8 +255,17 @@ class LocalIPCConfig(_StrictModel):
     port: int = 8125
 
 
+class VirtualUIConfig(_StrictModel):
+    """
+    Configuration for the virtual UI.
+    """
+
+    type: Literal[UIAdapter.VIRTUAL] = UIAdapter.VIRTUAL
+    inputs: list[Any]
+
+
 UIConfig = Annotated[
-    Union[PlainUIConfig, LocalIPCConfig],
+    Union[PlainUIConfig, LocalIPCConfig, VirtualUIConfig],
     Field(discriminator="type"),
 ]
 

core/ui/virtual.py

@@ -0,0 +1,123 @@
+from typing import Optional
+
+from core.log import get_logger
+from core.ui.base import ProjectStage, UIBase, UISource, UserInput
+
+log = get_logger(__name__)
+
+
+class VirtualUI(UIBase):
+    """
+    Testing UI adapter.
+    """
+
+    def __init__(self, inputs: list[dict[str, str]]):
+        self.virtual_inputs = [UserInput(**input) for input in inputs]
+
+    async def start(self) -> bool:
+        log.debug("Starting test UI")
+        return True
+
+    async def stop(self):
+        log.debug("Stopping test UI")
+
+    async def send_stream_chunk(self, chunk: Optional[str], *, source: Optional[UISource] = None):
+        if chunk is None:
+            # end of stream
+            print("", flush=True)
+        else:
+            print(chunk, end="", flush=True)
+
+    async def send_message(self, message: str, *, source: Optional[UISource] = None):
+        if source:
+            print(f"[{source}] {message}")
+        else:
+            print(message)
+
+    async def send_key_expired(self, message: Optional[str]):
+        pass
+
+    async def ask_question(
+        self,
+        question: str,
+        *,
+        buttons: Optional[dict[str, str]] = None,
+        default: Optional[str] = None,
+        buttons_only: bool = False,
+        allow_empty: bool = False,
+        hint: Optional[str] = None,
+        initial_text: Optional[str] = None,
+        source: Optional[UISource] = None,
+    ) -> UserInput:
+        if source:
+            print(f"[{source}] {question}")
+        else:
+            print(f"{question}")
+
+        if self.virtual_inputs:
+            ret = self.virtual_inputs[0]
+            self.virtual_inputs = self.virtual_inputs[1:]
+            return ret
+
+        if "continue" in buttons:
+            return UserInput(button="continue", text=None)
+        elif default:
+            if buttons:
+                return UserInput(button=default, text=None)
+            else:
+                return UserInput(text=default)
+        elif buttons_only:
+            return UserInput(button=list(buttons.keys)[0])
+        else:
+            return UserInput(text="")
+
+    async def send_project_stage(self, stage: ProjectStage):
+        pass
+
+    async def send_task_progress(
+        self,
+        index: int,
+        n_tasks: int,
+        description: str,
+        source: str,
+        status: str,
+        source_index: int = 1,
+        tasks: list[dict] = None,
+    ):
+        pass
+
+    async def send_step_progress(
+        self,
+        index: int,
+        n_steps: int,
+        step: dict,
+        task_source: str,
+    ):
+        pass
+
+    async def send_run_command(self, run_command: str):
+        pass
+
+    async def open_editor(self, file: str, line: Optional[int] = None):
+        pass
+
+    async def send_project_root(self, path: str):
+        pass
+
+    async def send_project_stats(self, stats: dict):
+        pass
+
+    async def loading_finished(self):
+        pass
+
+    async def send_project_description(self, description: str):
+        pass
+
+    async def send_features_list(self, features: list[str]):
+        pass
+
+    async def import_project(self, project_dir: str):
+        pass
+
+
+__all__ = ["VirtualUI"]