Persist OpenAI API key

main.py

@@ -8,11 +8,14 @@
 def main():
     # Check if the user wants to set an API key
     if len(sys.argv) >= 2 and sys.argv[1] == '--set-api-key':
-        if len(sys.argv) != 3:
-            print("Usage: python main.py --set-api-key <API_KEY>")
+        if len(sys.argv) < 3:
+            print("Usage: python main.py --set-api-key <API_KEY> [--use-keyring]")
             return
-
-        openai_config.api_key_manager.set_api_key(sys.argv[2])
+
+        use_keyring = '--use-keyring' in sys.argv[2:]
+        api_key = sys.argv[2]
+
+        openai_config.api_key_manager.set_api_key(api_key, use_keyring=use_keyring)
 
         print("API key set successfully.")
         return
@@ -21,8 +24,12 @@ def main():
     if len(sys.argv) != 2:
         # Check Python version and adjust the usage message accordingly
         try:
-            python_version = subprocess.check_output(['python', '--version'])
-            version_message = "python3" if "Python 3" in str(python_version) else "python"
+            python_version = subprocess.check_output(["python", "--version"])
+            version_message = (
+                "python3"
+                if "Python 3" in str(python_version)
+                else "python"
+            )
         except subprocess.CalledProcessError:
             # In case the above check fails, default to 'python3'
             version_message = "python3"

openai_config/api_key_manager.py

@@ -1,12 +1,106 @@
+"""Helpers for managing the OpenAI API key.
+
+The key can be stored in an ``.env`` file or in the system keyring so that the
+CLI does not require it to be entered every time.  These helpers keep the
+environment variable ``OPENAI_API_KEY`` in sync with the persisted value.
+"""
+
+from __future__ import annotations
+
 import os
+from pathlib import Path
+from typing import Optional
 
-def get_api_key():
-    api_key = os.getenv('OPENAI_API_KEY')
-    if not api_key:
-        raise ValueError("API key not found. Set the OPENAI_API_KEY environment variable.")
-    return api_key
+ENV_FILENAME = ".env"
+SERVICE_NAME = "openai_finetuning_framework"
+
+
+def _read_key_from_env_file(file_path: str | Path = ENV_FILENAME) -> Optional[str]:
+    """Return the API key stored in ``file_path`` if present."""
+
+    path = Path(file_path)
+    if not path.exists():
+        return None
+    for line in path.read_text().splitlines():
+        if line.startswith("OPENAI_API_KEY="):
+            value = line.partition("=")[2].strip()
+            return value.strip('"').strip("'")
+    return None
+
+
+def _write_key_to_env_file(api_key: str, file_path: str | Path = ENV_FILENAME) -> None:
+    """Persist ``api_key`` to ``file_path`` creating or updating the file."""
+
+    path = Path(file_path)
+    lines: list[str] = []
+    if path.exists():
+        lines = path.read_text().splitlines()
+    updated = False
+    for i, line in enumerate(lines):
+        if line.startswith("OPENAI_API_KEY="):
+            lines[i] = f"OPENAI_API_KEY={api_key}"
+            updated = True
+            break
+    if not updated:
+        lines.append(f"OPENAI_API_KEY={api_key}")
+    path.write_text("\n".join(lines) + "\n")
+
+
+def _read_key_from_keyring() -> Optional[str]:
+    """Return the API key stored in the system keyring if available."""
+
+    try:
+        import keyring  # type: ignore
+
+        return keyring.get_password(SERVICE_NAME, "OPENAI_API_KEY")
+    except Exception:
+        return None
+
+
+def _write_key_to_keyring(api_key: str) -> bool:
+    """Store ``api_key`` in the system keyring if possible."""
+
+    try:
+        import keyring  # type: ignore
+
+        keyring.set_password(SERVICE_NAME, "OPENAI_API_KEY", api_key)
+        return True
+    except Exception:
+        return False
+
+
+def get_api_key() -> str:
+    """Retrieve the API key from the environment, ``.env`` file or keyring."""
+
+    api_key = os.getenv("OPENAI_API_KEY")
+    if api_key:
+        return api_key
+
+    api_key = _read_key_from_env_file()
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+        return api_key
+
+    api_key = _read_key_from_keyring()
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+        return api_key
+
+    raise ValueError(
+        "API key not found. Set it using `python main.py --set-api-key <KEY>`"
+    )
+
+
+def set_api_key(api_key: str, use_keyring: bool = False) -> None:
+    """Persist ``api_key`` to ``.env`` or keyring and update the environment."""
 
-def set_api_key(api_key):
     if not api_key:
         raise ValueError("Invalid API key provided.")
-    os.environ['OPENAI_API_KEY'] = api_key
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    if use_keyring:
+        if not _write_key_to_keyring(api_key):
+            # Fall back to .env file if keyring is unavailable
+            _write_key_to_env_file(api_key)
+    else:
+        _write_key_to_env_file(api_key)

readme.md

@@ -6,13 +6,19 @@ You have two way to manipulate/enhance LLMs; that being RAG (Retrieval Augmented
 
 today's models are as bad as we allow them to be, if you finetune for a task and it performs better than gpt4 then congrats that's SOTA (state of the art).
 
-## How to use 
+## How to use
 ### install dependencies and set api key
 ```
 cd openai_finetuning_framework
 chmod +x build.sh
 ./build.sh
 ```
+### setting your OpenAI API key
+Persist the key in a `.env` file (default) or in your system keyring:
+```
+python main.py --set-api-key YOUR_KEY            # store in .env
+python main.py --set-api-key YOUR_KEY --use-keyring  # store in keyring
+```
 ### want to try on a dataset now?
 ```
  git clone https://huggingface.co/datasets/karan4d/machiavellian_synthetic_textbooks