from typing import Dict, Any
from aiflows.utils import logging
log = logging.get_logger(__name__)

from flow_modules.Tachi67.AbstractBossFlowModule import AbstractBossFlow
class ExtendLibraryFlow(AbstractBossFlow):
    """ExtendLibraryFlow is one branch executor of CoderFlow. At a higher level, it takes the goal
    from the Coder, writes functions in an interactive way, test the code and append the newly written function to
    the code library.

    *Input Interface (expected input)*
    - `goal` (str): The goal from the caller (source flow, i.e. CoderFlow)
    *Output Interface (expected output)*
    - `result` (str): The result of the flow, the result will be returned to the caller (i.e. CoderFlow).
    - `summary` (str): The summary of the flow, the summary will be logged into the logs of the caller flow (i.e. CoderFlow).
    """
    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        # ~~~ sets the input_data in the flow_state dict ~~~
        self._state_update_dict(update_data=input_data)

        # ~~~ set the memory file to the flow state ~~~
        self._state_update_dict(update_data={"memory_files": self.memory_files})

        max_rounds = self.flow_config.get("max_rounds", 1)
        if max_rounds is None:
            log.info(f"Running {self.flow_config['name']} without `max_rounds` until the early exit condition is met.")

        self._sequential_run(max_rounds=max_rounds)

        output = self._get_output_from_state()

        self.reset(full_reset=True, recursive=True, src_flow=self)

        return output