diff --git a/main.py b/main.py index 34ee213..09fd554 100644 --- a/main.py +++ b/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 ") + if len(sys.argv) < 3: + print("Usage: python main.py --set-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" diff --git a/openai_config/api_key_manager.py b/openai_config/api_key_manager.py index ffdb9ca..fae341d 100644 --- a/openai_config/api_key_manager.py +++ b/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 `" + ) + + +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 \ No newline at end of file + 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) diff --git a/readme.md b/readme.md index 9566958..18de1ad 100644 --- a/readme.md +++ b/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