Update api key file

flojoy/utils.py

@@ -213,57 +213,74 @@ def dump_str(result: Any, limit: int | None = None):
     )
 
 
-def get_frontier_api_key() -> Union[str, None]:
+def get_frontier_api_key(name: str) -> Union[str, None]:
+    try:
+        home = str(Path.home())
+        api_key = None
+        path = os.path.join(home, ".flojoy/credentials.yaml")
+        if not os.path.exists(path):
+            return api_key
+
+        with open(path, "r") as file:
+            load = yaml.safe_load(file)
+            if name == "cloud":
+                return load["CLOUD_API_KEY"]
+            elif name == "openai":
+                return load["OPENAI_API_KEY"]
+            else:
+                return {
+                    "s3Name": load[f"{name}"],
+                    "s3AccessKey": load[f"{name}accessKey"],
+                    "s3SecretKey": load[f"{name}secretKey"],
+                }
+    except Exception as e:
+        raise e
+
+
+def set_frontier_cloud_api(api_key: str):
     home = str(Path.home())
-    api_key = None
-    path = os.path.join(home, ".flojoy/credentials")
-    if not os.path.exists(path):
-        return api_key
-
-    stream = open(path, "r", encoding="utf-8")
-    yaml_dict = yaml.load(stream, Loader=yaml.FullLoader)
-    if yaml_dict is None:
-        return api_key
-    if isinstance(yaml_dict, str) == True:
-        split_by_line = yaml_dict.split("\n")
-        for line in split_by_line:
-            if "FRONTIER_API_KEY" in line:
-                api_key = line.split(":")[1]
-    else:
-        api_key = yaml_dict.get("FRONTIER_API_KEY", None)
-    return api_key
+    file_path = os.path.join(home, os.path.join(".flojoy", "credentials.yaml"))
 
+    data = {
+        "CLOUD_API_KEY": api_key,
+    }
+    if not os.path.exists(file_path):
+        # Create a new file and write the ACCSS_KEY to it
+        with open(file_path, "w") as file:
+            yaml.dump(data, file)
+        return
 
-def set_frontier_api_key(api_key: str):
-    try:
-        home = str(Path.home())
-        file_path = os.path.join(home, ".flojoy/credentials")
+    # Read the contents of the file
+    with open(file_path, "r") as file:
+        load = yaml.safe_load(file)
 
-        if not os.path.exists(file_path):
-            # Create a new file and write the API_KEY to it
-            with open(file_path, "w") as file:
-                file.write(f"FRONTIER_API_KEY:{api_key}\n")
-        else:
-            # Read the contents of the file
-            with open(file_path, "r") as file:
-                lines = file.readlines()
-
-            # Update the API key if it exists, otherwise append a new line
-            updated = False
-            for i, line in enumerate(lines):
-                if line.startswith("FRONTIER_API_KEY:"):
-                    lines[i] = f"FRONTIER_API_KEY:{api_key}\n"
-                    updated = True
-                    break
-
-            if not updated:
-                lines.append(f"FRONTIER_API_KEY:{api_key}\n")
-            # Write the updated contents to the file
-            with open(file_path, "w") as file:
-                file.writelines(lines)
+    load["CLOUD_API_KEY"] = api_key
 
-    except Exception as e:
-        raise e
+    with open(file_path, "w") as file:
+        yaml.dump(load, file)
+
+
+def set_frontier_openai_api(api_key: str):
+    home = str(Path.home())
+    file_path = os.path.join(home, os.path.join(".flojoy", "credentials.yaml"))
+
+    data = {
+        "OPENAI_API_KEY": api_key,
+    }
+    if not os.path.exists(file_path):
+        # Create a new file and write the ACCSS_KEY to it
+        with open(file_path, "w") as file:
+            yaml.dump(data, file)
+        return
+
+    # Read the contents of the file
+    with open(file_path, "r") as file:
+        load = yaml.safe_load(file)
+
+    load["OPENAI_API_KEY"] = api_key
+
+    with open(file_path, "w") as file:
+        yaml.dump(load, file)
 
 
 def set_frontier_s3_key(s3Name: str, s3AccessKey: str, s3SecretKey: str):