1<#
2.SYNOPSIS
3    Queries available capacity for an Azure OpenAI model and validates if a target is achievable.
4.PARAMETER ModelName
5    The model name (e.g., "gpt-4o", "o3-mini")
6.PARAMETER ModelVersion
7    The model version (e.g., "2025-01-31"). If omitted, lists available versions.
8.PARAMETER Region
9    Optional. Check capacity in a specific region only.
10.PARAMETER SKU
11    SKU to check (default: GlobalStandard)
12.EXAMPLE
13    .\query_capacity.ps1 -ModelName o3-mini
14    .\query_capacity.ps1 -ModelName o3-mini -ModelVersion 2025-01-31 -Region eastus2
15#>
16param(
17    [Parameter(Mandatory)][string]$ModelName,
18    [string]$ModelVersion,
19    [string]$Region,
20    [string]$SKU = "GlobalStandard"
21)
22 
23$ErrorActionPreference = "Stop"
24 
25$subId = az account show --query id -o tsv
26 
27# If no version provided, list available versions first
28if (-not $ModelVersion) {
29    Write-Host "Available versions for $ModelName`:"
30    $loc = if ($Region) { $Region } else { "eastus" }
31    az cognitiveservices model list --location $loc `
32        --query "[?model.name=='$ModelName'].{Version:model.version, Format:model.format}" `
33        --output table 2>$null
34    return
35}
36 
37# Build URL parameters
38$urlParams = @("api-version=2024-10-01", "modelFormat=OpenAI", "modelName=$ModelName", "modelVersion=$ModelVersion")
39 
40if ($Region) {
41    $url = "https://management.azure.com/subscriptions/$subId/providers/Microsoft.CognitiveServices/locations/$Region/modelCapacities"
42} else {
43    $url = "https://management.azure.com/subscriptions/$subId/providers/Microsoft.CognitiveServices/modelCapacities"
44}
45 
46$raw = az rest --method GET --url $url --url-parameters @urlParams 2>$null | Out-String | ConvertFrom-Json
47 
48# Filter by SKU
49$filtered = $raw.value | Where-Object { $_.properties.skuName -eq $SKU -and $_.properties.availableCapacity -gt 0 }
50 
51if (-not $filtered) {
52    Write-Host "No capacity found for $ModelName v$ModelVersion ($SKU)" -ForegroundColor Red
53    Write-Host "Try a different SKU or version."
54    return
55}
56 
57Write-Host "Capacity: $ModelName v$ModelVersion ($SKU)"
58Write-Host ""
59$filtered | ForEach-Object {
60    # Check subscription quota for this region
61    $quotaDisplay = "?"
62    try {
63        $usageData = az cognitiveservices usage list --location $_.location --subscription $subId -o json 2>$null | Out-String | ConvertFrom-Json
64        $usageEntry = $usageData | Where-Object { $_.name.value -eq "OpenAI.$SKU.$ModelName" }
65        if ($usageEntry) {
66            $quotaAvail = [int]$usageEntry.limit - [int]$usageEntry.currentValue
67            $quotaDisplay = if ($quotaAvail -gt 0) { "${quotaAvail}K" } else { "0 (at limit)" }
68        } else {
69            $quotaDisplay = "0 (none)"
70        }
71    } catch {
72        $quotaDisplay = "?"
73    }
74    [PSCustomObject]@{
75        Region    = $_.location
76        SKU       = $_.properties.skuName
77        Available = "$($_.properties.availableCapacity)K TPM"
78        Quota     = $quotaDisplay
79    }
80} | Sort-Object { [int]($_.Available -replace '[^\d]','') } -Descending | Format-Table -AutoSize
81