Use secrets manager for api key

This commit is contained in:
Aiden Dai
2025-02-10 15:25:12 +08:00
parent 74ca3b938e
commit c39f6bc942
4 changed files with 304 additions and 11 deletions

View File

@@ -1,22 +1,41 @@
import json
import os
from typing import Annotated
import boto3
from botocore.exceptions import ClientError
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from api.setting import DEFAULT_API_KEYS
api_key_param = os.environ.get("API_KEY_PARAM_NAME")
api_key_secret_arn = os.environ.get("API_KEY_SECRET_ARN")
api_key_env = os.environ.get("API_KEY")
if api_key_param:
# For backward compatibility.
# Please now use secrets manager instead.
ssm = boto3.client("ssm")
api_key = ssm.get_parameter(Name=api_key_param, WithDecryption=True)["Parameter"][
"Value"
]
elif api_key_secret_arn:
sm = boto3.client("secretsmanager")
try:
response = sm.get_secret_value(SecretId=api_key_secret_arn)
if "SecretString" in response:
secret = json.loads(response["SecretString"])
api_key = secret["api_key"]
except ClientError as e:
raise RuntimeError(
"Unable to retrieve API KEY, please ensure the secret ARN is correct"
)
except KeyError as e:
raise RuntimeError('Please ensure the secret contains a "api_key" field')
elif api_key_env:
api_key = api_key_env
else:
# For local use only.
api_key = DEFAULT_API_KEYS
security = HTTPBearer()

View File

@@ -105,7 +105,7 @@ def list_bedrock_models() -> dict:
status = model['modelLifecycle'].get('status', 'ACTIVE')
# currently, use this to filter out rerank models and legacy models
if not stream_supported or status != "ACTIVE":
if not stream_supported or status not in ["ACTIVE", "LEGACY"]:
continue
inference_types = model.get('inferenceTypesSupported', [])