data_processing.py

# data_processing.py

import os
import time
import json
import requests
from requests.adapters import HTTPAdapter
import logging
import functools
from urllib.parse import unquote
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
import numpy as np
import requests
from datasets import load_dataset

# =====================
# Configuration
# =====================
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

REPO_OWNER = "JonathanChavezTamales"
REPO_NAME = "LLMStats"
BRANCH = "main"
MODELS_PATH = "models/"
PROVIDERS_PATH = "providers/"

# Use an environment variable for the GitHub token to avoid rate limiting
GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
HEADERS = {'Authorization': f'token {GITHUB_TOKEN}'} if GITHUB_TOKEN else {}

# =====================
# Caching
# =====================
def cache_data(ttl_seconds=3600):
    """Simple in-memory cache decorator with a Time-To-Live (TTL)."""
    def decorator(func):
        cache = {}
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            key = f"{func.__name__}:{args}:{kwargs}"
            now = time.time()
            if key in cache and (now - cache[key]['timestamp']) < ttl_seconds:
                logger.info(f"Returning cached data for '{func.__name__}'.")
                return cache[key]['data']
            
            logger.info(f"Cache miss or expired. Fetching new data for '{func.__name__}'.")
            result = func(*args, **kwargs)
            cache[key] = {'data': result, 'timestamp': now}
            return result
        return wrapper
    return decorator

# =====================
# Data Fetching
# =====================

def _fetch_json_from_github(session, url):
    """Helper to fetch and parse a JSON file from a URL."""
    try:
        response = session.get(url, headers=HEADERS)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        logger.error(f"HTTP Error fetching {url}: {e}")
    except json.JSONDecodeError:
        logger.error(f"Failed to decode JSON from {url}")
    return None

def _process_model_item(item, session):
    """Processes a single model JSON file from the GitHub repo."""
    raw_url = f"https://raw.githubusercontent.com/{REPO_OWNER}/{REPO_NAME}/{BRANCH}/{item['path']}"
    json_data = _fetch_json_from_github(session, raw_url)
    if not json_data:
        return None

    path_parts = unquote(item['path']).split('/')
    organization = path_parts[1] if len(path_parts) > 1 else 'unknown'

    record = {
        "organization": organization.lower(),
        "name": json_data.get("name", "").lower(),
        "description": json_data.get("description"),
        "release_date": pd.to_datetime(json_data.get("release_date"), errors='coerce'),
        "input_context_size": json_data.get("input_context_size"),
        "output_context_size": json_data.get("output_context_size"),
        "license": json_data.get("license"),
        "GPQA_score": None # Default value
    }

    if "qualitative_metrics" in json_data:
        for metric in json_data["qualitative_metrics"]:
            if metric.get("dataset_name") == "GPQA":
                record["GPQA_score"] = metric.get("score")
                break # Found the one we care about for the timeline
    
    return record

def _process_provider_item(item, session):
    """Processes a single provider JSON file from the GitHub repo."""
    raw_url = f"https://raw.githubusercontent.com/{REPO_OWNER}/{REPO_NAME}/{BRANCH}/{item['path']}"
    provider_data = _fetch_json_from_github(session, raw_url)
    if not provider_data:
        return []

    provider_name = provider_data.get("name", "Unknown").lower()
    models_list = []
    for model in provider_data.get("providermodels", []):
        try:
            input_price = float(model.get("price_per_input_token", 0))
            models_list.append({
                "provider": provider_name,
                "model_id": model.get("model_id", "").lower(),
                "cost_per_million": round(input_price * 1_000_000, 2)
            })
        except (ValueError, TypeError):
            logger.warning(f"Could not process model cost for {model.get('model_id')} in {provider_name}")
    return models_list

def _fetch_github_tree():
    """Fetches the full recursive file tree for the repository."""
    api_url = f"https://api.github.com/repos/{REPO_OWNER}/{REPO_NAME}/git/trees/{BRANCH}?recursive=1"
    with requests.Session() as session:
        tree_data = _fetch_json_from_github(session, api_url)
    return tree_data.get("tree", []) if tree_data else []

@cache_data()
def fetch_all_github_data():
    """Fetches and processes all model and provider data from GitHub in parallel."""
    tree = _fetch_github_tree()
    if not tree:
        logger.warning("GitHub repository tree is empty. Cannot fetch data.")
        return pd.DataFrame(), pd.DataFrame()

    all_models_data = []
    all_costs_data = []

    # === Configure the HTTP Adapter ===
    # Set the number of workers for our thread pool
    MAX_WORKERS = 20
    
    # Create an adapter with a connection pool size that matches our worker count
    adapter = HTTPAdapter(pool_connections=MAX_WORKERS, pool_maxsize=MAX_WORKERS)
    
    # Create a session and mount the adapter
    session = requests.Session()
    session.mount('https://', adapter)
    # =======================================
    try:
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # We no longer need to pass the session into the executor, 
            # as it's created and managed within this function scope.
            
            # Submit model file processing tasks
            model_futures = [executor.submit(_process_model_item, item, session) 
                             for item in tree if item['path'].startswith(MODELS_PATH) and item['path'].endswith('.json')]
            # Submit provider file processing tasks
            provider_futures = [executor.submit(_process_provider_item, item, session) 
                                for item in tree if item['path'].startswith(PROVIDERS_PATH) and item['path'].endswith('.json')]

            for future in model_futures:
                result = future.result()
                if result:
                    all_models_data.append(result)
            
            for future in provider_futures:
                results = future.result()
                if results:
                    all_costs_data.extend(results)
    finally:
        # Ensure the session is closed to release all connections
        session.close()

    models_df = pd.DataFrame(all_models_data).dropna(subset=['name', 'release_date'])
    costs_df = pd.DataFrame(all_costs_data)
    
    logger.info(f"Successfully fetched {len(models_df)} models and {len(costs_df)} cost entries from GitHub.")
    return models_df, costs_df


@cache_data()
def load_leaderboard_data():
    """Loads and cleans the Open LLM Leaderboard data."""
    logger.info("Loading leaderboard data...")
    try:
        dataset = load_dataset("open-llm-leaderboard/contents")
        df = dataset["train"].to_pandas()
        
        # Filter for major organizations
        orgs = ['openai', 'anthropic', 'google', 'meta', 'qwen', 'mistral', 'deepseek', 'xai', 'cohere']
        df['organization'] = df['fullname'].str.lower().str.split('/').str[0]
        df = df[df['organization'].apply(lambda x: any(org in x for org in orgs))]
        
        # Ensure numeric types and drop rows with missing essential data
        metrics = ['BBH', 'MMLU-PRO', 'IFEval', 'MUSR', 'GPQA', 'MATH Lvl 5']
        df['#Params (B)'] = pd.to_numeric(df['#Params (B)'], errors='coerce')
        return df.dropna(subset=['Average ⬆️'] + metrics)
    except Exception as e:
        logger.error(f"Failed to load leaderboard data: {e}")
        return pd.DataFrame()

# =====================
# Data Processing/Merging
# =====================

def get_prepared_data():
    """Main function to fetch, process, and merge all data sources."""
    models_df, costs_df = fetch_all_github_data()
    leaderboard_df = load_leaderboard_data()

    # --- Prepare data for GPQA Timeline & Context Chart ---
    context_df = models_df.copy()
    context_df['input_context_size'] = context_df['input_context_size'].fillna(0)
    context_df['output_context_size'] = context_df['output_context_size'].fillna(0)
    context_df['total_context'] = context_df['input_context_size'] + context_df['output_context_size']
    context_df['prefix'] = context_df['name'].apply(
        lambda x: x.split(' ')[0] if ' ' in x else x.split('-')[0]
    )
    context_df = context_df.sort_values('total_context', ascending=False)
    
    # --- Prepare data for Value Analysis (Performance vs Cost) ---
    # The merge key is tricky. We'll merge on organization and a cleaned model name.
    # This is inherently fragile and a major challenge in real-world data integration.
    models_df['merge_key'] = models_df['name'].str.replace('-', '').str.replace(' ', '')
    costs_df['merge_key'] = costs_df['model_id'].str.replace('-', '').str.replace(' ', '')

    merged_df = pd.merge(
        models_df,
        costs_df,
        left_on=['organization', 'merge_key'],
        right_on=['provider', 'merge_key'],
        how='inner'
    ).drop(columns=['merge_key'])

    merged_df = merged_df.dropna(subset=['GPQA_score', 'cost_per_million'])
    
    # Avoid division by zero
    merged_df = merged_df[merged_df['cost_per_million'] > 0]
    if not merged_df.empty:
        merged_df['value_ratio'] = merged_df['GPQA_score'] / merged_df['cost_per_million']

    return {
        'timeline': models_df.dropna(subset=['GPQA_score']),
        'context': context_df,
        'value': merged_df,
        'leaderboard': leaderboard_df
    }