README

Pre-requisites

Use the branching_flow branch of flows package.

python -m Flow_dev.ReActFlow.ReActFlow to run the test case of ReActFlow.

ReActFlow Schema

Name: ReActFlow
Class name: ReActFlow
Type: circular
Description: ToDO: add description
Input keys: ['question']
Output keys: ['answer', 'status']
Subflows:
        Name: Agent
        Class name: AgentAtomicFlow
        Type: AtomicFlow
        Description: ToDO: add description
        Input keys: ['observation']
        Output keys: ['command', 'argument1', 'argument2']

        Name: Action
        Class name: ActionFlow
        Type: branching
        Description: ToDO: add description
        Input keys: None
        Output keys: ['observation']
        Subflows:
                Name: lookup
                Class name: WikiLookupAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['lookup_term', 'wiki_content']
                Output keys: ['observation']

                Name: search
                Class name: WikiSearchAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['search_term']
                Output keys: ['observation', 'wiki_content']

                Name: finish
                Class name: FinishAtomicFlow
                Type: AtomicFlow
                Description: A flow that outputs the final answer
                Input keys: None
                Output keys: ['observation', 'EARLY_EXIT']

ReAct protocol

As shown in the figure above, the ReAct protocol consists of an iterative loop of the following steps:

• Reason: Agent explains the current situation and what it needs to do next; agent gives a command and arguments in the form of command[argument]
• Act: LLM executes the command and returns an observation

Step 1 ReActFlow as a circular flow

Let's think what the interface of the ReAct Flow should look like.

Name: ReActFlow
Class name: ReActFlow
Type: circular
Input keys: ['question']
Output keys: ['answer', 'status']
Subflows:
    - [TO ADD]
    - [TO ADD]

Under the framework of Flows, the ReAct protocol is implemented as a circular flow. Cicular flow is nothing more than running subflows in a loop.

ReAct Flow takes a question as input and outputs an answer and a status(whether the flow is finished or early exited). The status is optional, but it is useful for debugging.

Step 2 two subflows: Agent and Action

Let's think about what subflows are needed. We need two subflows: Agent and Action. AgentFlow is responsible for reasoning and gives the next command to run. We will implement AgentFlow as a subclass of OpenAIChatAtomicFlow so that we can let GPT to decide what to do next.

ActionFlow is responsible for executing the command and returns an observation. In ReAct, three actions are available: lookup, search, and finish.

At each iteration, ActionFlow will execute one of the three actions depending on the command given by AgentFlow. This is implemented as a branching flow, which is a flow that runs one of the subflows depending on the input.

So now we have the following ReActFlow.

Name: ReActFlow
Class name: ReActFlow
Type: circular
Description: ToDO: add description
Input keys: ['question']
Output keys: ['answer', 'status']
Subflows:
        Name: Agent
        Class name: AgentAtomicFlow
        Type: AtomicFlow
        Description: ToDO: add description
        Input keys: ['observation']
        Output keys: ['command', 'argument1', 'argument2']

        Name: Action
        Class name: ActionFlow
        Type: branching
        Description: ToDO: add description
        Input keys: None
        Output keys: ['observation']
        Subflows:
                Name: lookup
                Class name: WikiLookupAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['lookup_term', 'wiki_content']
                Output keys: ['observation']

                Name: search
                Class name: WikiSearchAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['search_term']
                Output keys: ['observation', 'wiki_content']

                Name: finish
                Class name: FinishAtomicFlow
                Type: AtomicFlow
                Description: A flow that outputs the final answer
                Input keys: None
                Output keys: ['observation', 'EARLY_EXIT']

Name: ReActFlow
Class name: ReActFlow
Type: circular
Description: ToDO: add description
Input keys: ['question']
Output keys: ['answer', 'status']
Subflows:
        Name: Agent
        Class name: AgentAtomicFlow
        Type: AtomicFlow
        Description: ToDO: add description
        Input keys: ['observation']
        Output keys: ['command', 'argument1', 'argument2']

        Name: Action
        Class name: ActionFlow
        Type: branching
        Description: ToDO: add description
        Input keys: None
        Output keys: ['observation']
        Subflows:
            [TO ADD]
            [TO ADD]

The agent flow takes an observation as input and outputs a command and a number of arguments. The exact number of arguments depends on the command(or the api call) that the agent wants to execute. For example, if the agent wants to execute the lookup action, it will output a lookup command and a lookup_term argument. But in general, different commands may require different number of arguments. Hence, we need the output_keys to contains the maximum number of arguments that the agent may output. (This is not an ideal solution, because one needs to know the maximum number of arguments in advance. In the next version, we will use a more flexible solution. User simply leave the output_keys to be None, and the framework will automatically does the job. This requires to change the running logic of OutputDataTransformation)

The action flow takes input=None and outputs an observation. Here the input=None means that the action flow will take all inputs from the previous flow. This allows the action flow to take dynamic number of inputs.

Step 3 AgentFlow

Now we need to implement the AgentFlow, in particular, the prompts and the data transformations.

Here is the system prompt for the AgentFlow.

Your goal is to find the correct answer to the given question. 
You can take one of the three actions:
-  search[entity], which returns the first 5 sentences from the corresponding entity wiki page if it exists, or else suggests top-5 similar entities from the Wikipedia search engine
-  lookup[term], which would return the next sentence in the page containing term.
-  finish[answer], which would finish the current task with answer. 

You need to return the action in the above format, i.e. search[entity], lookup[term] or finish[answer].
For example, if you want to search for the entity "Albert Einstein", you need to return "search[Albert Einstein]".
If you want to lookup the term "relativity", you need to return "lookup[relativity]".
If you want to finish the task with the answer "Einstein's theory of relativity", you need to return "finish[Einstein's theory of relativity]".

With this prompt, GPT4 will know what actions are available and what format it should output.

Here is the user prompt for the AgentFlow. Note that the user prompt is different at the first iteration and the following iterations. The first iteration prompt contains the question, while the following iterations prompt contains the observation from the previous iteration's action

Here is the prompt for the first iteration.

init_human_message_prompt_template:
  _target_: langchain.PromptTemplate
  template: |2-
    Here is the question you need to answer:
    {{question}}

  input_variables:
    - "question"
  template_format: jinja2

Here is the prompt for the following iterations.

human_message_prompt_template:
  _target_: langchain.PromptTemplate
  template: |2-
    Here is the response to your last action:
    {{observation}}

  input_variables:
    - "observation"
  template_format: jinja2

Now we have the prompts, we need to implement the data transformations.

There is not much to do for the input_data_transformations. But we need to implement the output_data_transformations to parse the output from GPT4.

We need to have two RegexFirstOccurrenceExtractor to extract the command and the argument1.

output_data_transformations:
  - _target_: flows.data_transformations.RegexFirstOccurrenceExtractor
    regex: '^(search|lookup|finish)' # this extracts the command, one of search, lookup, finish
    input_key: "api_output"
    output_key: "command"
    strip: True
    assert_unique: True
  - _target_: flows.data_transformations.RegexFirstOccurrenceExtractor
    regex: '\[(.*?)\]' # this extracts the string inside the square brackets
    input_key: "api_output"
    output_key: "argument1"
    strip: True
    assert_unique: True
    match_group: 1
  - _target_: flows.data_transformations.KeyRename
    old_key2new_key:
      observation: "argument2" # this renaming is optional, but it is good to have a consistent naming convention

Step 4 ActionFlow

The ActionFlow is a branching flow, which means that it has multiple subflows. In this example, we have three subflows: lookup, search and finish. The three flows are all atomic flows, which means that they do not have any subflows. We will implement the three atomic flows in the next step.

Here let's focus on the config yaml for the ActionFlow.

The class BranchingFlow requires an input variable branch to decide which subflow to take. Here we need to do a renaming to use the command from the AgentFlow as the branch input variable for the ActionFlow.

input_data_transformations:
  - _target_: flows.data_transformations.KeyRename
    old_key2new_key:
      command: "branch"

Another thing to note is that the class BranchingFlow requires the subflows to be named as the values of the branch input variable. In our case, the branch input variable can be one of lookup, search and finish. Hence, we need to rename the subflows to lookup, search and finish.

subflows_config:
  - _target_: FlowsEpfl.ReActFlow.WikiLookupAtomicFlow.instantiate_from_default_config
    overrides:
      name: "lookup"
  - _target_: FlowsEpfl.ReActFlow.WikiSearchAtomicFlow.instantiate_from_default_config
    overrides:
      name: "search"
  - _target_: FlowsEpfl.ReActFlow.FinishAtomicFlow.instantiate_from_default_config
    overrides:
      name: "finish"

Step 5 AtomicFlows of the ActionFlow

Now we need to implement the three atomic flows: lookup, search and finish.

Step 5.1 SearchAtomicFlow

The SearchAtomicFlow is a simple flow that takes an term as input and search the Wikipedia for the page and return the result. A slight complication is that we need to handle cases where the page does not exist. In this case, we need to return the top-5 similar entities from the Wikipedia search engine.

class WikiSearchAtomicFlow(AtomicFlow):
    REQUIRED_KEYS_CONFIG = ["lang", "top_k_results", "doc_content_chars_max"]
    REQUIRED_KEYS_CONSTRUCTOR = []

    SUPPORTS_CACHING: bool = True

    api_wrapper: WikipediaAPIWrapper

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def run(self,
            input_data: Dict[str, Any]) -> Dict[str, Any]:
        # ~~~ Process input ~~~
        term = input_data.get("search_term", None)
        api_wrapper = WikipediaAPIWrapper(
            lang=self.flow_config["lang"],
            top_k_results=self.flow_config["top_k_results"],
            doc_content_chars_max=self.flow_config["doc_content_chars_max"]
        )

        # ~~~ Call ~~~
        if page_content := api_wrapper._fetch_page(term):
            search_response = {"wiki_content": page_content, "relevant_pages": None}
        else:
            page_titles = api_wrapper.search_page_titles(term)
            search_response = {"wiki_content": None, "relevant_pages": f"Could not find [{term}]. similar: {page_titles}"}
        # Log the update to the flow messages list
        observation = search_response["wiki_content"] if search_response["wiki_content"] else search_response["relevant_pages"]
        return {"observation": observation, "wiki_content": search_response["wiki_content"]}

Note that in the run method, we return the observation and the wiki_content separately. In case the page exists, the two values are the same. This is because the observation is used as the input to the next iteration of the AgentFlow and the wiki_content is used as the input to the WikiLookupAtomicFlow.

Step 5.2 LookupAtomicFlow

The LookupAtomicFlow is a simple flow that takes a term as input and search the term in the wiki_content and we have retrieved from the WikiSearchAtomicFlow. The output is the next sentence containing the term in the wiki_content.

class WikiLookupAtomicFlow(AtomicFlow):

    SUPPORTS_CACHING: bool = True

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def run(self,
            input_data: Dict[str, Any]) -> Dict[str, Any]:
        # ~~~ Process input ~~~
        term = input_data["lookup_term"].lower()# lookup is case insensitive

        wiki_content = input_data.get("wiki_content", None)
        if wiki_content is None:
            raise ValueError("Cannot lookup without a successful search. The wiki_content is None. You must run WikiSearchAtomicFlow first.")

        if wiki_content != self.flow_state["wiki_content"]:
            self.flow_state["wiki_content"] = wiki_content
            self.flow_state["lookup_idx"] = 0

        sentences: List[str] = split_into_sentences(wiki_content)

        lookup_idx = self.flow_state["lookup_idx"]

        lookups = [s for s in sentences if term in s.lower()]
        if len(lookups) == 0:
            return {"lookup_response": "No Results"}
        elif lookup_idx >= len(lookups):
            return {"lookup_response": "No More Results"}
        else:
            result_prefix = f"(Result {lookup_idx + 1}/{len(lookups)})"
            self.flow_state["lookup_idx"] += 1
            return {"lookup_response": f"{result_prefix} {lookups[lookup_idx]}"}

    def set_up_flow_state(self):
        super().set_up_flow_state()
        self.flow_state["wiki_content"] = None
        self.flow_state["lookup_idx"] = 0

One important point is we have overridden the set_up_flow_state method. This is because the WikiLookupAtomicFlow needs to keep track of the wiki_content and the lookup_idx across multiple iterations of the AgentFlow. It should return the next sentence containing the term in the wiki_content, and the lookup_idx should be incremented by 1. If the Agent call the WikiSearchAtomicFlow again with a different page, the wiki_content should be updated and the lookup_idx should be reset to 0.

Step 5.3 FinishAtomicFlow

The FinishAtomicFlow is a placeholder flow that simply rename the argument1 to observation to be consistent with the other actions.

class FinishAtomicFlow(AtomicFlow):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        return input_data

Although this is a placeholder flow, it is important to note that the FinishAtomicFlow has an important role in the AgentFlow. It will trigger the early_exit of the ReAct.

name: FinishAtomicFlow
description: A flow that outputs the final answer

end_of_interaction: "" # empty string means, any input will be considered as end of interaction, including empty string

output_data_transformations:
  - _target_: flows.data_transformations.EndOfInteraction
    end_of_interaction_string: ""
    input_key: "observation"
    output_key: "EARLY_EXIT"


output_keys:
  - "observation"
  - "EARLY_EXIT"

Step 5.4 Data Transformations

Now we have implemented all the atomic flows, we can put them together in the ActionFlow. The only thing left is to define DataTransformations to make sure the input and output of the atomic flows are consistent.

Look in the below how we renamed the argument1 and argument2 to be consistent with each action flow's run method's signature.

subflows_config:
  - _target_: FlowsEpfl.ReActFlow.WikiLookupAtomicFlow.instantiate_from_default_config
    overrides:
      name: "lookup"
      input_data_transformations:
        - _target_: flows.data_transformations.KeyRename
          old_key2new_key:
            argument1: "lookup_term"
        - _target_: flows.data_transformations.KeyRename
          old_key2new_key:
            argument2: "wiki_content"
  - _target_: FlowsEpfl.ReActFlow.WikiSearchAtomicFlow.instantiate_from_default_config
    overrides:
      name: "search"
      input_data_transformations:
        - _target_: flows.data_transformations.KeyRename
          old_key2new_key:
            argument1: "search_term"
  - _target_: FlowsEpfl.ReActFlow.FinishAtomicFlow.instantiate_from_default_config
    overrides:
      name: "finish"

Step 6 Putting everything together

Now we have implemented the ReActFlow:

Name: ReActFlow
Class name: ReActFlow
Type: circular
Description: ToDO: add description
Input keys: ['question']
Output keys: ['answer', 'status']
Subflows:
        Name: Agent
        Class name: AgentAtomicFlow
        Type: AtomicFlow
        Description: ToDO: add description
        Input keys: ['observation']
        Output keys: ['command', 'argument1', 'argument2']

        Name: Action
        Class name: ActionFlow
        Type: branching
        Description: ToDO: add description
        Input keys: None
        Output keys: ['observation']
        Subflows:
                Name: lookup
                Class name: WikiLookupAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['lookup_term', 'wiki_content']
                Output keys: ['observation']

                Name: search
                Class name: WikiSearchAtomicFlow
                Type: AtomicFlow
                Description: A flow that query wikipedia API for page content
                Input keys: ['search_term']
                Output keys: ['observation', 'wiki_content']

                Name: finish
                Class name: FinishAtomicFlow
                Type: AtomicFlow
                Description: A flow that outputs the final answer
                Input keys: None
                Output keys: ['observation', 'EARLY_EXIT']