API key management in Prompt Engineering / GenAI - ML Experiment: Train & Evaluate

Experiment - API key management

Problem:You have built a machine learning model that uses an external API for data enrichment. Currently, the API key is hardcoded in your code, which poses security risks and makes it difficult to rotate keys.

Current Metrics:Model runs successfully but API key exposure risk is high; no automated key rotation; potential downtime if key is compromised.

Issue:Hardcoded API keys increase security risks and reduce maintainability. There is no system to manage or rotate keys safely.

Your Task

Implement a secure API key management system that keeps keys out of source code, supports easy rotation, and prevents unauthorized access.

Do not hardcode API keys in the source code.

Use environment variables or secure vaults for storing keys.

Ensure the model can access the API key at runtime without exposing it.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

Prompt Engineering / GenAI

import os
import requests

# Load API key from environment variable
API_KEY = os.getenv('API_KEY')

if not API_KEY:
    raise ValueError('API_KEY environment variable not set')

# Example function to call external API using the key

def get_enriched_data(data_id):
    url = f'https://api.example.com/data/{data_id}'
    headers = {'Authorization': f'Bearer {API_KEY}'}
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    return response.json()

# Usage example
if __name__ == '__main__':
    sample_id = '12345'
    result = get_enriched_data(sample_id)
    print(f'Enriched data for {sample_id}:', result)

# To run this code, set the API_KEY environment variable in your shell:
# export API_KEY='your_actual_api_key_here'

Removed hardcoded API key from source code.

Added code to load API key from environment variable.

Added error handling if API key is missing.

Demonstrated usage of API key in request headers securely.

Results Interpretation

Before: API key was hardcoded in the code, risking exposure and making rotation difficult.

After: API key is loaded securely from environment variables, reducing exposure risk and enabling easy rotation.

Managing API keys securely by keeping them out of source code and loading them at runtime improves security and maintainability of machine learning projects that rely on external APIs.

Bonus Experiment

Try integrating a secrets manager (like AWS Secrets Manager or HashiCorp Vault) to store and retrieve API keys dynamically instead of environment variables.

💡 Hint

Use the SDK of your chosen secrets manager to fetch the API key at runtime securely, and update your code to handle this retrieval.