[WIP] Support fireworks login

Author: benjibc

README.md

@@ -70,6 +70,28 @@ Install with pip:
 pip install eval-protocol
 ```
 
+## Fireworks Login (REST)
+
+Use the CLI to sign in without gRPC.
+
+```
+# API key flow
+eval-protocol login --api-key YOUR_KEY --account-id YOUR_ACCOUNT_ID --validate
+
+# OAuth2 device flow (like firectl)
+eval-protocol login --oauth --issuer https://YOUR_ISSUER --client-id YOUR_PUBLIC_CLIENT_ID \
+  --account-id YOUR_ACCOUNT_ID --open-browser
+```
+
+- Omit `--api-key` to be prompted securely.
+- Omit `--account-id` to save only the key; you can add it later.
+- Add `--api-base https://api.fireworks.ai` for a custom base, if needed.
+- For OAuth2, you can also set env vars: `FIREWORKS_OIDC_ISSUER`, `FIREWORKS_OAUTH_CLIENT_ID`, `FIREWORKS_OAUTH_SCOPE`.
+
+Credentials are stored at `~/.fireworks/auth.ini` with 600 permissions and are read automatically by the SDK.
+
+Note: Model/LLM calls still require a Fireworks API key. OAuth login alone does not enable LLM calls yet; ensure `FIREWORKS_API_KEY` is set or saved via `eval-protocol login --api-key ...`.
+
 ## License
 
 [MIT](LICENSE)

eval_protocol/auth.py

@@ -2,8 +2,11 @@
 import logging
 import os
 from pathlib import Path
+import time
 from typing import Dict, Optional  # Added Dict
 
+import requests
+
 logger = logging.getLogger(__name__)
 
 FIREWORKS_CONFIG_DIR = Path.home() / ".fireworks"
@@ -36,7 +39,19 @@ def _parse_simple_auth_file(file_path: Path) -> Dict[str, str]:
                     ):
                         value = value[1:-1]
 
-                    if key in ["api_key", "account_id"] and value:
+                    if key in [
+                        "api_key",
+                        "account_id",
+                        "api_base",
+                        # OAuth2-related keys
+                        "issuer",
+                        "client_id",
+                        "access_token",
+                        "refresh_token",
+                        "expires_at",
+                        "scope",
+                        "token_type",
+                    ] and value:
                         creds[key] = value
     except Exception as e:
         logger.warning(f"Error during simple parsing of {file_path}: {e}")
@@ -142,15 +157,135 @@ def get_fireworks_api_base() -> str:
     """
     Retrieves the Fireworks API base URL.
 
-    The base URL is sourced from the FIREWORKS_API_BASE environment variable.
-    If not set, it defaults to "https://api.fireworks.ai".
+    The base URL is sourced in the following order:
+    1. FIREWORKS_API_BASE environment variable.
+    2. 'api_base' from the [fireworks] section of ~/.fireworks/auth.ini (or simple key=val).
+    3. Defaults to "https://api.fireworks.ai".
 
     Returns:
         The API base URL.
     """
-    api_base = os.environ.get("FIREWORKS_API_BASE", "https://api.fireworks.ai")
-    if os.environ.get("FIREWORKS_API_BASE"):
+    env_api_base = os.environ.get("FIREWORKS_API_BASE")
+    if env_api_base:
         logger.debug("Using FIREWORKS_API_BASE from environment variable.")
-    else:
-        logger.debug(f"FIREWORKS_API_BASE not set in environment, defaulting to {api_base}.")
-    return api_base
+        return env_api_base
+
+    file_api_base = _get_credential_from_config_file("api_base")
+    if file_api_base:
+        logger.debug("Using api_base from auth.ini configuration.")
+        return file_api_base
+
+    default_base = "https://api.fireworks.ai"
+    logger.debug(f"FIREWORKS_API_BASE not set; defaulting to {default_base}.")
+    return default_base
+
+
+def _get_from_env_or_file(key_name: str) -> Optional[str]:
+    # 1. Check env
+    env_val = os.environ.get(key_name.upper())
+    if env_val:
+        return env_val
+    # 2. Check config file
+    return _get_credential_from_config_file(key_name.lower())
+
+
+def _write_auth_config(updates: Dict[str, str]) -> None:
+    """Merge-write simple key=value pairs into AUTH_INI_FILE preserving existing values."""
+    FIREWORKS_CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    existing = _parse_simple_auth_file(AUTH_INI_FILE)
+    existing.update({k: v for k, v in updates.items() if v is not None})
+    lines = [f"{k}={v}" for k, v in existing.items()]
+    AUTH_INI_FILE.write_text("\n".join(lines) + "\n")
+    try:
+        os.chmod(AUTH_INI_FILE, 0o600)
+    except Exception:
+        pass
+
+
+def _discover_oidc(issuer: str) -> Dict[str, str]:
+    """Fetch OIDC discovery doc. Returns empty dict on failure."""
+    try:
+        url = issuer.rstrip("/") + "/.well-known/openid-configuration"
+        resp = requests.get(url, timeout=10)
+        if resp.ok:
+            return resp.json()
+    except Exception:
+        return {}
+    return {}
+
+
+def _refresh_oauth_token_if_needed() -> Optional[str]:
+    """Refresh OAuth access token if expired and refresh token available. Returns current/new token or None."""
+    cfg = _parse_simple_auth_file(AUTH_INI_FILE)
+    access_token = cfg.get("access_token")
+    refresh_token = cfg.get("refresh_token")
+    expires_at_str = cfg.get("expires_at")
+    issuer = cfg.get("issuer") or os.environ.get("FIREWORKS_OIDC_ISSUER")
+    client_id = cfg.get("client_id") or os.environ.get("FIREWORKS_OAUTH_CLIENT_ID")
+
+    # If we have no expiry, just return access token (best effort)
+    if not refresh_token or not issuer or not client_id:
+        return access_token
+
+    now = int(time.time())
+    try:
+        expires_at = int(expires_at_str) if expires_at_str else None
+    except ValueError:
+        expires_at = None
+
+    # If not expired (with 60s buffer), return current token
+    if access_token and expires_at and expires_at - 60 > now:
+        return access_token
+
+    # Attempt refresh
+    discovery = _discover_oidc(issuer)
+    token_endpoint = discovery.get("token_endpoint") or issuer.rstrip("/") + "/oauth/token"
+    data = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token,
+        "client_id": client_id,
+    }
+    try:
+        resp = requests.post(token_endpoint, data=data, timeout=15)
+        if not resp.ok:
+            logger.warning(f"OAuth token refresh failed: {resp.status_code} {resp.text[:200]}")
+            return access_token
+        tok = resp.json()
+        new_access = tok.get("access_token")
+        new_refresh = tok.get("refresh_token") or refresh_token
+        expires_in = tok.get("expires_in")
+        new_expires_at = str(now + int(expires_in)) if expires_in else expires_at_str
+        _write_auth_config(
+            {
+                "access_token": new_access,
+                "refresh_token": new_refresh,
+                "expires_at": new_expires_at,
+                "token_type": tok.get("token_type") or cfg.get("token_type") or "Bearer",
+                "scope": tok.get("scope") or cfg.get("scope") or "",
+            }
+        )
+        return new_access or access_token
+    except Exception as e:
+        logger.debug(f"Exception during oauth refresh: {e}")
+        return access_token
+
+
+def get_auth_bearer() -> Optional[str]:
+    """Return a bearer token to use in Authorization.
+
+    Priority:
+    1. FIREWORKS_ACCESS_TOKEN env
+    2. FIREWORKS_API_KEY env
+    3. Refreshed OAuth access_token from auth.ini (if present)
+    4. api_key from auth.ini
+    """
+    env_access = os.environ.get("FIREWORKS_ACCESS_TOKEN")
+    if env_access:
+        return env_access
+    env_key = os.environ.get("FIREWORKS_API_KEY")
+    if env_key:
+        return env_key
+    refreshed = _refresh_oauth_token_if_needed()
+    if refreshed:
+        return refreshed
+    return _get_credential_from_config_file("api_key")

eval_protocol/cli.py

@@ -28,6 +28,7 @@
 from .cli_commands.logs import logs_command
 from .cli_commands.preview import preview_command
 from .cli_commands.run_eval_cmd import hydra_cli_entry_point
+from .cli_commands.login import login_command
 
 
 def parse_args(args=None):
@@ -37,6 +38,30 @@ def parse_args(args=None):
 
     subparsers = parser.add_subparsers(dest="command", help="Command to run")
 
+    # Login command
+    login_parser = subparsers.add_parser(
+        "login", help="Sign in to Fireworks via API key or OAuth2 device flow"
+    )
+    # API key flow
+    login_parser.add_argument("--api-key", help="Fireworks API key (prompted if not provided)")
+    # OAuth2 flow toggles
+    login_parser.add_argument("--oauth", action="store_true", help="Use OAuth2 device flow (like firectl)")
+    login_parser.add_argument("--issuer", help="OIDC issuer URL (e.g., https://auth.fireworks.ai)")
+    login_parser.add_argument("--client-id", help="OAuth2 public client id for device flow")
+    login_parser.add_argument(
+        "--scope",
+        help="OAuth2 scopes (default: 'openid offline_access email profile')",
+    )
+    login_parser.add_argument(
+        "--open-browser", action="store_true", help="Attempt to open the verification URL in a browser"
+    )
+    # Common options
+    login_parser.add_argument("--account-id", help="Fireworks Account ID to associate with this login")
+    login_parser.add_argument("--api-base", help="Custom API base (defaults to https://api.fireworks.ai)")
+    vgroup = login_parser.add_mutually_exclusive_group()
+    vgroup.add_argument("--validate", action="store_true", help="Validate account with a test API call (API key flow)")
+    vgroup.add_argument("--no-validate", action="store_true", help="Do not validate; just write the file")
+
     # Preview command
     preview_parser = subparsers.add_parser("preview", help="Preview an evaluator with sample data")
     preview_parser.add_argument(
@@ -338,6 +363,10 @@ def main():
 
     if args.command == "preview":
         return preview_command(args)
+    elif args.command == "login":
+        # translate mutually exclusive group into a single boolean
+        setattr(args, "validate", bool(getattr(args, "validate", False) and not getattr(args, "no_validate", False)))
+        return login_command(args)
     elif args.command == "deploy":
         return deploy_command(args)
     elif args.command == "deploy-mcp":

eval_protocol/cli_commands/common.py

@@ -7,6 +7,8 @@
 import os
 from typing import Any, Dict, Iterator, List, Optional
 
+from eval_protocol.auth import get_auth_bearer
+
 logger = logging.getLogger(__name__)
 
 
@@ -42,13 +44,22 @@ def setup_logging(verbose=False, debug=False):
 
 
 def check_environment():
-    """Check if required environment variables are set for general commands."""
-    if not os.environ.get("FIREWORKS_API_KEY"):
-        logger.warning("FIREWORKS_API_KEY environment variable is not set.")
-        logger.warning("This is required for API calls. Set this variable before running the command.")
-        logger.warning("Example: FIREWORKS_API_KEY=$DEV_FIREWORKS_API_KEY reward-kit [command]")
-        return False
-    return True
+    """Check if credentials are available for non-LLM API calls.
+
+    Accepts either FIREWORKS_API_KEY or an OAuth bearer (FIREWORKS_ACCESS_TOKEN or tokens in auth.ini).
+    LLM calls elsewhere still explicitly require FIREWORKS_API_KEY.
+    """
+    if os.environ.get("FIREWORKS_API_KEY"):
+        return True
+    bearer = get_auth_bearer()
+    if bearer:
+        if not os.environ.get("FIREWORKS_API_KEY"):
+            logger.info(
+                "Using OAuth bearer for non-LLM API calls. Note: LLM/model calls still require FIREWORKS_API_KEY."
+            )
+        return True
+    logger.warning("No credentials found. Set FIREWORKS_API_KEY or login via OAuth: eval-protocol login --oauth ...")
+    return False
 
 
 def check_agent_environment(test_mode=False):