Update instantiation of flow.

OpenAIChatAtomicFlow.py

@@ -1,4 +1,6 @@
 import pprint
+from copy import deepcopy
+
 import hydra
 
 import colorama
@@ -7,15 +9,17 @@ import time
 from typing import List, Dict, Optional, Any
 
 from langchain import PromptTemplate
-from langchain.chat_models import ChatOpenAI
+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.utils.caching_utils import flow_run_cache
 
 log = utils.get_pylogger(__name__)
 
@@ -40,55 +44,98 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     response_annotators: Optional[Dict[str, MessageAnnotator]] = {}
 
     def __init__(self, **kwargs):
+        self._validate_parameters(kwargs)
+        super().__init__(**kwargs)
+
+        assert self.flow_config["name"] not in [
+            "system",
+            "user",
+            "assistant",
+        ], f"Flow name '{self.flow_config['name']}' cannot be 'system', 'user' or 'assistant'"
+
+    def set_up_flow_state(self):
+        super().set_up_flow_state()
+        self.flow_state["conversation_initialized"] = False
+
+    @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:
-            raise KeyError
+        if "model_name" not in kwargs["flow_config"]:
+            raise KeyError("model_name not specified in the flow_config.")
 
-        if "generation_parameters" not in kwargs:
-            raise KeyError
+        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
+            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
+            raise KeyError("human_message_prompt_template not passed to the constructor.")
 
-        super().__init__(**kwargs)
-        self._instantiate()
+    @classmethod
+    def _set_up_prompts(cls, config):
+        kwargs = {}
 
-        assert self.name not in [
-            "system",
-            "user",
-            "assistant",
-        ], f"Flow name '{self.name}' cannot be 'system', 'user' or 'assistant'"
-
-    def _instantiate(self):
-        # ~~~ Instantiate prompts ~~~
-        self.system_message_prompt_template = \
-            hydra.utils.instantiate(self.flow_config['system_message_prompt_template'], _convert_="partial")
-        self.query_message_prompt_template = \
-            hydra.utils.instantiate(self.flow_config['query_message_prompt_template'], _convert_="partial")
-        if self.flow_config["human_message_prompt_template"] is not None:
-            self.human_message_prompt_template = \
-                hydra.utils.instantiate(self.flow_config['human_message_prompt_template'], _convert_="partial")
-
-        # ~~~ Instantiate response annotators ~~~
-        if self.flow_config["response_annotators"] and len(self.flow_config["response_annotators"]) > 0:
-            for key, config in self.flow_config["response_annotators"].items():
-                self.response_annotators[key] = hydra.utils.instantiate(config, _convert_="partial")
-
-    def is_initialized(self):
-        conv_init = False
-        if "conversation_initialized" in self.flow_state:
-            conv_init = self.flow_state["conversation_initialized"]
-        return conv_init
+        kwargs["system_message_prompt_template"] = \
+            hydra.utils.instantiate(config['system_message_prompt_template'], _convert_="partial")
+        kwargs["query_message_prompt_template"] = \
+            hydra.utils.instantiate(config['query_message_prompt_template'], _convert_="partial")
+        kwargs["human_message_prompt_template"] = \
+            hydra.utils.instantiate(config['human_message_prompt_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", {})
+        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}
+
+        # ~~~ 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 response annotators ~~~
+        kwargs.update(cls._set_up_response_annotators(flow_config))
+
+        # ~~~ Instantiate flow ~~~
+        return cls(**kwargs)
+
+    def _is_conversation_initialized(self):
+        return self.flow_state["conversation_initialized"]
 
     def expected_inputs_given_state(self):
-        if self.is_initialized():
+        if self._is_conversation_initialized():
             return ["query"]
         else:
-            return self.expected_inputs
+            return self.flow_config["expected_inputs"]
 
     @staticmethod
     def _get_message(prompt_template, input_data: Dict[str, Any]):
@@ -100,10 +147,12 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         return msg_content
 
     def _get_demonstration_query_message_content(self, sample_data: Dict):
-        return self.query_message_prompt_template.format(**sample_data), []
+        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):
-        return self.demonstrations_response_template.format(**sample_data), []
+        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():
@@ -113,6 +162,9 @@ 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)
@@ -137,7 +189,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         chat_message = ChatMessage(
             message_creator=message_creator,
             parent_message_ids=parent_message_ids,
-            flow_runner=self.name,
+            flow_runner=self.flow_config["name"],
             flow_run_id=self.flow_run_id,
             content=content
         )
@@ -155,10 +207,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         self._update_state(update_data={"conversation_initialized": True})
 
     def get_conversation_messages(self, message_format: Optional[str] = None):
-        assert message_format is None or message_format in [
-            "open_ai"
-        ], f"Currently supported conversation message formats: 'open_ai'. '{message_format}' is not supported"
-
         messages = self.flow_state["history"].get_chat_messages()
 
         if message_format is None:
@@ -178,15 +226,16 @@ class OpenAIChatAtomicFlow(AtomicFlow):
                     raise ValueError(f"Unknown name: {message.message_creator}")
             return processed_messages
         else:
-            raise ValueError(f"Unknown message format: {message_format}")
+            raise ValueError(
+                f"Currently supported conversation message formats: 'open_ai'. '{message_format}' is not supported")
 
     def _call(self):
         api_key = self.flow_state["api_key"]
 
-        backend = ChatOpenAI(
-            model_name=self.model_name,
+        backend = langchain.chat_models.ChatOpenAI(
+            model_name=self.flow_config["model_name"],
             openai_api_key=api_key,
-            **self.generation_parameters,
+            **self.flow_config["generation_parameters"],
         )
 
         messages = self.get_conversation_messages(
@@ -218,17 +267,17 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         if not _success:
             raise error
 
-        if self.verbose:
+        if self.flow_config["verbose"]:
             messages_str = self.flow_state["history"].to_string()
             log.info(
-                f"\n{colorama.Fore.MAGENTA}~~~ History [{self.name}] ~~~\n"
+                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]):
-        if self.is_initialized():
+        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"])
 
@@ -239,14 +288,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         self._log_chat_message(message_creator=self.user_name,
                                content=user_message_content)
 
-        # if self.flow_state["dry_run"]:
-        #     messages_str = self.flow_state["history"].to_string()
-        #     log.info(
-        #         f"\n{colorama.Fore.MAGENTA}~~~ Messages [{self.name} -- {self.flow_run_id}] ~~~\n"
-        #         f"{colorama.Style.RESET_ALL}{messages_str}"
-        #     )
-        #     exit(0)
-
+    @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)
@@ -254,7 +296,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         # ~~~ Call ~~~
         response = self._call()
         answer_message = self._log_chat_message(
-            message_creator=self.assistant_name,
+            message_creator=self.flow_config["assistant_name"],
             content=response
         )
 
@@ -265,7 +307,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         )
         self._update_state(update_data=parsed_outputs)
 
-        if self.verbose:
+        if self.flow_config["verbose"]:
             parsed_output_messages_str = pprint.pformat({k: m for k, m in parsed_outputs.items()},
                                                         indent=4)
             log.info(

OpenAIChatAtomicFlow.yaml

@@ -1,4 +1,4 @@
-# This is an abstract flow, therefore some required fields are missing (not defined)
+# This is an abstract flow, therefore some required fields are not defined (and must be defined by the concrete flow)
 
 n_api_retries: 6
 wait_time_between_retries: 20