Generalize tasks through their inputs

actions.py

@@ -2,7 +2,7 @@ import re
 
 import gradio as gr
 
-from components import AITask, all_inputs, all_tasks, Input, MAX_INPUTS, MAX_TASKS
+from components import all_inputs, all_tasks, Input, MAX_INPUTS, MAX_TASKS, Task
 
 
 def add_input(*visibility):
@@ -59,19 +59,32 @@ def _clear_error():
     return gr.HighlightedText.update(value=None, visible=False)
 
 
-def execute_task(id_: int, prompt: str, prev_error_value, *vars):
-    inputs = vars[:MAX_INPUTS]
-    task_outputs = vars[MAX_INPUTS:]
-
-    prompt_vars = set(re.findall("{(.*?)}", prompt))
-    vars_in_scope = {
-        f"{Input.vname}{i}": input_ for i, input_ in enumerate(inputs) if input_
+def execute_task(id_: int, prev_error_value, n_inputs, *vars_in_scope):
+    """
+    Params:
+        - id_: This will tell us which task to execute.
+        - prev_error_value: I carry around whether there is an error in the execution, to be displayed at the end.
+        - n_inputs: How many inputs does this task have?
+        - vars: All variables in scope. This can be a) task inputs, input varaibles or previous task outputs.
+    """
+    n_inputs = int(n_inputs)
+    task_inputs = vars_in_scope[:n_inputs]
+    input_vars = vars_in_scope[n_inputs:MAX_INPUTS]
+    task_outputs = vars_in_scope[MAX_INPUTS:]
+    non_empty_task_inputs = [ti for ti in task_inputs if ti]
+
+    # Put all defined variables into a dict, with names (except task inputs)
+    vars = {
+        f"{Input.vname}{i}": input_ for i, input_ in enumerate(input_vars) if input_
     }
-    vars_in_scope.update(
-        {f"{AITask.vname}{i}": task for i, task in enumerate(task_outputs)}
+    vars.update(
+        {f"{Task.vname}{i}": task for i, task in enumerate(task_outputs)}
     )
-    undefined_vars = prompt_vars - vars_in_scope.keys()
+    # Get all variables referenced within the task inputs
+    prompt_vars = {v for ti in non_empty_task_inputs for v in re.findall("{(.*?)}", ti)}
 
+    # If there is an undefined variable referenced, HighlightedText will signal the error.
+    undefined_vars = prompt_vars - vars.keys()
     if len(undefined_vars) > 0:
         return None, gr.HighlightedText.update(
             value=[
@@ -86,7 +99,12 @@ def execute_task(id_: int, prompt: str, prev_error_value, *vars):
         value=prev_error_value, visible=prev_error_value is not None
     )
 
-    if prompt:
-        return all_tasks[id_].execute(prompt, vars_in_scope), error_update
-
-    return None, error_update
+    if non_empty_task_inputs:
+        # Execute the task logic
+        return (
+            all_tasks[id_].execute(*non_empty_task_inputs, vars),
+            error_update,
+        )
+    else:
+        # There is no actionf or this task.
+        return None, error_update

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.prompt, error_message] + a._get_all_vars_up_to(i),  # type: ignore
+            inputs=[task.component_id, error_message, task.n_inputs] + task.inputs() + a._get_all_vars_up_to(i),  # type: ignore
             outputs=[task.output, error_message],
         )
 

components.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import gradio as gr
 
@@ -23,11 +23,11 @@ class Component(ABC):
 
     @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)
@@ -35,7 +35,7 @@ class Component(ABC):
         self.gr_component = self._render(self._id, self._initial_visibility)
 
     def execute(self, *args):
-        print(f"Executing component :: {self._source}.{self._id}")
+        print(f"Executing {self._source} :: {self._id}")
         return self._execute(*args)
 
 
@@ -51,13 +51,31 @@ class Input(Component):
         )
         return self.output
 
-    def _execute(self):
+    def _execute(self) -> None:
         pass
 
 
-class AITask(Component):
+class Task(Component, ABC):
     vname = "t"
 
+    def __init__(self, id_: int, visible: bool = False):
+        super().__init__(id_, visible)
+        self.output: gr.Textbox
+
+    @abstractmethod
+    def inputs(self) -> List:
+        ...
+
+    @property
+    def _n_inputs(self) -> int:
+        return len(self.inputs())
+
+    def render(self) -> None:
+        super().render()
+        self.n_inputs = gr.Number(value=self._n_inputs, visible=False)
+
+
+class AITask(Task):
     def _render(self, id_: int, visible: bool) -> gr.Box:
         with gr.Box(visible=visible) as gr_component:
             gr.Markdown(f"AI task")
@@ -80,6 +98,9 @@ class AITask(Component):
             formatted_prompt = prompt.format(**prompt_vars)
             return ai.llm.next([{"role": "user", "content": formatted_prompt}])
 
+    def inputs(self) -> List[gr.Textbox]:
+        return [self.prompt]
+
 
 MAX_INPUTS = 10
 MAX_TASKS = 10