flow.py

from enum import Enum
from typing import Tuple, Callable
from utils import load_prompt
import os

class Step:
    def __init__(self, key, value):
        self.key = key
        self.value = value
        
    def __str__(self):
        return f"{self.key} : {self.value}"
    
    def __repr__(self):
        return f"{self.__class__.__name__} - ({self.key}:{self.value})"
    
    def __iter__(self):
        return self.key, self.value
    
class ConditionableStep(Step):
    def __init__(self, key, value, condition_judge = None):
        super().__init__(key, value)
        if condition_judge is None:
            condition_judge = lambda : True
        elif isinstance(condition_judge, Callable):
            self.condition_judge = condition_judge
        else:
            raise ValueError("Condition judge is not a valid function")
    
    def is_condition_met(self) -> bool:
        return self.condition_judge()

class PromptStep(Step):
    def __init__(self, key, value, path : str):
        super().__init__(key, value)
        self.prompt = load_prompt(path)
    
    @property
    def prompt(self):
        return self._prompt
    
    @prompt.setter
    def prompt(self, prompt : str):
        self._prompt = prompt
        
class FlowStep(PromptStep, ConditionableStep):
    def __init__(self, key, value, path: str = None, condition_judge: Callable = None):
        PromptStep.__init__(self, key, value, path)
        ConditionableStep.__init__(self, key, value, condition_judge)

class LinearFlow:
    def __init__(self, steps : Tuple[Step] = None):
        self.steps = steps
        if steps is None:
            self.steps = tuple()
        self.current_index = 0
        
    def __iter__(self):
        return self.steps.__iter__()
    
    def __getitem__(self, index):
        if index < len(self.steps):
            return self.steps[index]
        else:
            raise StopIteration("No more steps in the flow")
        
    def __str__(self):
        return f"{self.__class__.__name__} : {self.steps}"

    def add_step(self, step : Step):
        if isinstance(step, Step):
            self.steps += (step,)
        else:
            raise ValueError("Step is not a valid step")
    
    def next(self) -> Step:
        if isinstance(self.current(), FlowStep) or isinstance(self.current(), ConditionableStep):
            condition_met = self.current().is_condition_met()
            result = self.current()
            self.current_index += 1 if condition_met else 0
            return result
        elif self.current_index < len(self.steps):
            step = self.steps[self.current_index]
            self.current_index += 1
            return step
        else:
            raise StopIteration("No more steps in the flow")

    def current(self) -> Step:
        if self.current_index < len(self.steps):
            return self.steps[self.current_index]
        else:
            raise StopIteration("No more steps in the flow")