Files
2026-07-13 12:28:55 +08:00

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!"