llama-ping

#!/usr/bin/env bash

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &> /dev/null && pwd)"
source "$SCRIPT_DIR/utils.sh"

# Configuration
API_URL="http://0.0.0.0:1234/v1/chat/completions"
MODELS_URL="http://0.0.0.0:1234/v1/models"

fetch_models() {
    local response
    response=$(curl -s "$MODELS_URL")
    if [ -z "$response" ]; then
        die "Could not fetch models from server. Is it running?"
    fi

    log_step "Available models:"
    echo "$response" | jq -r '.data[] | "\(.id) \(.status.value) \(.status.args | join(" "))"' | while read -r id loaded args; do
        local context="N/A"
        if [ -n "$args" ]; then
            context=$(echo "$args" | grep -oP '(?<=\--ctx-size )\K\d+' | head -1)
        fi
        
        if [ "$loaded" = "loaded" ]; then
            echo -e "  ${COLOR_GREEN}${COLOR_BOLD}★ ${id} (LOADED)${COLOR_RESET} - Context: ${context:-N/A}"
        else
            echo -e "  - ${COLOR_BLUE}${id}${COLOR_RESET} (Context: ${context:-N/A})"
        fi
    done
}

get_payload() {
    local model_alias=$1
    local prompt=$2
    jq -n \
        --arg model "$model_alias" \
        --arg prompt "$prompt" \
        '{model: $model, messages: [{role: "user", content: $prompt}], temperature: 0}'
}

call_api() {
    local payload=$1
    curl -s -X POST "$API_URL" \
        -H "Content-Type: application/json" \
        -d "$payload"
}

handle_response() {
    local response=$1
    local prompt=$2

    if [ -z "$response" ]; then
        die "Empty response from server."
    fi

    local error_msg=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_msg" ]; then
        die "Server Error: $error_msg"
    fi

    local content=$(echo "$response" | jq -r '.choices[0].message.content')
    local prompt_n=$(echo "$response" | jq -r '.timings.prompt_n')
    # Trunca as métricas em 2 casas decimais já na origem (ponto como separador
    # decimal), evitando que consumidores como o llama-verify tenham que
    # reformatar valores longos sob locales com vírgula decimal.
    local prompt_tps=$(echo "$response" | jq -r '((.timings.prompt_per_second // 0) * 100 | floor / 100) | tostring')
    local pred_n=$(echo "$response" | jq -r '.timings.predicted_n')
    local pred_tps=$(echo "$response" | jq -r '((.timings.predicted_per_second // 0) * 100 | floor / 100) | tostring')

    # Note: Using echo here to avoid timestamp prefixes if any, keeping it raw
    echo "$content"
    echo
    log_info "Prompt: $prompt_n tokens @ $prompt_tps tokens/s"
    log_info "Generated: $pred_n tokens @ $pred_tps tokens/s"
}

main() {
    if [ -z "$1" ]; then
        fetch_models
        exit 0
    fi

    local model_alias=$1
    local prompt=${2:-"Hello"}

    local payload=$(get_payload "$model_alias" "$prompt")
    
    echo -e "${COLOR_YELLOW}User: ${prompt}${COLOR_RESET}"
    echo -ne "${COLOR_GREEN}  AI: ${COLOR_RESET}"
    
    local response=$(call_api "$payload")

    handle_response "$response" "$prompt"
}

main "$@"