Refactor into class hierarchy

app.py

@@ -64,7 +64,7 @@ with gr.Blocks() as demo:
     for i, task in all_tasks.items():
         execution_event = execution_event.then(
             a.execute_task,
-            inputs=[task.component_id, task.internal.prompt, error_message] + a._get_all_vars_up_to(i),  # type: ignore
+            inputs=[task.component_id, task.prompt, error_message] + a._get_all_vars_up_to(i),  # type: ignore
             outputs=[task.output, error_message],
         )
 

components.py

@@ -1,14 +1,48 @@
-from typing import Dict, Optional, Union
+from abc import ABC, abstractmethod
+from typing import Dict, Optional
 
 import gradio as gr
 
 import ai
 
 
-class Input:
+class Component(ABC):
+    vname = None
+
+    def __init__(self, id_: int, visible: bool = False):
+        # Internal state
+        self._id = id_
+        self._source = self.__class__.__name__
+        self._initial_visibility = visible
+
+        # Gradio state
+        self.component_id: gr.Number
+        self.visible: gr.Number
+        self.gr_component = gr.Box
+        self.output: gr.Textbox
+
+    @abstractmethod
+    def _render(self, id_: int, visible: bool):
+        pass
+
+    @abstractmethod
+    def _execute(self):
+        pass
+
+    def render(self) -> None:
+        self.component_id = gr.Number(value=self._id, visible=False)
+        self.visible = gr.Number(int(self._initial_visibility), visible=False)
+        self.gr_component = self._render(self._id, self._initial_visibility)
+
+    def execute(self, *args):
+        print(f"Executing component :: {self._source}.{self._id}")
+        return self._execute(*args)
+
+
+class Input(Component):
     vname = "v"
 
-    def render(self, id_: int, visible: bool) -> gr.Textbox:
+    def _render(self, id_: int, visible: bool) -> gr.Textbox:
         self.output = gr.Textbox(
             label=f"Input: {{{self.vname}{id_}}}",
             interactive=True,
@@ -17,14 +51,14 @@ class Input:
         )
         return self.output
 
-    def execute(self) -> None:
+    def _execute(self):
         pass
 
 
-class AITask:
+class AITask(Component):
     vname = "t"
 
-    def render(self, id_: int, visible: bool) -> gr.Box:
+    def _render(self, id_: int, visible: bool) -> gr.Box:
         with gr.Box(visible=visible) as gr_component:
             gr.Markdown(f"AI task")
             with gr.Row():
@@ -41,43 +75,18 @@ class AITask:
                 )
             return gr_component
 
-    def execute(self, prompt: str, prompt_vars: Dict[str, str]) -> Optional[str]:
+    def _execute(self, prompt: str, prompt_vars: Dict[str, str]) -> Optional[str]:
         if prompt:
             formatted_prompt = prompt.format(**prompt_vars)
             return ai.llm.next([{"role": "user", "content": formatted_prompt}])
 
 
-class Component:
-    def __init__(self, id_: int, internal: Union[Input, AITask], visible: bool = False):
-        # Internal state
-        self._id = id_
-        self.internal = internal
-        self._source = self.internal.__class__.__name__
-        self._initial_visibility = visible
-
-        # Gradio state
-        self.component_id: gr.Number
-        self.visible: gr.Number
-        self.gr_component = gr.Box
-        self.output: gr.Textbox
-
-    def render(self) -> None:
-        self.component_id = gr.Number(value=self._id, visible=False)
-        self.visible = gr.Number(int(self._initial_visibility), visible=False)
-        self.gr_component = self.internal.render(self._id, self._initial_visibility)
-        self.output = self.internal.output
-
-    def execute(self, *args):
-        print(f"Executing component :: {self._source}.{self._id}")
-        return self.internal.execute(*args)
-
-
 MAX_INPUTS = 10
 MAX_TASKS = 10
 
 
-all_inputs = {i: Component(i, Input()) for i in range(MAX_INPUTS)}
-all_tasks = {i: Component(i, AITask()) for i in range(MAX_TASKS)}
+all_inputs = {i: Input(i) for i in range(MAX_INPUTS)}
+all_tasks = {i: AITask(i) for i in range(MAX_TASKS)}
 
 all_inputs[0]._initial_visibility = True
 all_tasks[0]._initial_visibility = True