braintrust_add_metadata.py

#!/usr/bin/env python3
"""
Script to iterate over all logs in a Braintrust project and add metadata to each one.
Uses the Braintrust REST API directly.
"""

import os
import requests
from typing import Dict, Any, Optional, List
import braintrust


class BraintrustLogUpdater:
    def __init__(self, api_key: Optional[str] = None, project_name: Optional[str] = None):
        self.api_key = api_key or os.environ.get("BRAINTRUST_API_KEY")
        self.base_url = "https://api.braintrust.dev"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        # Initialize braintrust logger for updating spans
        braintrust.init(api_key=self.api_key, project=project_name)
        self.logger = None
    
    def get_project_id(self, project_name: str) -> Optional[str]:
        """Get project ID by name."""
        url = f"{self.base_url}/v1/project"
        response = requests.get(url, headers=self.headers)
        
        if response.status_code != 200:
            print(f"Failed to fetch projects: {response.status_code} {response.text}")
            return None
        
        projects = response.json().get("objects", [])
        for project in projects:
            if project.get("name") == project_name:
                return project.get("id")
        
        return None
    
    def fetch_project_logs(self, project_id: str, limit: int = 1000) -> List[Dict[str, Any]]:
        """Fetch all logs from a project."""
        url = f"{self.base_url}/v1/project_logs/{project_id}/fetch"
        params = {"limit": limit}
        
        response = requests.get(url, headers=self.headers, params=params)
        
        if response.status_code != 200:
            print(f"Failed to fetch project logs: {response.status_code} {response.text}")
            return []
        
        data = response.json()
        return data.get("events", [])
    
    def init_project_logger(self, project_name: str):
        """Initialize logger for the specific project."""
        if not self.logger:
            self.logger = braintrust.init_logger(project=project_name)
        return self.logger
    
    def update_log_metadata(
        self, 
        project_name: str,
        log_id: str, 
        metadata_update: Dict[str, Any]
    ) -> bool:
        """Update metadata for a specific log using the update_span method."""
        try:
            # Initialize logger for this project if not already done
            logger = self.init_project_logger(project_name)
            
            # Use update_span to update the metadata
            logger.update_span(id=log_id, metadata=metadata_update)
            
            return True
        except Exception as e:
            print(f"Failed to update span {log_id}: {e}")
            return False
    
    def add_metadata_to_all_logs(
        self,
        project_name: str,
        metadata_to_add: Dict[str, Any]
    ) -> int:
        """Add metadata to all logs in a project."""
        # Get project ID
        project_id = self.get_project_id(project_name)
        if not project_id:
            print(f"Project '{project_name}' not found")
            return 0
        
        print(f"Found project '{project_name}' with ID: {project_id}")
        
        # Fetch all logs
        print("Fetching project logs...")
        logs = self.fetch_project_logs(project_id)
        
        if not logs:
            print("No logs found in project")
            return 0
        
        print(f"Found {len(logs)} logs to update")
        
        # Update each log
        updated_count = 0
        for i, log in enumerate(logs, 1):
            log_id = log.get("id")
            if not log_id:
                print(f"Log #{i} has no ID, skipping")
                continue
            
            # Merge existing metadata with new metadata
            existing_metadata = log.get("metadata", {}) or {}
            updated_metadata = {**existing_metadata, **metadata_to_add}
            
            if self.update_log_metadata(project_name, log_id, updated_metadata):
                updated_count += 1
                print(f"✓ Updated log {log_id} (#{i}/{len(logs)})")
            else:
                print(f"✗ Failed to update log {log_id} (#{i}/{len(logs)})")
        
        return updated_count


def main():
    """Main function with example usage."""
    
    # Configuration
    PROJECT_NAME = "Natty Test 1"  # Replace with your actual project name
    
    # Metadata to add to each log
    METADATA_TO_ADD = {
        "processed_by": "metadata_script",
        # Add any additional metadata fields you need
    }
    
    # Check if API key is available
    api_key = os.environ.get("BRAINTRUST_API_KEY")
    if not api_key:
        print("Error: BRAINTRUST_API_KEY environment variable is not set")
        print("Set it with: export BRAINTRUST_API_KEY=your_api_key")
        return
    
    try:
        updater = BraintrustLogUpdater(api_key, PROJECT_NAME)
        
        print(f"Starting to process logs in project: {PROJECT_NAME}")
        print(f"Metadata to add: {METADATA_TO_ADD}")
        
        updated_count = updater.add_metadata_to_all_logs(PROJECT_NAME, METADATA_TO_ADD)
        
        print(f"\nCompleted! Successfully updated {updated_count} logs.")
        
    except Exception as e:
        print(f"Script failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()