Remove dependence on response annotators and replace them with data tranformations.

OpenAIChatAtomicFlow.py

@@ -10,14 +10,12 @@ from langchain import PromptTemplate
 import langchain
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 
-from flows.message_annotators.abstract import MessageAnnotator
 from flows.base_flows.abstract import AtomicFlow
 from flows.datasets import GenericDemonstrationsDataset
 
 from flows import utils
 from flows.messages.flow_message import UpdateMessage_ChatMessage
 from flows.utils.caching_utils import flow_run_cache
-from flows.utils.general_helpers import validate_parameters
 
 log = utils.get_pylogger(__name__)
 
@@ -40,7 +38,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     query_message_prompt_template: Optional[PromptTemplate] = None
     demonstrations: GenericDemonstrationsDataset = None
     demonstrations_response_template: PromptTemplate = None
-    response_annotators: Optional[Dict[str, MessageAnnotator]] = {}
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -57,10 +54,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         super().set_up_flow_state()
         self.flow_state["previous_messages"] = []
 
-    @classmethod
-    def _validate_parameters(cls, kwargs):
-        validate_parameters(cls, kwargs)
-
     @classmethod
     def _set_up_prompts(cls, config):
         kwargs = {}
@@ -84,21 +77,13 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     #
     #     return kwargs
 
-    @classmethod
-    def _set_up_response_annotators(cls, config):
-        response_annotators = config.get("response_annotators", {})
-        response_annotators = deepcopy(response_annotators)
-        if len(response_annotators) > 0:
-            for key, config in response_annotators.items():
-                response_annotators[key] = hydra.utils.instantiate(config, _convert_="partial")
-        return {"response_annotators": response_annotators}
-
     @classmethod
     def instantiate_from_config(cls, config):
         flow_config = deepcopy(config)
 
         kwargs = {"flow_config": flow_config}
-        kwargs["outputs_transformations"] = cls._set_up_outputs_transformations(flow_config)
+        kwargs["input_data_transformations"] = cls._set_up_data_transformations(config["input_data_transformations"])
+        kwargs["output_data_transformations"] = cls._set_up_data_transformations(config["output_data_transformations"])
 
         # ~~~ Set up prompts ~~~
         kwargs.update(cls._set_up_prompts(flow_config))
@@ -106,9 +91,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         # # ~~~ Set up demonstration templates ~~~
         # kwargs.update(cls._set_up_demonstration_templates(flow_config))
 
-        # ~~~ Set up response annotators ~~~
-        kwargs.update(cls._set_up_response_annotators(flow_config))
-
         # ~~~ Instantiate flow ~~~
         return cls(**kwargs)
 
@@ -142,30 +124,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     #     input_variables = self.demonstrations_response_template.input_variables
     #     return self.demonstrations_response_template.format(**{k: sample_data[k] for k in input_variables}), []
 
-    def _get_annotator_with_key(self, key: str):
-        for _, ra in self.response_annotators.items():
-            if ra.key == key:
-                return ra
-
-    def _response_parsing(self, response: str, output_keys: List[str]):
-        target_annotators = [ra for _, ra in self.response_annotators.items() if ra.key in output_keys]
-
-        parsed_outputs = {}
-        for ra in target_annotators:
-            parsed_out = ra(response)
-            parsed_outputs.update(parsed_out)
-
-        if "raw_response" in output_keys:
-            parsed_outputs["raw_response"] = response
-        else:
-            log.warning("The raw response is not logged because it was not requested as per the expected output.")
-
-        if len(parsed_outputs) == 0:
-            raise Exception(f"The output dictionary is empty. "
-                            f"None of the expected outputs: `{str(output_keys)}` were found.")
-
-        return parsed_outputs
-
     # def _add_demonstrations(self):
     #     if self.demonstrations is not None:
     #         for example in self.demonstrations:
@@ -288,11 +246,4 @@ class OpenAIChatAtomicFlow(AtomicFlow):
             content=response
         )
 
-        # ~~~ Response parsing ~~~
-        output_data = self._response_parsing(
-            response=response,
-            output_keys=input_data["output_keys"]
-        )
-        # self._state_update_dict(update_data=output_data)  # ToDo: Is this necessary? When?
-
-        return output_data
+        return response

OpenAIChatAtomicFlow.yaml

@@ -18,7 +18,7 @@ system_name: system
 user_name: user
 assistant_name: assistant
 
-response_annotators: {}
+output_data_transformations: []
 
 system_message_prompt_template:
   _target_: langchain.PromptTemplate
@@ -35,6 +35,7 @@ human_message_prompt_template:
     - "query"
   template_format: jinja2
 default_human_input_key: "query"
+input_data_transformations: []
 
 query_message_prompt_template:
   _target_: langchain.PromptTemplate