99import sys
1010from dataclasses import dataclass
1111from pathlib import Path
12- from typing import Any , Callable , Iterable , Optional
12+ from typing import Any , Dict , Iterable
1313
1414import pytest
1515from eval_protocol .auth import (
@@ -551,6 +551,35 @@ def _prompt_select(tests: list[DiscoveredTest], non_interactive: bool) -> list[D
551551 return _prompt_select_interactive (tests )
552552
553553
554+ def _load_secrets_from_env_file (env_file_path : str ) -> Dict [str , str ]:
555+ """
556+ Load secrets from a .env file that should be uploaded to Fireworks.
557+
558+ Returns a dictionary of secret key-value pairs that contain 'API_KEY' in the name.
559+ """
560+ if not os .path .exists (env_file_path ):
561+ return {}
562+
563+ # Load the .env file into a temporary environment
564+ env_vars = {}
565+ with open (env_file_path , "r" ) as f :
566+ for line in f :
567+ line = line .strip ()
568+ if line and not line .startswith ("#" ) and "=" in line :
569+ key , value = line .split ("=" , 1 )
570+ key = key .strip ()
571+ value = value .strip ().strip ('"' ).strip ("'" ) # Remove quotes
572+ env_vars [key ] = value
573+
574+ # Filter for secrets that look like API keys
575+ secrets = {}
576+ for key , value in env_vars .items ():
577+ if "API_KEY" in key .upper () and value :
578+ secrets [key ] = value
579+
580+ return secrets
581+
582+
554583def upload_command (args : argparse .Namespace ) -> int :
555584 root = os .path .abspath (getattr (args , "path" , "." ))
556585 entries_arg = getattr (args , "entry" , None )
@@ -585,11 +614,26 @@ def upload_command(args: argparse.Namespace) -> int:
585614 display_name = getattr (args , "display_name" , None )
586615 description = getattr (args , "description" , None )
587616 force = bool (getattr (args , "force" , False ))
617+ env_file = getattr (args , "env_file" , None )
588618
589- # Ensure FIREWORKS_API_KEY is available to the remote by storing it as a Fireworks secret
619+ # Load secrets from .env file and ensure they're available on Fireworks
590620 try :
591621 fw_account_id = get_fireworks_account_id ()
622+
623+ # Determine .env file path
624+ if env_file :
625+ env_file_path = env_file
626+ else :
627+ env_file_path = os .path .join (root , ".env" )
628+
629+ # Load secrets from .env file
630+ secrets_from_file = _load_secrets_from_env_file (env_file_path )
631+
632+ # Also ensure FIREWORKS_API_KEY from environment is included
592633 fw_api_key_value = get_fireworks_api_key ()
634+ if fw_api_key_value :
635+ secrets_from_file ["FIREWORKS_API_KEY" ] = fw_api_key_value
636+
593637 if not fw_account_id and fw_api_key_value :
594638 # Attempt to verify and resolve account id from server headers
595639 resolved = verify_api_key_and_get_account_id (api_key = fw_api_key_value , api_base = get_fireworks_api_base ())
@@ -598,21 +642,27 @@ def upload_command(args: argparse.Namespace) -> int:
598642 # Propagate to environment so downstream calls use it if needed
599643 os .environ ["FIREWORKS_ACCOUNT_ID" ] = fw_account_id
600644 print (f"Resolved FIREWORKS_ACCOUNT_ID via API verification: { fw_account_id } " )
601- if fw_account_id and fw_api_key_value :
602- print ("Ensuring FIREWORKS_API_KEY is registered as a secret on Fireworks for rollout..." )
603- if create_or_update_fireworks_secret (
604- account_id = fw_account_id ,
605- key_name = "FIREWORKS_API_KEY" ,
606- secret_value = fw_api_key_value ,
607- ):
608- print ("✓ FIREWORKS_API_KEY secret created/updated on Fireworks." )
609- else :
610- print ("Warning: Failed to create/update FIREWORKS_API_KEY secret on Fireworks." )
645+
646+ if fw_account_id and secrets_from_file :
647+ print (f"Found { len (secrets_from_file )} API keys to upload as Fireworks secrets..." )
648+ if env_file or os .path .exists (env_file_path ):
649+ print (f"Loading secrets from: { env_file_path } " )
650+
651+ for secret_name , secret_value in secrets_from_file .items ():
652+ print (f"Ensuring { secret_name } is registered as a secret on Fireworks for rollout..." )
653+ if create_or_update_fireworks_secret (
654+ account_id = fw_account_id ,
655+ key_name = secret_name ,
656+ secret_value = secret_value ,
657+ ):
658+ print (f"✓ { secret_name } secret created/updated on Fireworks." )
659+ else :
660+ print (f"Warning: Failed to create/update { secret_name } secret on Fireworks." )
611661 else :
612662 if not fw_account_id :
613- print ("Warning: FIREWORKS_ACCOUNT_ID not found; cannot register FIREWORKS_API_KEY secret ." )
614- if not fw_api_key_value :
615- print ("Warning: FIREWORKS_API_KEY not found locally; cannot register secret ." )
663+ print ("Warning: FIREWORKS_ACCOUNT_ID not found; cannot register secrets ." )
664+ if not secrets_from_file :
665+ print ("Warning: No API keys found in environment or .env file; no secrets to register ." )
616666 except Exception as e :
617667 print (f"Warning: Skipped Fireworks secret registration due to error: { e } " )
618668
0 commit comments