proper account detection (#322)

Author: benjibc

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
 
 

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

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",

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)