aider/scripts/badges.py

#!/usr/bin/env python3

import argparse
import os
import sys

import requests
import yaml
from dotenv import load_dotenv
from google.cloud import bigquery
from google.oauth2 import service_account

TOKENS_PER_WEEK = "15B"

# Badge tooltip texts
GITHUB_STARS_TOOLTIP = "Total number of GitHub stars the Aider project has received"
PYPI_DOWNLOADS_TOOLTIP = "Total number of installations via pip from PyPI"
TOKENS_WEEKLY_TOOLTIP = "Number of tokens processed weekly by Aider users"
OPENROUTER_TOOLTIP = "Aider's ranking among applications on the OpenRouter platform"
SINGULARITY_TOOLTIP = "Percentage of the new code in Aider's last release written by Aider itself"


def get_downloads_from_bigquery(credentials_path=None, package_name="aider-chat"):
    """
    Fetch download statistics for a package from Google BigQuery PyPI dataset
    """
    try:
        # Initialize credentials if path provided
        credentials = None
        if credentials_path:
            credentials = service_account.Credentials.from_service_account_file(
                credentials_path, scopes=["https://www.googleapis.com/auth/cloud-platform"]
            )

        # Create a client
        client = bigquery.Client(credentials=credentials)

        # Query to get total downloads for the package, excluding CI/CD systems
        query = f"""
            SELECT COUNT(*) as total_downloads
            FROM `bigquery-public-data.pypi.file_downloads`
            WHERE file.project = '{package_name}'
            AND NOT (
                -- Exclude common CI/CD systems based on installer name patterns
                LOWER(details.installer.name) LIKE '%github%' OR
                LOWER(details.installer.name) LIKE '%travis%' OR
                LOWER(details.installer.name) LIKE '%circle%' OR
                LOWER(details.installer.name) LIKE '%jenkins%' OR
                LOWER(details.installer.name) LIKE '%gitlab%' OR
                LOWER(details.installer.name) LIKE '%azure%' OR
                LOWER(details.installer.name) LIKE '%ci%' OR
                LOWER(details.installer.name) LIKE '%cd%' OR
                LOWER(details.installer.name) LIKE '%bot%' OR
                LOWER(details.installer.name) LIKE '%build%'
            )
        """

        # Execute the query
        query_job = client.query(query)
        results = query_job.result()

        # Get the first (and only) row
        for row in results:
            return row.total_downloads

        return 0
    except Exception as e:
        print(f"Error fetching download statistics from BigQuery: {e}", file=sys.stderr)
        return None


def get_total_downloads(
    api_key=None, package_name="aider-chat", use_bigquery=False, credentials_path=None
):
    """
    Fetch total downloads for a Python package

    If use_bigquery is True, fetches from BigQuery.
    Otherwise uses pepy.tech API (requires api_key).
    """
    if use_bigquery:
        return get_downloads_from_bigquery(credentials_path, package_name)

    # Fall back to pepy.tech API
    if not api_key:
        print("API key not provided for pepy.tech", file=sys.stderr)
        sys.exit(1)

    url = f"https://api.pepy.tech/api/v2/projects/{package_name}"
    headers = {"X-API-Key": api_key}

    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()  # Raise an exception for HTTP errors

        data = response.json()
        total_downloads = data.get("total_downloads", 0)

        return total_downloads
    except requests.exceptions.RequestException as e:
        print(f"Error fetching download statistics from pepy.tech: {e}", file=sys.stderr)
        sys.exit(1)


def get_github_stars(repo="paul-gauthier/aider"):
    """
    Fetch the number of GitHub stars for a repository
    """
    url = f"https://api.github.com/repos/{repo}"
    headers = {"Accept": "application/vnd.github.v3+json"}

    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()  # Raise an exception for HTTP errors

        data = response.json()
        stars = data.get("stargazers_count", 0)

        return stars
    except requests.exceptions.RequestException as e:
        print(f"Error fetching GitHub stars: {e}", file=sys.stderr)
        return None


def get_latest_release_aider_percentage():
    """
    Get the percentage of code written by Aider in the LATEST release
    from the blame.yml file
    """
    blame_path = os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        "aider",
        "website",
        "_data",
        "blame.yml",
    )

    try:
        with open(blame_path, "r") as f:
            blame_data = yaml.safe_load(f)

        if not blame_data or len(blame_data) == 0:
            return 0, "unknown"

        # Find the latest release by parsing version numbers
        latest_version = None
        latest_release = None

        for release in blame_data:
            version_tag = release.get("end_tag", "")
            if not version_tag.startswith("v"):
                continue

            # Parse version like "v0.77.0" into a tuple (0, 77, 0)
            try:
                version_parts = tuple(int(part) for part in version_tag[1:].split("."))
                if latest_version is None or version_parts > latest_version:
                    latest_version = version_parts
                    latest_release = release
            except ValueError:
                # Skip if version can't be parsed as integers
                continue

        if latest_release:
            percentage = latest_release.get("aider_percentage", 0)
            version = latest_release.get("end_tag", "unknown")
            return percentage, version

        return 0, "unknown"
    except Exception as e:
        print(f"Error reading blame data: {e}", file=sys.stderr)
        return 0, "unknown"


def format_number(number):
    """
    Format a large number with K, M, B suffixes with 1 decimal place
    """
    if number is None:
        return "0"

    if number >= 1_000_000_000:
        return f"{number / 1_000_000_000:.1f}B"
    elif number >= 1_000_000:
        return f"{number / 1_000_000:.1f}M"
    elif number >= 1_000:
        return f"{number / 1_000:.1f}K"
    else:
        return str(number)


def generate_badges_md(downloads, stars, aider_percentage):
    """
    Generate markdown for badges with updated values
    """
    # Format downloads to 1 decimal place with M suffix
    downloads_formatted = format_number(downloads)

    # Round aider percentage to whole number
    aider_percent_rounded = round(aider_percentage)

    markdown = f"""  <a href="https://github.com/Aider-AI/aider/stargazers"><img alt="GitHub Stars" title="{GITHUB_STARS_TOOLTIP}"
src="https://img.shields.io/github/stars/Aider-AI/aider?style=flat-square&logo=github&color=f1c40f&labelColor=555555"/></a>
  <a href="https://pypi.org/project/aider-chat/"><img alt="PyPI Downloads" title="{PYPI_DOWNLOADS_TOOLTIP}"
src="https://img.shields.io/badge/📦%20Installs-{downloads_formatted}-2ecc71?style=flat-square&labelColor=555555"/></a>
  <img alt="Tokens per week" title="{TOKENS_WEEKLY_TOOLTIP}"
src="https://img.shields.io/badge/📈%20Tokens%2Fweek-{TOKENS_PER_WEEK}-e74c3c?style=flat-square&labelColor=555555"/>
  <a href="https://openrouter.ai/"><img alt="OpenRouter Ranking" title="{OPENROUTER_TOOLTIP}"
src="https://img.shields.io/badge/🏆%20OpenRouter-Top%2020-9b59b6?style=flat-square&labelColor=555555"/></a>
  <a href="https://aider.chat/HISTORY.html"><img alt="Singularity" title="{SINGULARITY_TOOLTIP}"
src="https://img.shields.io/badge/🔄%20Singularity-{aider_percent_rounded}%25-3498db?style=flat-square&labelColor=555555"/></a>"""  # noqa

    return markdown


def get_badges_md():
    """
    Get all statistics and return the generated badges markdown
    """
    # Load environment variables from .env file
    load_dotenv()

    # Check if we should use BigQuery and get credentials path
    bigquery_env = os.environ.get("USE_BIGQUERY", "false")
    use_bigquery = bigquery_env.lower() in ("true", "1", "yes") or os.path.exists(bigquery_env)
    credentials_path = bigquery_env if os.path.exists(bigquery_env) else None

    # Get API key from environment variable if not using BigQuery
    api_key = None
    if not use_bigquery:
        api_key = os.environ.get("PEPY_API_KEY")
        if not api_key:
            print(
                (
                    "API key not provided and BigQuery not enabled. Please set PEPY_API_KEY"
                    " environment variable"
                ),
                file=sys.stderr,
            )
            sys.exit(1)

    # Get PyPI downloads for the default package
    total_downloads = get_total_downloads(api_key, "aider-chat", use_bigquery, credentials_path)

    # Get GitHub stars for the default repo
    stars = get_github_stars("paul-gauthier/aider")

    # Get Aider contribution percentage in latest release
    percentage, _ = get_latest_release_aider_percentage()

    # Generate and return badges markdown
    return generate_badges_md(total_downloads, stars, percentage)


def get_badges_html():
    """
    Get all statistics and return HTML-formatted badges
    """
    # Load environment variables from .env file
    load_dotenv()

    # Check if we should use BigQuery and get credentials path
    bigquery_env = os.environ.get("USE_BIGQUERY", "false")
    use_bigquery = bigquery_env.lower() in ("true", "1", "yes") or os.path.exists(bigquery_env)
    credentials_path = bigquery_env if os.path.exists(bigquery_env) else None

    # Get API key from environment variable if not using BigQuery
    api_key = None
    if not use_bigquery:
        api_key = os.environ.get("PEPY_API_KEY")
        if not api_key:
            print(
                (
                    "API key not provided and BigQuery not enabled. Please set PEPY_API_KEY"
                    " environment variable"
                ),
                file=sys.stderr,
            )
            sys.exit(1)

    # Get PyPI downloads for the default package
    total_downloads = get_total_downloads(api_key, "aider-chat", use_bigquery, credentials_path)

    # Get GitHub stars for the default repo
    stars = get_github_stars("paul-gauthier/aider")

    # Get Aider contribution percentage in latest release
    percentage, _ = get_latest_release_aider_percentage()

    # Format values
    downloads_formatted = format_number(total_downloads)
    # Stars should be rounded to whole numbers
    if stars is None:
        stars_formatted = "0"
    elif stars >= 1_000_000_000:
        stars_formatted = f"{round(stars / 1_000_000_000)}B"
    elif stars >= 1_000_000:
        stars_formatted = f"{round(stars / 1_000_000)}M"
    elif stars >= 1_000:
        stars_formatted = f"{round(stars / 1_000)}K"
    else:
        stars_formatted = str(int(round(stars)))
    aider_percent_rounded = round(percentage)

    # Generate HTML badges
    html = f"""<a href="https://github.com/Aider-AI/aider" class="github-badge badge-stars" title="{GITHUB_STARS_TOOLTIP}">
    <span class="badge-label">⭐ GitHub Stars</span>
    <span class="badge-value">{stars_formatted}</span>
</a>
<a href="https://pypi.org/project/aider-chat/" class="github-badge badge-installs" title="{PYPI_DOWNLOADS_TOOLTIP}">
    <span class="badge-label">📦 Installs</span>
    <span class="badge-value">{downloads_formatted}</span>
</a>
<div class="github-badge badge-tokens" title="{TOKENS_WEEKLY_TOOLTIP}">
    <span class="badge-label">📈 Tokens/week</span>
    <span class="badge-value">{TOKENS_PER_WEEK}</span>
</div>
<a href="https://openrouter.ai/" class="github-badge badge-router" title="{OPENROUTER_TOOLTIP}">
    <span class="badge-label">🏆 OpenRouter</span>
    <span class="badge-value">Top 20</span>
</a>
<a href="/HISTORY.html" class="github-badge badge-coded" title="{SINGULARITY_TOOLTIP}">
    <span class="badge-label">🔄 Singularity</span>
    <span class="badge-value">{aider_percent_rounded}%</span>
</a>"""  # noqa

    return html


def main():
    # Load environment variables from .env file
    load_dotenv()

    parser = argparse.ArgumentParser(description="Get total downloads and GitHub stars for aider")
    parser.add_argument(
        "--api-key",
        help=(
            "pepy.tech API key (can also be set via PEPY_API_KEY in .env file or environment"
            " variable)"
        ),
    )
    parser.add_argument(
        "--package", default="aider-chat", help="Package name (default: aider-chat)"
    )
    parser.add_argument(
        "--github-repo",
        default="paul-gauthier/aider",
        help="GitHub repository (default: paul-gauthier/aider)",
    )
    parser.add_argument("--markdown", action="store_true", help="Generate markdown badges block")
    parser.add_argument(
        "--use-bigquery",
        action="store_true",
        help="Use BigQuery to fetch download statistics instead of pepy.tech",
    )
    parser.add_argument(
        "--credentials-path", help="Path to Google Cloud service account credentials JSON file"
    )
    args = parser.parse_args()

    # Determine whether to use BigQuery and get credentials path
    bigquery_env = os.environ.get("USE_BIGQUERY", "false")
    use_bigquery = (
        args.use_bigquery
        or bigquery_env.lower() in ("true", "1", "yes")
        or os.path.exists(bigquery_env)
    )
    credentials_path = args.credentials_path or (
        bigquery_env if os.path.exists(bigquery_env) else None
    )

    # Check for required parameters
    api_key = None
    if not use_bigquery:
        # Get API key from args or environment variable
        api_key = args.api_key or os.environ.get("PEPY_API_KEY")
        if not api_key:
            print(
                (
                    "API key not provided and BigQuery not enabled. Please set PEPY_API_KEY"
                    " environment variable, use --api-key, or enable BigQuery with --use-bigquery"
                ),
                file=sys.stderr,
            )
            sys.exit(1)
    elif use_bigquery and not credentials_path and not args.credentials_path:
        print(
            (
                "BigQuery enabled but no credentials provided. Please set"
                " USE_BIGQUERY to path of credentials file or use --credentials-path"
            ),
            file=sys.stderr,
        )
        # Continue execution - BigQuery might work without explicit credentials in some environments

    # Get PyPI downloads
    total_downloads = get_total_downloads(api_key, args.package, use_bigquery, credentials_path)
    print(f"Total downloads for {args.package}: {total_downloads:,}")

    # Get GitHub stars
    stars = get_github_stars(args.github_repo)
    if stars is not None:
        print(f"GitHub stars for {args.github_repo}: {stars:,}")

    # Get Aider contribution percentage in latest release
    percentage, version = get_latest_release_aider_percentage()
    print(f"Aider wrote {percentage:.2f}% of code in the LATEST release ({version})")


if __name__ == "__main__":
    main()