test pr

.gitignore

@@ -158,4 +158,8 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
-.*cache*/
+.*cache*/
+
+
+# auto-generated by flows, all synced moduels will be ignored by default
+FLOW_MODULE_ID

OpenAIChatAtomicFlow.py

@@ -2,7 +2,7 @@ from copy import deepcopy
 
 import hydra
 
-import time
+import os, time
 
 from typing import List, Dict, Optional, Any
 
@@ -30,8 +30,6 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
     SUPPORTS_CACHING: bool = True
 
-    api_keys: Dict[str, str]
-
     system_message_prompt_template: PromptTemplate
     human_message_prompt_template: PromptTemplate
 
@@ -41,8 +39,8 @@ class OpenAIChatAtomicFlow(AtomicFlow):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-
-        self.api_keys = None
+        # TODO(yeeef): add documentation, what is guaranteed to be in flow_config?
+        # TODO(yeeef): everything set here as self.xxx=yyy will be invalidated after reset
 
         assert self.flow_config["name"] not in [
             "system",
@@ -165,8 +163,7 @@ class OpenAIChatAtomicFlow(AtomicFlow):
         )
         self._log_message(chat_message)
 
-    def _call(self):
-        api_key = self.api_keys["openai"]
+    def _call(self, api_key: str):
 
         backend = langchain.chat_models.ChatOpenAI(
             model_name=self.flow_config["model_name"],
@@ -233,14 +230,16 @@ class OpenAIChatAtomicFlow(AtomicFlow):
             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"]
+        
+        api_key = self.flow_config.get("api_key", "")
+        if "api_key" in input_data:
+            api_key = input_data.pop("api_key")
 
         # ~~~ Process input ~~~
         self._process_input(input_data)
 
         # ~~~ Call ~~~
-        response = self._call()
+        response = self._call(api_key)
         self._state_update_add_chat_message(
             role=self.flow_config["assistant_name"],
             content=response

OpenAIChatAtomicFlow.yaml

@@ -11,6 +11,8 @@ generation_parameters:
     frequency_penalty: 0
     presence_penalty: 0
 
+api_key:  "YOUR_API_KEY"
+
 n_api_retries: 6
 wait_time_between_retries: 20