Jina cli bash script for multiple api keys

jina-cli.sh

#!/bin/bash

#######################
# Jina CLI Tool
# Save as ~/.local/bin/jina-cli.sh
#######################

#######################
# Configuration
#######################
JINA_CONFIG_DIR="$HOME/.config/jina"
JINA_STATS_FILE="$JINA_CONFIG_DIR/stats.json"
JINA_CACHE_FILE="$JINA_CONFIG_DIR/cache.json"
CACHE_TTL=3600  # Cache validity in seconds (1 hour)
declare -A JINA_KEY_USAGE

#######################
# Setup Functions
#######################
setup_directories() {
    mkdir -p "$JINA_CONFIG_DIR"
    
    # Initialize stats file if not exists
    if [ ! -f "$JINA_STATS_FILE" ]; then
        echo '{"current_key":1,"usage":{}}' > "$JINA_STATS_FILE"
    fi
    
    # Initialize cache file if not exists
    if [ ! -f "$JINA_CACHE_FILE" ]; then
        echo '{"last_update":"","balances":{}}' > "$JINA_CACHE_FILE"
    fi
}

check_dependencies() {
    local missing_deps=()
    
    if ! command -v jq >/dev/null 2>&1; then
        missing_deps+=("jq")
    fi
    if ! command -v xclip >/dev/null 2>&1; then
        missing_deps+=("xclip")
    fi
    
    if [ ${#missing_deps[@]} -ne 0 ]; then
        echo "Error: Missing dependencies: ${missing_deps[*]}"
        echo "Please install them with: sudo apt install ${missing_deps[*]}"
        return 1
    fi
    return 0
}

#######################
# API Key Management
#######################
count_api_keys() {
    local count=1
    while true; do
        local key_var="JINA_API_KEY_${count}"
        [ -z "${!key_var}" ] && break
        ((count++))
    done
    echo $((count - 1))
}

get_api_key() {
    local key_num=$1
    local key_var="JINA_API_KEY_${key_num}"
    echo "${!key_var}"
}

fetch_token_balance() {
    local api_key=$1
    local response
    
    response=$(curl -s "https://embeddings-dashboard-api.jina.ai/api/v1/api_key/user?api_key=${api_key}" \
        -H 'accept: */*' \
        -H 'origin: https://jina.ai' \
        -H 'referer: https://jina.ai/')
    
    echo "$response" | jq -r '.wallet.total_balance // 0'
}

#######################
# Cache Management
#######################
update_token_cache() {
    local current_time=$(date +%s)
    local last_update=$(jq -r '.last_update' "$JINA_CACHE_FILE")
    local force_update=${1:-false}
    
    if [ "$force_update" = true ] || [ -z "$last_update" ] || [ $((current_time - last_update)) -gt $CACHE_TTL ]; then
        local balances="{}"
        local key_count=$(count_api_keys)
        
        for ((i=1; i<=key_count; i++)); do
            local api_key=$(get_api_key $i)
            if [ -n "$api_key" ]; then
                local balance=$(fetch_token_balance "$api_key")
                balances=$(echo $balances | jq --arg key "$i" --arg balance "$balance" '. + {($key): ($balance)}')
            fi
        done
        
        echo "{\"last_update\":$current_time,\"balances\":$balances}" > "$JINA_CACHE_FILE"
    fi
}

get_sorted_keys() {
    update_token_cache
    # Convert string balances to numbers before sorting
    jq -r '.balances | to_entries | map({key: .key, value: (.value | tonumber)}) | sort_by(-.value) | map(.key)[]' "$JINA_CACHE_FILE"
}

#######################
# Usage Statistics
#######################
load_usage_stats() {
    if [ -f "$JINA_STATS_FILE" ]; then
        while IFS="=" read -r key value; do
            JINA_KEY_USAGE[$key]=$value
        done < <(jq -r '.usage | to_entries | .[] | "\(.key)=\(.value)"' "$JINA_STATS_FILE")
    fi
}

save_usage_stats() {
    local key=$1
    JINA_KEY_USAGE[$key]=$((${JINA_KEY_USAGE[$key]:-0} + 1))
    
    local usage_json="{"
    for k in "${!JINA_KEY_USAGE[@]}"; do
        usage_json+="\"$k\":${JINA_KEY_USAGE[$k]},"
    done
    usage_json="${usage_json%,}}"
    
    echo "{\"current_key\":$key,\"usage\":$usage_json}" > "$JINA_STATS_FILE"
}

reset_usage_stats() {
    declare -A JINA_KEY_USAGE
    echo '{"current_key":1,"usage":{}}' > "$JINA_STATS_FILE"
}

#######################
# Command Handlers
#######################
handle_status() {
    update_token_cache
    local key_count=$(count_api_keys)
    
    echo "Available API Keys: ${key_count}"
    echo "Usage Statistics and Balances:"
    
    while IFS= read -r key; do
        local usage=${JINA_KEY_USAGE[$key]:-0}
        local balance=$(jq -r ".balances.\"$key\"" "$JINA_CACHE_FILE")
        echo "  Key #$key: $usage calls, $balance tokens available"
    done < <(get_sorted_keys)
}

handle_reset() {
    reset_usage_stats
    echo "Usage statistics reset"
}

handle_update() {
    update_token_cache true
    echo "Token balance cache updated"
}

handle_request() {
    local url=$1
    local key_num=$(get_sorted_keys | head -n1)
    local api_key=$(get_api_key $key_num)
    
    if [ -z "$api_key" ]; then
        echo "Error: No valid API key found"
        return 1
    fi
    
    # Make request and copy to clipboard
    curl "https://r.jina.ai/$url" \
        -H "Authorization: Bearer ${api_key}" \
        -s | xclip -selection clipboard
    
    save_usage_stats $key_num
    echo "Result copied to clipboard (used key #${key_num})"
}

#######################
# Main Function
#######################
jina() {
    setup_directories
    check_dependencies || return 1
    load_usage_stats
    
    if [ -z "$1" ]; then
        echo "Usage: jina <url>"
        echo "Commands:"
        echo "  jina status    - Show key usage statistics"
        echo "  jina reset     - Reset usage statistics"
        echo "  jina update    - Force update token balance cache"
        return 1
    fi
    
    case "$1" in
        status)
            handle_status
            ;;
        reset)
            handle_reset
            ;;
        update)
            handle_update
            ;;
        *)
            handle_request "$1"
            ;;
    esac
}