File size: 3,614 Bytes
d63b11a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import netrc
import os
from typing import List, Optional, Tuple
import streamlit as st
def initialize_env_variables(prod_mode: bool = False, additional_env_vars: Optional[List[str]] = None) -> None:
if additional_env_vars is None:
additional_env_vars = []
if not prod_mode:
# In non-prod mode, prioritize environment variables
st.session_state.SAMBANOVA_API_KEY = os.environ.get(
'SAMBANOVA_API_KEY', st.session_state.get('SMABANOVA_API_KEY', '')
)
for var in additional_env_vars:
st.session_state[var] = os.environ.get(var, st.session_state.get(var, ''))
else:
# In prod mode, only use session state
if 'SAMBANOVA_API_KEY' not in st.session_state:
st.session_state.SAMBANOVA_API_KEY = ''
for var in additional_env_vars:
if var not in st.session_state:
st.session_state[var] = ''
def set_env_variables(api_key, additional_vars=None, prod_mode=False):
st.session_state.SAMBANOVA_API_KEY = api_key
if additional_vars:
for key, value in additional_vars.items():
st.session_state[key] = value
if not prod_mode:
# In non-prod mode, also set environment variables
os.environ['SAMBANOVA_API_KEY'] = api_key
if additional_vars:
for key, value in additional_vars.items():
os.environ[key] = value
def env_input_fields(additional_env_vars=None) -> Tuple[str, str]:
if additional_env_vars is None:
additional_env_vars = []
api_key = st.text_input('Sambanova API Key', value=st.session_state.SAMBANOVA_API_KEY, type='password')
additional_vars = {}
for var in additional_env_vars:
additional_vars[var] = st.text_input(f'{var}', value=st.session_state.get(var, ''), type='password')
return api_key, additional_vars
def are_credentials_set(additional_env_vars=None) -> bool:
if additional_env_vars is None:
additional_env_vars = []
base_creds_set = bool(st.session_state.SAMBANOVA_API_KEY)
additional_creds_set = all(bool(st.session_state.get(var, '')) for var in additional_env_vars)
return base_creds_set and additional_creds_set
def save_credentials(api_key, additional_vars=None, prod_mode=False) -> str:
set_env_variables(api_key, additional_vars, prod_mode)
return 'Credentials saved successfully!'
def get_wandb_key():
# Check for WANDB_API_KEY in environment variables
env_wandb_api_key = os.getenv('WANDB_API_KEY')
# Check for WANDB_API_KEY in ~/.netrc
try:
netrc_path = os.path.expanduser('~/.netrc')
netrc_data = netrc.netrc(netrc_path)
netrc_wandb_api_key = netrc_data.authenticators('api.wandb.ai')
except (FileNotFoundError, netrc.NetrcParseError):
netrc_wandb_api_key = None
# If both are set, handle the conflict
if env_wandb_api_key and netrc_wandb_api_key:
print('WANDB_API_KEY is set in both the environment and ~/.netrc. Prioritizing environment variable.')
# Optionally, you can choose to remove one of them, here we remove the env variable
del os.environ['WANDB_API_KEY'] # Remove from environment to prioritize ~/.netrc
return netrc_wandb_api_key[2] if netrc_wandb_api_key else None # Return the key from .netrc
# Return the key from environment if available, otherwise from .netrc
if env_wandb_api_key:
return env_wandb_api_key
elif netrc_wandb_api_key:
return netrc_wandb_api_key[2] if netrc_wandb_api_key else None
# If neither is set, return None
return None |