246 lines
7.7 KiB
Bash
Executable File
246 lines
7.7 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
#
|
|
# Generate model TOML files for Cloudflare AI Gateway
|
|
#
|
|
# This script reads the API response from data/api_response.json and generates
|
|
# TOML files for each model.
|
|
#
|
|
# Usage:
|
|
# ./generate_model_toml.sh
|
|
#
|
|
|
|
set -eo pipefail
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
PROVIDER_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
|
|
source "${SCRIPT_DIR}/utils.sh"
|
|
|
|
# =============================================================================
|
|
# Step 3: Generate TOML files
|
|
# =============================================================================
|
|
|
|
echo ""
|
|
echo "=== Step 3: Generating model TOML files ==="
|
|
|
|
# =============================================================================
|
|
# Main script
|
|
# =============================================================================
|
|
MODELS_DIR="${PROVIDER_DIR}/models"
|
|
DATA_DIR="${PROVIDER_DIR}/data"
|
|
PROVIDERS_DIR="${PROVIDER_DIR}/.."
|
|
MODEL_NAMES_FILE="${DATA_DIR}/model_names.json"
|
|
MODEL_FAMILIES_FILE="${DATA_DIR}/model_families.json"
|
|
API_RESPONSE_FILE="${DATA_DIR}/api_response.json"
|
|
|
|
# Check if API response file exists
|
|
if [[ ! -f "${API_RESPONSE_FILE}" ]]; then
|
|
echo "Error: API response file not found at ${API_RESPONSE_FILE}" >&2
|
|
echo "Run generate_models.sh first to fetch the API response." >&2
|
|
exit 1
|
|
fi
|
|
|
|
# Check if the response is valid JSON with data
|
|
if ! jq -e '.data' "${API_RESPONSE_FILE}" > /dev/null 2>&1; then
|
|
echo "Error: Invalid API response or no data in ${API_RESPONSE_FILE}" >&2
|
|
exit 1
|
|
fi
|
|
|
|
MODEL_COUNT=$(jq '.data | length' "${API_RESPONSE_FILE}")
|
|
|
|
if [[ "${MODEL_COUNT}" -eq 0 ]]; then
|
|
echo "Error: No models found in API response" >&2
|
|
exit 1
|
|
fi
|
|
|
|
echo "Found ${MODEL_COUNT} models in API response"
|
|
|
|
# Create a temporary file to track API model files
|
|
API_MODEL_FILES=$(mktemp)
|
|
trap "rm -f ${API_MODEL_FILES}" EXIT
|
|
|
|
GENERATED_COUNT=0
|
|
SKIPPED_COUNT=0
|
|
CROSS_REF_COUNT=0
|
|
|
|
# =============================================================================
|
|
# Helper function to get model family from mapping file
|
|
# =============================================================================
|
|
get_model_family() {
|
|
local model_id="$1"
|
|
|
|
# Return empty if families file doesn't exist
|
|
if [[ ! -f "${MODEL_FAMILIES_FILE}" ]]; then
|
|
echo ""
|
|
return
|
|
fi
|
|
|
|
# Try to match patterns in order
|
|
local family
|
|
family=$(jq -r --arg id "${model_id}" '
|
|
.patterns[] | .pattern as $p | .family as $f | select($id | test($p)) | $f
|
|
' "${MODEL_FAMILIES_FILE}" | head -n 1)
|
|
|
|
# Return family if found, empty otherwise
|
|
if [[ -n "${family}" && "${family}" != "null" ]]; then
|
|
echo "${family}"
|
|
else
|
|
echo ""
|
|
fi
|
|
}
|
|
|
|
# Process each model from the API response
|
|
while IFS= read -r MODEL_JSON; do
|
|
MODEL_ID=$(echo "${MODEL_JSON}" | jq -r '.id')
|
|
COST_IN=$(echo "${MODEL_JSON}" | jq -r '.cost_in // 0')
|
|
COST_OUT=$(echo "${MODEL_JSON}" | jq -r '.cost_out // 0')
|
|
CREATED_AT=$(echo "${MODEL_JSON}" | jq -r '.created_at // 0')
|
|
|
|
# Skip empty IDs
|
|
[[ -z "${MODEL_ID}" || "${MODEL_ID}" == "null" ]] && continue
|
|
|
|
# Check if this model should be included
|
|
if ! should_include_model "${MODEL_ID}"; then
|
|
((SKIPPED_COUNT++)) || true
|
|
echo "Skipping: ${MODEL_ID}"
|
|
continue
|
|
fi
|
|
|
|
((INCLUDED_COUNT++)) || true
|
|
|
|
# Extract provider and model name
|
|
PROVIDER=$(echo "${MODEL_ID}" | cut -d'/' -f1)
|
|
MODEL_NAME=$(echo "${MODEL_ID}" | cut -d'/' -f2-)
|
|
MODEL_PATH="${MODEL_ID}.toml"
|
|
|
|
FULL_PATH="${MODELS_DIR}/${MODEL_PATH}"
|
|
echo "${FULL_PATH}" >> "${API_MODEL_FILES}"
|
|
|
|
# Create directory if needed
|
|
MODEL_DIR=$(dirname "${FULL_PATH}")
|
|
mkdir -p "${MODEL_DIR}"
|
|
|
|
# Check if we should cross-reference from source provider
|
|
SOURCE_FILE=$(find_source_file "${PROVIDER}" "${MODEL_NAME}" "${PROVIDERS_DIR}" || true);
|
|
|
|
if [[ -n "${SOURCE_FILE}" && -f "${SOURCE_FILE}" ]]; then
|
|
echo "Cross-referencing: ${MODEL_PATH} <- ${SOURCE_FILE#${PROVIDERS_DIR}/}"
|
|
cp "${SOURCE_FILE}" "${FULL_PATH}"
|
|
((CROSS_REF_COUNT++)) || true
|
|
else
|
|
# Generate file with defaults for workers-ai, replicate, etc.
|
|
echo "Generating: ${MODEL_PATH}"
|
|
((GENERATED_COUNT++)) || true
|
|
|
|
# Get display name from mapping file, or add model ID to file and use as-is
|
|
DISPLAY_NAME=$(jq -r --arg id "${MODEL_ID}" '.[$id] // empty' "${MODEL_NAMES_FILE}")
|
|
if [[ -z "${DISPLAY_NAME}" ]]; then
|
|
# Add model ID to model_names.json with ID as placeholder
|
|
TMP_FILE=$(mktemp)
|
|
jq --arg id "${MODEL_ID}" '.[$id] = $id' "${MODEL_NAMES_FILE}" > "${TMP_FILE}"
|
|
mv "${TMP_FILE}" "${MODEL_NAMES_FILE}"
|
|
DISPLAY_NAME="${MODEL_ID}"
|
|
fi
|
|
# Escape double quotes for TOML string
|
|
DISPLAY_NAME="${DISPLAY_NAME//\"/\\\"}"
|
|
|
|
# Get model family from mapping file
|
|
MODEL_FAMILY=$(get_model_family "${MODEL_ID}")
|
|
|
|
# Convert created_at timestamp to date (YYYY-MM-DD)
|
|
if [[ "${CREATED_AT}" != "0" && "${CREATED_AT}" != "null" ]]; then
|
|
RELEASE_DATE=$(date -r "${CREATED_AT}" +%Y-%m-%d 2>/dev/null || date -d "@${CREATED_AT}" +%Y-%m-%d 2>/dev/null || date +%Y-%m-%d)
|
|
else
|
|
RELEASE_DATE=$(date +%Y-%m-%d)
|
|
fi
|
|
|
|
# Convert cost per token to cost per million tokens
|
|
# API returns cost per token, we need cost per 1M tokens
|
|
# Treat negative or invalid costs as 0
|
|
# Note: jq outputs scientific notation with uppercase 'E', but bc requires lowercase 'e'
|
|
if [[ "${COST_IN}" != "0" && "${COST_IN}" != "null" ]]; then
|
|
COST_IN_BC=$(echo "${COST_IN}" | tr 'E' 'e')
|
|
COST_IN_PER_M=$(echo "${COST_IN_BC} * 1000000" | bc -l | sed 's/^\./0./' | sed 's/0*$//' | sed 's/\.$//')
|
|
# If negative, set to 0
|
|
if (( $(echo "${COST_IN_PER_M} < 0" | bc -l) )); then
|
|
COST_IN_PER_M="0"
|
|
fi
|
|
else
|
|
COST_IN_PER_M="0"
|
|
fi
|
|
|
|
if [[ "${COST_OUT}" != "0" && "${COST_OUT}" != "null" ]]; then
|
|
COST_OUT_BC=$(echo "${COST_OUT}" | tr 'E' 'e')
|
|
COST_OUT_PER_M=$(echo "${COST_OUT_BC} * 1000000" | bc -l | sed 's/^\./0./' | sed 's/0*$//' | sed 's/\.$//')
|
|
# If negative, set to 0
|
|
if (( $(echo "${COST_OUT_PER_M} < 0" | bc -l) )); then
|
|
COST_OUT_PER_M="0"
|
|
fi
|
|
else
|
|
COST_OUT_PER_M="0"
|
|
fi
|
|
|
|
# Always overwrite to ensure data is up to date
|
|
# Build the TOML content with optional family field
|
|
TOML_CONTENT="name = \"${DISPLAY_NAME}\""
|
|
if [[ -n "${MODEL_FAMILY}" ]]; then
|
|
TOML_CONTENT="${TOML_CONTENT}
|
|
family = \"${MODEL_FAMILY}\""
|
|
fi
|
|
TOML_CONTENT="${TOML_CONTENT}
|
|
release_date = \"${RELEASE_DATE}\"
|
|
last_updated = \"${RELEASE_DATE}\"
|
|
attachment = false
|
|
reasoning = false
|
|
temperature = true
|
|
tool_call = false
|
|
open_weights = false
|
|
|
|
[cost]
|
|
input = ${COST_IN_PER_M}
|
|
output = ${COST_OUT_PER_M}
|
|
|
|
[limit]
|
|
context = 128000
|
|
output = 16384
|
|
|
|
[modalities]
|
|
input = [\"text\"]
|
|
output = [\"text\"]"
|
|
|
|
echo "${TOML_CONTENT}" > "${FULL_PATH}"
|
|
fi
|
|
done < <(jq -c '.data[]' "${API_RESPONSE_FILE}")
|
|
|
|
# Find and remove models that are not in the API response
|
|
echo ""
|
|
echo "Checking for models to remove..."
|
|
REMOVED_COUNT=0
|
|
|
|
# Find all existing .toml files in models directory
|
|
if [[ -d "${MODELS_DIR}" ]]; then
|
|
while IFS= read -r -d '' EXISTING_FILE; do
|
|
if ! grep -qxF "${EXISTING_FILE}" "${API_MODEL_FILES}"; then
|
|
REL_PATH="${EXISTING_FILE#${MODELS_DIR}/}"
|
|
echo "Removing model not in API: ${REL_PATH}"
|
|
rm -f "${EXISTING_FILE}"
|
|
((REMOVED_COUNT++)) || true
|
|
fi
|
|
done < <(find "${MODELS_DIR}" -name "*.toml" -type f -print0)
|
|
fi
|
|
|
|
# Clean up empty directories
|
|
find "${MODELS_DIR}" -type d -empty -delete 2>/dev/null || true
|
|
|
|
FINAL_COUNT=$(find "${MODELS_DIR}" -name "*.toml" -type f | wc -l | tr -d ' ')
|
|
|
|
echo ""
|
|
echo "Summary:"
|
|
echo " Models from API: ${MODEL_COUNT}"
|
|
echo " Models skipped (filtered): ${SKIPPED_COUNT}"
|
|
echo " Models cross-referenced: ${CROSS_REF_COUNT}"
|
|
echo " Models generated: ${GENERATED_COUNT}"
|
|
echo " Models removed: ${REMOVED_COUNT}"
|
|
echo " Total models: ${FINAL_COUNT}"
|
|
echo ""
|
|
echo "Done!"
|