diff --git a/eval_protocol/auth.py b/eval_protocol/auth.py index 3002bd4a..f1d6c922 100644 --- a/eval_protocol/auth.py +++ b/eval_protocol/auth.py @@ -136,6 +136,56 @@ def _get_credential_from_config_file(key_name: str) -> Optional[str]: return None +def _get_credentials_from_config_file() -> Dict[str, Optional[str]]: + """ + Retrieve both api_key and account_id from auth.ini with a single read/parse. + Tries simple parsing first for both keys, then falls back to configparser for any missing ones. + Returns a dict with up to two keys: 'api_key' and 'account_id'. + """ + results: Dict[str, Optional[str]] = {} + auth_ini_path = _get_auth_ini_file() + if not auth_ini_path.exists(): + return results + + # 1) Simple key=value parsing + try: + simple_creds = _parse_simple_auth_file(auth_ini_path) + if "api_key" in simple_creds and simple_creds["api_key"]: + results["api_key"] = simple_creds["api_key"] + if "account_id" in simple_creds and simple_creds["account_id"]: + results["account_id"] = simple_creds["account_id"] + if "api_key" in results and "account_id" in results: + return results + except Exception as e: + logger.warning("Error during simple parsing of %s: %s", str(auth_ini_path), e) + + # 2) ConfigParser for any missing keys + try: + config = configparser.ConfigParser() + config.read(auth_ini_path) + for key_name in ("api_key", "account_id"): + if key_name in results and results[key_name]: + continue + if "fireworks" in config and config.has_option("fireworks", key_name): + value_from_file = config.get("fireworks", key_name) + if value_from_file: + results[key_name] = value_from_file + continue + if config.has_option(config.default_section, key_name): + value_from_default = config.get(config.default_section, key_name) + if value_from_default: + results[key_name] = value_from_default + except configparser.MissingSectionHeaderError: + # Purely key=value file without section headers; simple parsing should have handled it already. + logger.debug("%s has no section headers; falling back to simple parsing results.", str(auth_ini_path)) + except configparser.Error as e_config: + logger.warning("Configparser error reading %s: %s", str(auth_ini_path), e_config) + except Exception as e_general: + logger.warning("Unexpected error reading %s: %s", str(auth_ini_path), e_general) + + return results + + def get_fireworks_api_key() -> Optional[str]: """ Retrieves the Fireworks API key. @@ -177,13 +227,15 @@ def get_fireworks_account_id() -> Optional[str]: The Account ID is sourced in the following order: 1. FIREWORKS_ACCOUNT_ID environment variable. 2. 'account_id' from the [fireworks] section of ~/.fireworks/auth.ini. + 3. If an API key is available (env or auth.ini), resolve via verifyApiKey. Returns: The Account ID if found, otherwise None. """ # If a profile is active, prefer profile file first, then env if _is_profile_active(): - account_id_from_file = _get_credential_from_config_file("account_id") + creds = _get_credentials_from_config_file() + account_id_from_file = creds.get("account_id") if account_id_from_file: return account_id_from_file account_id = os.environ.get("FIREWORKS_ACCOUNT_ID") @@ -196,11 +248,24 @@ def get_fireworks_account_id() -> Optional[str]: if account_id: logger.debug("Using FIREWORKS_ACCOUNT_ID from environment variable.") return account_id - account_id_from_file = _get_credential_from_config_file("account_id") + creds = _get_credentials_from_config_file() + account_id_from_file = creds.get("account_id") if account_id_from_file: return account_id_from_file - logger.debug("Fireworks Account ID not found in environment variables or auth.ini.") + # 3) Fallback: if API key is present, attempt to resolve via verifyApiKey (env or auth.ini) + try: + # Intentionally use get_fireworks_api_key to centralize precedence (env vs file) + api_key_for_verify = get_fireworks_api_key() + if api_key_for_verify: + resolved = verify_api_key_and_get_account_id(api_key=api_key_for_verify, api_base=get_fireworks_api_base()) + if resolved: + logger.debug("Using FIREWORKS_ACCOUNT_ID resolved via verifyApiKey: %s", resolved) + return resolved + except Exception as e: + logger.debug("Failed to resolve FIREWORKS_ACCOUNT_ID via verifyApiKey: %s", e) + + logger.debug("Fireworks Account ID not found in environment variables, auth.ini, or via verifyApiKey.") return None diff --git a/eval_protocol/evaluation.py b/eval_protocol/evaluation.py index 7d7ef3da..6123f15d 100644 --- a/eval_protocol/evaluation.py +++ b/eval_protocol/evaluation.py @@ -595,7 +595,9 @@ def _ensure_requirements_present(source_dir: str) -> None: logger.error("Missing requirements.txt in upload directory: %s", source_dir) raise ValueError( "Upload requires requirements.txt in the project root. " - "Please add requirements.txt and re-run ep upload." + "Create a requirements.txt (it can be empty) and rerun 'eval-protocol upload' " + "or 'eval-protocol create rft'. If you're running in a notebook (e.g., Colab), " + f"create the file in your working directory (e.g., {source_dir}/requirements.txt)." ) @staticmethod diff --git a/eval_protocol/pytest/handle_persist_flow.py b/eval_protocol/pytest/handle_persist_flow.py index 07a627f3..66538903 100644 --- a/eval_protocol/pytest/handle_persist_flow.py +++ b/eval_protocol/pytest/handle_persist_flow.py @@ -11,6 +11,12 @@ from eval_protocol.directory_utils import find_eval_protocol_dir from eval_protocol.models import EvaluationRow from eval_protocol.pytest.store_experiment_link import store_experiment_link +from eval_protocol.auth import ( + get_fireworks_api_key, + get_fireworks_account_id, + verify_api_key_and_get_account_id, + get_fireworks_api_base, +) import requests @@ -90,22 +96,16 @@ def handle_persist_flow(all_results: list[list[EvaluationRow]], test_func_name: if not should_upload: continue - def get_auth_value(key: str) -> str | None: - """Get auth value from config file or environment.""" + # Resolve credentials using centralized auth helpers with verification fallback + fireworks_api_key = get_fireworks_api_key() + fireworks_account_id = get_fireworks_account_id() + if not fireworks_account_id and fireworks_api_key: try: - config_path = Path.home() / ".fireworks" / "auth.ini" - if config_path.exists(): - config = configparser.ConfigParser() # noqa: F821 - config.read(config_path) - for section in ["DEFAULT", "auth"]: - if config.has_section(section) and config.has_option(section, key): - return config.get(section, key) + fireworks_account_id = verify_api_key_and_get_account_id( + api_key=fireworks_api_key, api_base=get_fireworks_api_base() + ) except Exception: - pass - return os.getenv(key) - - fireworks_api_key = get_auth_value("FIREWORKS_API_KEY") - fireworks_account_id = get_auth_value("FIREWORKS_ACCOUNT_ID") + fireworks_account_id = None if not fireworks_api_key and not fireworks_account_id: store_experiment_link( @@ -129,7 +129,7 @@ def get_auth_value(key: str) -> str | None: ) continue - api_base = "https://api.fireworks.ai" + api_base = get_fireworks_api_base() headers = { "Authorization": f"Bearer {fireworks_api_key}", "Content-Type": "application/json", diff --git a/tests/test_auth.py b/tests/test_auth.py index 27ea3417..72dd54ca 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -255,7 +255,8 @@ def test_get_account_id_not_found(mock_path_exists): with patch("eval_protocol.auth._parse_simple_auth_file", return_value={}) as mock_parse_simple: assert get_fireworks_account_id() is None mock_parse_simple.assert_not_called() - mock_path_exists.assert_called_once_with() + # With verify fallback using get_fireworks_api_key, exists() may be checked more than once + assert mock_path_exists.call_count >= 1 @patch("pathlib.Path.exists", return_value=True) @@ -269,7 +270,8 @@ def test_get_account_id_ini_exists_no_section(mock_parse_simple, mock_ConfigPars mock_open(read_data="other_key = some_val_but_no_section_header\nanother=val"), ): assert get_fireworks_account_id() is None - mock_parse_simple.assert_called_once_with(AUTH_INI_FILE) + # Fallback verify path may trigger a second simple parse for api_key; ensure at least one call + assert mock_parse_simple.call_count >= 1 @patch("pathlib.Path.exists", return_value=True) @@ -283,7 +285,8 @@ def test_get_account_id_ini_exists_no_id_option(mock_parse_simple, mock_ConfigPa with patch("builtins.open", mock_open(read_data="[fireworks]\nsome_other_key=foo")): assert get_fireworks_account_id() is None - mock_parse_simple.assert_called_once_with(AUTH_INI_FILE) + # Fallback verify path may trigger a second simple parse for api_key; ensure at least one call + assert mock_parse_simple.call_count >= 1 @patch("pathlib.Path.exists", return_value=True) @@ -301,7 +304,8 @@ def test_get_account_id_ini_empty_value(mock_parse_simple, mock_ConfigParser_cla ) with patch("builtins.open", mock_open(read_data="[fireworks]\naccount_id=")): assert get_fireworks_account_id() is None - mock_parse_simple.assert_called_once_with(AUTH_INI_FILE) + # Fallback verify path may trigger a second simple parse for api_key; ensure at least one call + assert mock_parse_simple.call_count >= 1 @patch("pathlib.Path.exists", return_value=True) @@ -372,7 +376,8 @@ def test_get_account_id_ini_parse_error(mock_parse_simple, mock_ConfigParser_cla assert get_fireworks_account_id() is None assert "Configparser error reading" in caplog.text assert "Mocked Parsing Error" in caplog.text - mock_parse_simple.assert_called_once_with(AUTH_INI_FILE) + # Fallback verify path may trigger a second simple parse for api_key; ensure at least one call + assert mock_parse_simple.call_count >= 1 @patch("pathlib.Path.exists", return_value=True)