Populate project names in list prompts API

curl --location 'https://api.braintrust.dev/v1/project' \
--header 'Authorization: Bearer sk-your-api-key'

curl --location 'https://api.braintrust.dev/v1/prompt' \
--header 'Authorization: Bearer sk-your-api-key'

#!/usr/bin/env python3
"""
Script to list all prompts per project using the Braintrust API.

This script:
1. Fetches all projects you have access to
2. For each project, fetches all prompts
3. Displays the results in a structured format
"""

import os
import requests
from typing import List, Dict, Any
import json

class BraintrustAPI:
    """Client for interacting with the Braintrust API."""

    def __init__(self, api_key: str = None):
        """
        Initialize the Braintrust API client.

        Args:
            api_key: Braintrust API key. If not provided, will look for BRAINTRUST_API_KEY env var.
        """
        self.api_key = api_key or os.environ.get('BRAINTRUST_API_KEY')
        if not self.api_key:
            raise ValueError(
                "API key is required. Set BRAINTRUST_API_KEY environment variable "
                "or pass api_key parameter."
            )

        self.base_url = "https://api.braintrust.dev"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def list_projects(self, limit: int = 1000) -> List[Dict[str, Any]]:
        """
        Fetch all projects with pagination.

        Args:
            limit: Number of projects to fetch per page (default: 1000)

        Returns:
            List of project objects
        """
        url = f"{self.base_url}/v1/project"
        all_projects = []
        starting_after = None
        page = 1

        try:
            while True:
                params = {"limit": limit}
                if starting_after:
                    params["starting_after"] = starting_after

                response = requests.get(url, headers=self.headers, params=params)
                response.raise_for_status()
                data = response.json()

                projects = data.get('objects', [])
                if not projects:
                    break

                all_projects.extend(projects)
                print(f"  Page {page}: Fetched {len(projects)} projects")

                # Check if we got fewer results than the limit, meaning we're done
                if len(projects) < limit:
                    break

                # Use the last project's ID for the next page
                starting_after = projects[-1].get('id')
                page += 1

            print(f"✓ Total projects fetched: {len(all_projects)}")
            return all_projects

        except requests.exceptions.RequestException as e:
            print(f"✗ Error fetching projects: {e}")
            return []

    def list_all_prompts(self, limit: int = 1000) -> List[Dict[str, Any]]:
        """
        Fetch all prompts with pagination (no filtering).

        Args:
            limit: Number of prompts to fetch per page (default: 1000)

        Returns:
            List of prompt objects
        """
        url = f"{self.base_url}/v1/prompt"
        all_prompts = []
        starting_after = None
        page = 1

        try:
            while True:
                params = {"limit": limit}

                if starting_after:
                    params['starting_after'] = starting_after

                response = requests.get(url, headers=self.headers, params=params)
                response.raise_for_status()
                data = response.json()

                prompts = data.get('objects', [])
                if not prompts:
                    break

                all_prompts.extend(prompts)
                print(f"  Page {page}: Fetched {len(prompts)} prompts")

                # Check if we got fewer results than the limit, meaning we're done
                if len(prompts) < limit:
                    break

                # Use the last prompt's ID for the next page
                starting_after = prompts[-1].get('id')
                page += 1

            print(f"✓ Total prompts fetched: {len(all_prompts)}")
            return all_prompts

        except requests.exceptions.RequestException as e:
            print(f"✗ Error fetching prompts: {e}")
            return []

def main():
    """Main function to list prompts per project."""

    print("=" * 80)
    print("Braintrust: Prompts per Project")
    print("=" * 80)
    print()

    # Initialize API client
    try:
        client = BraintrustAPI()
    except ValueError as e:
        print(f"Error: {e}")
        return

    # Fetch all projects
    print("Fetching projects...")
    projects = client.list_projects()

    if not projects:
        print("No projects found or error occurred.")
        return

    print()

    # Create a mapping of project_id to project_name
    project_map = {project.get('id'): project.get('name') for project in projects}

    print("Fetching all prompts...")
    all_prompts = client.list_all_prompts()

    if not all_prompts:
        print("No prompts found.")
        return

    print()
    print("=" * 80)
    print()

    # Filter and organize prompts by project
    results = {}
    prompts_with_projects = []

    for prompt in all_prompts:
        project_id = prompt.get('project_id')
        project_name = project_map.get(project_id, 'Unknown Project')

        # Add to results dictionary
        if project_name not in results:
            results[project_name] = []
        results[project_name].append(prompt)

        # Add to display list
        prompts_with_projects.append({
            'project_name': project_name,
            'prompt_name': prompt.get('name', 'Unnamed'),
            'prompt_id': prompt.get('id'),
            'description': prompt.get('description')
        })

    # Display prompts with their project names
    if prompts_with_projects:
        print("Prompts and their Projects:")
        print("-" * 80)
        for item in prompts_with_projects:
            print(f"• {item['prompt_name']} → Project: {item['project_name']}")
    else:
        print("No prompts found across all projects.")

    print()

    # Summary
    print("=" * 80)
    print("Summary")
    print("=" * 80)
    print(f"Total Projects: {len(projects)}")
    print(f"Projects with Prompts: {sum(1 for p in results.values() if p)}")
    print(f"Total Prompts: {sum(len(p) for p in results.values())}")
    print()

    # Optional: Save to JSON file
    output_file = "prompts_per_project.json"
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)
    print(f"✓ Results saved to {output_file}")

if __name__ == "__main__":
    main()