Major overhaul to the OpenAIChatAtomicFlow.

OpenAIChatAtomicFlow.py

@@ -1,9 +1,7 @@
-import pprint
 from copy import deepcopy
 
 import hydra
 
-import colorama
 import time
 
 from typing import List, Dict, Optional, Any
@@ -12,31 +10,32 @@ from langchain import PromptTemplate
 import langchain
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 
-from flows.history import FlowHistory
 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.chat_message import ChatMessage
+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__)
 
+# ToDo: Add support for demonstrations
+
 
 class OpenAIChatAtomicFlow(AtomicFlow):
-    model_name: str
-    generation_parameters: Dict
+    REQUIRED_KEYS_CONFIG = ["model_name", "generation_parameters"]
+    REQUIRED_KEYS_KWARGS = ["system_message_prompt_template",
+                            "human_message_prompt_template",
+                            "query_message_prompt_template"]
 
-    system_message_prompt_template: PromptTemplate
-    human_message_prompt_template: PromptTemplate
+    SUPPORTS_CACHING: bool = True
 
-    system_name: str = "system"
-    user_name: str = "user"
-    assistant_name: str = "assistant"
+    api_keys: Dict[str, str]
 
-    n_api_retries: int = 6
-    wait_time_between_retries: int = 20
+    system_message_prompt_template: PromptTemplate
+    human_message_prompt_template: PromptTemplate
 
     query_message_prompt_template: Optional[PromptTemplate] = None
     demonstrations: GenericDemonstrationsDataset = None
@@ -44,9 +43,10 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     response_annotators: Optional[Dict[str, MessageAnnotator]] = {}
 
     def __init__(self, **kwargs):
-        self._validate_parameters(kwargs)
         super().__init__(**kwargs)
 
+        self.api_keys = None
+
         assert self.flow_config["name"] not in [
             "system",
             "user",
@@ -55,29 +55,11 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
     def set_up_flow_state(self):
         super().set_up_flow_state()
-        self.flow_state["conversation_initialized"] = False
+        self.flow_state["previous_messages"] = []
 
     @classmethod
     def _validate_parameters(cls, kwargs):
-        # ToDo: Deal with this in a cleaner way (with less repetition)
-        super()._validate_parameters(kwargs)
-
-        # ~~~ Model generation ~~~
-        if "model_name" not in kwargs["flow_config"]:
-            raise KeyError("model_name not specified in the flow_config.")
-
-        if "generation_parameters" not in kwargs["flow_config"]:
-            raise KeyError("generation_parameters not specified in the flow_config.")
-
-        # ~~~ Prompting ~~~
-        if "system_message_prompt_template" not in kwargs:
-            raise KeyError("system_message_prompt_template not passed to the constructor.")
-
-        if "query_message_prompt_template" not in kwargs:
-            raise KeyError("query_message_prompt_template not passed to the constructor.")
-
-        if "human_message_prompt_template" not in kwargs:
-            raise KeyError("human_message_prompt_template not passed to the constructor.")
+        validate_parameters(cls, kwargs)
 
     @classmethod
     def _set_up_prompts(cls, config):
@@ -92,19 +74,20 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
         return kwargs
 
-    @classmethod
-    def _set_up_demonstration_templates(cls, config):
-        kwargs = {}
-
-        if "demonstrations_response_template" in config:
-            kwargs["demonstrations_response_template"] = \
-                hydra.utils.instantiate(config['demonstrations_response_template'], _convert_="partial")
-
-        return kwargs
+    # @classmethod
+    # def _set_up_demonstration_templates(cls, config):
+    #     kwargs = {}
+    #
+    #     if "demonstrations_response_template" in config:
+    #         kwargs["demonstrations_response_template"] = \
+    #             hydra.utils.instantiate(config['demonstrations_response_template'], _convert_="partial")
+    #
+    #     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")
@@ -119,8 +102,8 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         # ~~~ Set up prompts ~~~
         kwargs.update(cls._set_up_prompts(flow_config))
 
-        # ~~~ Set up demonstration templates ~~~
-        kwargs.update(cls._set_up_demonstration_templates(flow_config))
+        # # ~~~ 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))
@@ -129,9 +112,13 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         return cls(**kwargs)
 
     def _is_conversation_initialized(self):
-        return self.flow_state["conversation_initialized"]
+        if len(self.flow_state["previous_messages"]) > 0:
+            return True
+
+        return False
 
-    def expected_inputs_given_state(self):
+    def get_expected_inputs(self, data: Optional[Dict[str, Any]] = None):
+        """Returns the expected inputs for the flow given the current state and, optionally, the input data"""
         if self._is_conversation_initialized():
             return ["query"]
         else:
@@ -146,13 +133,13 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         msg_content = prompt_template.format(**template_kwargs)
         return msg_content
 
-    def _get_demonstration_query_message_content(self, sample_data: Dict):
-        input_variables = self.query_message_prompt_template.input_variables
-        return self.query_message_prompt_template.format(**{k: sample_data[k] for k in input_variables}), []
-
-    def _get_demonstration_response_message_content(self, sample_data: Dict):
-        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_demonstration_query_message_content(self, sample_data: Dict):
+    #     input_variables = self.query_message_prompt_template.input_variables
+    #     return self.query_message_prompt_template.format(**{k: sample_data[k] for k in input_variables}), []
+    #
+    # def _get_demonstration_response_message_content(self, sample_data: Dict):
+    #     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():
@@ -162,75 +149,65 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     def _response_parsing(self, response: str, expected_outputs: List[str]):
         target_annotators = [ra for _, ra in self.response_annotators.items() if ra.key in expected_outputs]
 
-        if len(target_annotators) == 0:
-            return {expected_outputs[0]: response}
-
         parsed_outputs = {}
         for ra in target_annotators:
             parsed_out = ra(response)
             parsed_outputs.update(parsed_out)
-        return parsed_outputs
-
-    def _add_demonstrations(self):
-        if self.demonstrations is not None:
-            for example in self.demonstrations:
-                query, parents = self._get_demonstration_query_message_content(example)
-                response, parents = self._get_demonstration_response_message_content(example)
-
-                self._log_chat_message(content=query,
-                                       message_creator=self.user_name,
-                                       parent_message_ids=parents)
-
-                self._log_chat_message(content=response,
-                                       message_creator=self.assistant_name,
-                                       parent_message_ids=parents)
-
-    def _log_chat_message(self, message_creator: str, content: str, parent_message_ids: List[str] = None):
-        chat_message = ChatMessage(
-            message_creator=message_creator,
-            parent_message_ids=parent_message_ids,
-            flow_runner=self.flow_config["name"],
-            flow_run_id=self.flow_run_id,
-            content=content
-        )
-        return self._log_message(chat_message)
 
-    def _initialize_conversation(self, input_data: Dict[str, Any]):
-        # ~~~ Add the system message ~~~
-        system_message_content = self._get_message(self.system_message_prompt_template, input_data)
-
-        self._log_chat_message(content=system_message_content,
-                               message_creator=self.system_name)
-
-        # ~~~ Add the demonstration query-response tuples (if any) ~~~
-        self._add_demonstrations()
-        self._update_state(update_data={"conversation_initialized": True})
-
-    def get_conversation_messages(self, message_format: Optional[str] = None):
-        messages = self.flow_state["history"].get_chat_messages()
+        if "raw_response" in expected_outputs:
+            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 message_format is None:
-            return messages
+        if len(parsed_outputs) == 0:
+            raise Exception(f"The output dictionary is empty. "
+                            f"None of the expected outputs: `{str(expected_outputs)}` were found.")
 
-        elif message_format == "open_ai":
-            processed_messages = []
+        return parsed_outputs
 
-            for message in messages:
-                if message.message_creator == self.system_name:
-                    processed_messages.append(SystemMessage(content=message.content))
-                elif message.message_creator == self.assistant_name:
-                    processed_messages.append(AIMessage(content=message.content))
-                elif message.message_creator == self.user_name:
-                    processed_messages.append(HumanMessage(content=message.content))
-                else:
-                    raise ValueError(f"Unknown name: {message.message_creator}")
-            return processed_messages
+    # def _add_demonstrations(self):
+    #     if self.demonstrations is not None:
+    #         for example in self.demonstrations:
+    #             query, parents = self._get_demonstration_query_message_content(example)
+    #             response, parents = self._get_demonstration_response_message_content(example)
+    #
+    #             self._log_chat_message(content=query,
+    #                                    role=self.user_name,
+    #                                    parent_message_ids=parents)
+    #
+    #             self._log_chat_message(content=response,
+    #                                    role=self.assistant_name,
+    #                                    parent_message_ids=parents)
+
+    def _state_update_add_chat_message(self,
+                                       role: str,
+                                       content: str) -> None:
+
+        # Add the message to the previous messages list
+        if role == self.flow_config["system_name"]:
+            self.flow_state["previous_messages"].append(SystemMessage(content=content))
+        elif role == self.flow_config["user_name"]:
+            self.flow_state["previous_messages"].append(HumanMessage(content=content))
+        elif role == self.flow_config["assistant_name"]:
+            self.flow_state["previous_messages"].append(AIMessage(content=content))
         else:
-            raise ValueError(
-                f"Currently supported conversation message formats: 'open_ai'. '{message_format}' is not supported")
+            raise Exception(f"Invalid role: `{role}`.\n"
+                            f"Role should be one of: "
+                            f"`{self.flow_config['system_name']}`, "
+                            f"`{self.flow_config['user_name']}`, "
+                            f"`{self.flow_config['assistant_name']}`")
+
+        # Log the update to the flow messages list
+        chat_message = UpdateMessage_ChatMessage(
+            created_by=self.flow_config["name"],
+            updated_flow=self.flow_config["name"],
+            role=role,
+            content=content,
+        )
+        self._log_message(chat_message)
 
     def _call(self):
-        api_key = self.flow_state["api_key"]
+        api_key = self.api_keys["openai"]
 
         backend = langchain.chat_models.ChatOpenAI(
             model_name=self.flow_config["model_name"],
@@ -238,15 +215,13 @@ class OpenAIChatAtomicFlow(AtomicFlow):
             **self.flow_config["generation_parameters"],
         )
 
-        messages = self.get_conversation_messages(
-            message_format="open_ai"
-        )
+        messages = self.flow_state["previous_messages"]
 
         _success = False
         attempts = 1
         error = None
         response = None
-        while attempts <= self.n_api_retries:
+        while attempts <= self.flow_config['n_api_retries']:
             try:
                 response = backend(messages).content
                 _success = True
@@ -254,70 +229,69 @@ class OpenAIChatAtomicFlow(AtomicFlow):
             except Exception as e:
                 log.error(
                     f"Error {attempts} in calling backend: {e}. Key used: `{api_key}`. "
-                    f"Retrying in {self.wait_time_between_retries} seconds..."
-                )
-                log.error(
-                    f"API call raised Exception with the following arguments arguments: "
-                    f"\n{self.flow_state['history'].to_string()}"
+                    f"Retrying in {self.flow_config['wait_time_between_retries']} seconds..."
                 )
+                # log.error(
+                #     f"The API call raised an exception with the following arguments: "
+                #     f"\n{self.flow_state['history'].to_string()}"
+                # ) # ToDo: Make this message more user-friendly
                 attempts += 1
-                time.sleep(self.wait_time_between_retries)
+                time.sleep(self.flow_config['wait_time_between_retries'])
                 error = e
 
         if not _success:
             raise error
 
-        if self.flow_config["verbose"]:
-            messages_str = self.flow_state["history"].to_string()
-            log.info(
-                f"\n{colorama.Fore.MAGENTA}~~~ History [{self.flow_config['name']}] ~~~\n"
-                f"{colorama.Style.RESET_ALL}{messages_str}"
-            )
-
         return response
 
-    def _prepare_conversation(self, input_data: Dict[str, Any]):
+    def _initialize_conversation(self, input_data: Dict[str, Any]):
+        # ~~~ Add the system message ~~~
+        system_message_content = self._get_message(self.system_message_prompt_template, input_data)
+
+        self._state_update_add_chat_message(content=system_message_content,
+                                            role=self.flow_config["system_name"])
+
+        # # ~~~ Add the demonstration query-response tuples (if any) ~~~
+        # self._add_demonstrations()
+        # self._update_state(update_data={"conversation_initialized": True})
+
+    def _process_input(self, input_data: Dict[str, Any]):
         if self._is_conversation_initialized():
-            # ~~~ Check that the message has a `query` field ~~~
-            user_message_content = self.human_message_prompt_template.format(query=input_data["query"])
+            # Construct the message using the human message prompt template
+            user_message_content = self._get_message(self.human_message_prompt_template, input_data)
 
         else:
+            # Initialize the conversation (add the system message, and potentially the demonstrations)
             self._initialize_conversation(input_data)
+            # Construct the message using the query message prompt template
             user_message_content = self._get_message(self.query_message_prompt_template, input_data)
 
-        self._log_chat_message(message_creator=self.user_name,
-                               content=user_message_content)
+        self._state_update_add_chat_message(role=self.flow_config["user_name"],
+                                            content=user_message_content)
 
     @flow_run_cache()
-    def run(self, input_data: Dict[str, Any], expected_outputs: List[str]) -> Dict[str, Any]:
-        # ~~~ Chat-specific preparation ~~~
-        self._prepare_conversation(input_data)
+    def run(self,
+            input_data: Dict[str, Any],
+            private_keys: Optional[List[str]] = [],
+            keys_to_ignore_for_hash: Optional[List[str]] = []) -> Dict[str, Any]:
+        self.api_keys = input_data["api_keys"]
+        del input_data["api_keys"]
+
+        # ~~~ Process input ~~~
+        self._process_input(input_data)
 
         # ~~~ Call ~~~
         response = self._call()
-        answer_message = self._log_chat_message(
-            message_creator=self.flow_config["assistant_name"],
+        self._state_update_add_chat_message(
+            role=self.flow_config["assistant_name"],
             content=response
         )
 
         # ~~~ Response parsing ~~~
-        parsed_outputs = self._response_parsing(
+        output_data = self._response_parsing(
             response=response,
-            expected_outputs=expected_outputs
+            expected_outputs=input_data["expected_outputs"]
         )
-        self._update_state(update_data=parsed_outputs)
-
-        if self.flow_config["verbose"]:
-            parsed_output_messages_str = pprint.pformat({k: m for k, m in parsed_outputs.items()},
-                                                        indent=4)
-            log.info(
-                f"\n{colorama.Fore.MAGENTA}~~~ "
-                f"Response [{answer_message.message_creator} -- "
-                f"{answer_message.message_id} -- "
-                f"{answer_message.flow_run_id}] ~~~"
-                f"\n{colorama.Fore.YELLOW}Content: {answer_message}{colorama.Style.RESET_ALL}"
-                f"\n{colorama.Fore.YELLOW}Parsed Outputs: {parsed_output_messages_str}{colorama.Style.RESET_ALL}"
-            )
-
-        # ~~~ The final answer should be in self.flow_state, thus allow_class_namespace=False ~~~
-        return self._get_keys_from_state(keys=expected_outputs, allow_class_namespace=False)
+        # self._state_update_dict(update_data=output_data)  # ToDo: Is this necessary? When?
+
+        return output_data

OpenAIChatAtomicFlow.yaml

@@ -1,5 +1,16 @@
 # This is an abstract flow, therefore some required fields are not defined (and must be defined by the concrete flow)
 
+model_name: "gpt-4"
+generation_parameters:
+  n: 1
+  max_tokens: 3000
+  temperature: 0.3
+
+  model_kwargs:
+    top_p: 0.2
+    frequency_penalty: 0
+    presence_penalty: 0
+
 n_api_retries: 6
 wait_time_between_retries: 20
 
@@ -9,6 +20,24 @@ assistant_name: assistant
 
 response_annotators: {}
 
-query_message_prompt_template: null  # ToDo: When will this be null?
+system_message_prompt_template:
+  _target_: langchain.PromptTemplate
+  template_format: jinja2
+
+user_message_prompt_template:
+  _target_: langchain.PromptTemplate
+  template_format: jinja2
+
+human_message_prompt_template:
+  _target_: langchain.PromptTemplate
+  template_format: jinja2
+
+query_message_prompt_template:
+  _target_: langchain.PromptTemplate
+  template: "{{query}}"
+  input_variables:
+      - "query"
+  template_format: jinja2
+
 demonstrations: null
 demonstrations_response_template: null