Private/Invoke-OpenAICall.ps1
|
function Invoke-OpenAICall { <# .SYNOPSIS Make an API call to OpenAI/Azure OpenAI with rate limiting and retry logic .DESCRIPTION Executes an API call to the configured OpenAI endpoint with automatic retry on rate limits #> [CmdletBinding()] param( [Parameter(Mandatory = $true)] [string]$Prompt, [Parameter(Mandatory = $true)] [PSCustomObject]$AIConfig, [Parameter(Mandatory = $false)] [int]$MaxRetries = 5, [Parameter(Mandatory = $false)] [int]$BaseDelaySeconds = 2 ) $retryCount = 0 $lastError = $null while ($retryCount -le $MaxRetries) { try { # Determine if it's Azure OpenAI or OpenAI $isAzure = $AIConfig.endpoint -like "*azure.com*" -or $AIConfig.endpoint -like "*openai.azure.com*" if ($isAzure) { # Azure OpenAI - use latest stable API version $apiVersion = "2024-08-01-preview" $deploymentName = $AIConfig.model # Extract resource name from endpoint $uri = "$($AIConfig.endpoint.TrimEnd('/'))/openai/deployments/$deploymentName/chat/completions?api-version=$apiVersion" $headers = @{ "api-key" = $AIConfig.apiKey "Content-Type" = "application/json" } } else { # Standard OpenAI $uri = "$($AIConfig.endpoint.TrimEnd('/'))/v1/chat/completions" $headers = @{ "Authorization" = "Bearer $($AIConfig.apiKey)" "Content-Type" = "application/json" } } $body = @{ model = $AIConfig.model messages = @( @{ role = "system" content = "You are a security expert specialized in Microsoft 365 Conditional Access policies and the Baseline Secure Cloud." }, @{ role = "user" content = $Prompt } ) temperature = 0.7 max_tokens = 4000 } | ConvertTo-Json -Depth 10 $response = Invoke-RestMethod -Uri $uri -Method Post -Headers $headers -Body $body -ErrorAction Stop if ($response.choices -and $response.choices.Count -gt 0) { return $response.choices[0].message.content } else { Write-Warning "No response from AI" return $null } } catch { $lastError = $_ $errorDetails = $_.Exception.Message # Check if it's a rate limit error $isRateLimit = $false $retryAfter = $BaseDelaySeconds if ($errorDetails -match "RateLimitReached" -or $errorDetails -match "429" -or $_.Exception.Response.StatusCode -eq 429) { $isRateLimit = $true # Try to extract retry-after time from error message if ($errorDetails -match "retry after (\d+) seconds?") { $retryAfter = [int]$matches[1] # Add buffer $retryAfter = [Math]::Max($retryAfter + 2, $BaseDelaySeconds) } elseif ($_.Exception.Response.Headers -and $_.Exception.Response.Headers['Retry-After']) { $retryAfter = [int]$_.Exception.Response.Headers['Retry-After'].Value $retryAfter = [Math]::Max($retryAfter + 2, $BaseDelaySeconds) } } # Check for other retryable errors $isRetryable = $isRateLimit -or ($errorDetails -match "503" -or $_.Exception.Response.StatusCode -eq 503) -or ($errorDetails -match "500" -or $_.Exception.Response.StatusCode -eq 500) -or ($errorDetails -match "502" -or $_.Exception.Response.StatusCode -eq 502) if ($isRetryable -and $retryCount -lt $MaxRetries) { $retryCount++ # More aggressive exponential backoff for rate limits if ($isRateLimit) { $exponentialDelay = $BaseDelaySeconds * [Math]::Pow(2, $retryCount - 1) # For rate limits, use longer delays: min 10 seconds, max 120 seconds $actualDelay = [Math]::Max([Math]::Max($retryAfter + 5, $exponentialDelay), 10) $actualDelay = [Math]::Min($actualDelay, 120) } else { # For other errors, use standard exponential backoff $exponentialDelay = $BaseDelaySeconds * [Math]::Pow(2, $retryCount - 1) $actualDelay = [Math]::Max($retryAfter, $exponentialDelay) } Write-Host " Rate limit or temporary error encountered. Retrying in $actualDelay seconds... (Attempt $retryCount/$MaxRetries)" -ForegroundColor Yellow Start-Sleep -Seconds $actualDelay continue } else { # Not retryable or max retries reached if ($isRateLimit) { Write-Error "Rate limit exceeded after $retryCount retries. Please wait and try again later, or consider using gpt-4o which has higher rate limits." } else { Write-Error "Error in OpenAI API call: $errorDetails" } return $null } } } # If we get here, all retries failed Write-Error "Failed to complete OpenAI API call after $MaxRetries retries. Last error: $lastError" return $null } |