ExtendLibraryFlow.py

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

from flow_modules.aiflows.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.

    Workflow of ExtendLibrary:
    0. Coder calls ExtendLibrary with a goal.
    1. MemoryReading reads logs, plan, and code library.
    2. Planner makes a plan based on the goal.
    3. Write code in an interactive fashion.
    4. Test code.
    5. Finish, writes code to the library.
    Step 3-5 is done via prompting the controller.

    ExtendLibraryFlow inherits from AbstractBossFlow, visit https://huggingface.co/aiflows/AbstractBossFlowModule/

    *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).

    *Configuration Parameters*:
    - `memory_files`: The memory files to be read by the flow. By convention, the memory files should contain plan, logs, and code library.
    - `input_interface`: The input interface of the flow. By default, they should be `goal`.
    - `output_interface`: The output interface of the flow. By default, they should be `result` and `summary`.
    - `subflows_config`: Configurations w/ the subflows.

    """
    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """Run the flow.
        :param input_data: The input data of the flow.
        :type input_data: Dict[str, Any]
        :return: The output data of the flow.
        :rtype: 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