Add support for default_human_input_key and outputs transformations.

OpenAIChatAtomicFlow.py

@@ -98,6 +98,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         flow_config = deepcopy(config)
 
         kwargs = {"flow_config": flow_config}
+        kwargs["outputs_transformations"] = cls._set_up_outputs_transformations(flow_config)
 
         # ~~~ Set up prompts ~~~
         kwargs.update(cls._set_up_prompts(flow_config))
@@ -120,7 +121,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
     def get_input_keys(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"]
+            return [self.flow_config["default_human_input_key"]]
         else:
             return self.flow_config["input_keys"]
 
@@ -294,4 +295,4 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         )
         # self._state_update_dict(update_data=output_data)  # ToDo: Is this necessary? When?
 
-        return output_data
+        return output_data

OpenAIChatAtomicFlow.yaml

@@ -30,13 +30,14 @@ user_message_prompt_template:
 
 human_message_prompt_template:
   _target_: langchain.PromptTemplate
+  template: "{{query}}"
+  input_variables:
+    - "query"
   template_format: jinja2
+default_human_input_key: "query"
 
 query_message_prompt_template:
   _target_: langchain.PromptTemplate
-  template: "{{query}}"
-  input_variables:
-      - "query"
   template_format: jinja2
 
 demonstrations: null