Apply log retention policies across all projects via API

import requests

api_key = 'sk-your-api-key'
project_id = "your project id"

# Retrieves the current retention policy configuration for a specific project
# This function queries the Braintrust API to get all project automations and filters
# for the retention policy automation specifically
# Args:
#   api_key (str): The API key for authentication with Braintrust API
#   project_id (str): The unique identifier of the project to query
# Returns:
#   dict: The retention policy automation configuration, or empty string if none exists
def get_retention_policy(api_key, project_id):
    url = "https://api.braintrust.dev/api/project_automation/get"

    headers = {
        "Authorization": "Bearer " + api_key,
        "Content-Type": "application/json"
    }

    data = {
        "project_id": project_id
    }

    response = requests.post(url, json=data, headers=headers)

    retention_config = ""
    for automation in response.json():
        if (automation['config']['event_type'] == 'retention'):
            retention_config = automation

    return(retention_config)

# Creates a new retention policy or updates an existing one for a project
# This function first checks if a retention policy already exists for the project.
# If no policy exists, it creates a new one with the specified parameters.
# If a policy exists, it updates the retention period while preserving other settings.
# Args:
#   api_key (str): The API key for authentication with Braintrust API
#   project_id (str): The unique identifier of the project
#   retention_period (int): Number of days to retain logs before deletion
#   retention_target (str): Type of objects to apply retention to (e.g., 'project_logs')
# Returns:
#   dict: The API response containing the created or updated policy information
def create_or_update_retention_policy(api_key, project_id, retention_period, retention_target):
    url = "https://api.braintrust.dev/api/project_automation/register"

    headers = {
        "Authorization": "Bearer " + api_key,
        "Content-Type": "application/json"
    }
    policy = get_retention_policy(api_key, project_id)
    if policy == "": 
        data = {
            "project_automation_name": str(retention_period) + "-Day Log Retention Policy",
            "description": "Automatically delete project logs older than + str(retention_period) + days",
            "project_id": project_id,
            "config": {
                "event_type": "retention",
                "object_type": retention_target,
                "retention_days": retention_period
            }
        }
        print("Retention policy for project ID " + project_id + " does not exist, creating")
    else:
        data = {
            "project_automation_name": policy['name'],
            "description": policy['description'],
            "project_id": project_id,
            "config": {
                "event_type": "retention",
                "object_type": policy['config']['object_type'],
                "retention_days": retention_period
            },
            "update": True
        }
        print("Retention policy for project ID " + project_id + " exist, updating")

    response = requests.post(url, json=data, headers=headers)
    return response.json()

# Retrieves a complete list of all projects accessible with the given API key
# This function handles pagination automatically to ensure all projects are returned,
# even if there are more than 500 projects (the API limit per request).
# It continues making requests until all projects have been retrieved.
# Args:
#   api_key (str): The API key for authentication with Braintrust API
# Returns:
#   list: A list of all project objects, each containing project details like id, name, etc.
def get_project_list(api_key):
    url = "https://api.braintrust.dev/v1/project"
    params = {
        "limit": 500
    }

    headers = {
        "Authorization": "Bearer " + api_key
    }

    response = requests.get(url, headers=headers, params=params)
    project_list = response.json()['objects']
    while len(response.json()['objects']) == 500:
        params = {
            "limit": 500,
            "starting_after": response.json()['objects'][499]['id']
        }
        response = requests.get(url, headers=headers, params=params)
        project_list.extend(response.json()['objects'])
    return project_list

# retention_config = get_retention_policy(api_key, project_id)
# print(retention_config)

# retention_config = create_or_update_retention_policy(api_key, project_id, 93, 'project_logs')
# print(retention_config)

# project_list = get_project_list(api_key)
# print(str(project_list))

# Get all projects
project_list = get_project_list(api_key)

# Apply 30-day retention to each project
for project in project_list:
    retention_config = create_or_update_retention_policy(api_key, project['id'], 30, 'project_logs')
    print(f"Project {project['name']}: {retention_config}")

retention_config = get_retention_policy(api_key, "your-project-id")
print(retention_config)