Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 68 additions & 3 deletions eval_protocol/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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")
Expand All @@ -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


Expand Down
4 changes: 3 additions & 1 deletion eval_protocol/evaluation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 15 additions & 15 deletions eval_protocol/pytest/handle_persist_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand All @@ -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",
Expand Down
15 changes: 10 additions & 5 deletions tests/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading