src/AIPowerShellAssistant.ps1
|
# ========================================== # AI PowerShell Assistant # Ollama + Intelligent Model Selector # ========================================== $script:AssistantVersion = "1.3.7" $script:ConfigDir = "$HOME\.ai-powershell-assistant" $script:ConfigPath = Join-Path $script:ConfigDir "config.json" $script:ConfigLoadError = $null $script:Model = $null $script:OllamaUrl = "http://localhost:11434/api/generate" function Update-ScriptConfigFromDisk { if (Test-Path $script:ConfigPath) { try { $cfg = Get-Content $script:ConfigPath -Raw -Encoding UTF8 | ConvertFrom-Json if ($cfg.Model) { $script:Model = [string]$cfg.Model if ($cfg.OllamaUrl) { $script:OllamaUrl = [string]$cfg.OllamaUrl } $script:ConfigLoadError = $null } } catch { $script:ConfigLoadError = "Config load error" } } } Update-ScriptConfigFromDisk $script:LastProcessedHistoryId = -1 $script:LastErrorCount = $global:Error.Count $script:AgentModeEnabled = $false [Console]::OutputEncoding = [System.Text.Encoding]::UTF8 # ========================================== # PREREQUISITES & SETUP # ========================================== function Get-OllamaModelsSafe { $ollamaBase = "http://localhost:11434" if ($script:OllamaUrl) { $ollamaBase = ($script:OllamaUrl -replace '/api/.*$', '') } try { $tagsResponse = Invoke-RestMethod -Uri "$ollamaBase/api/tags" -Method GET -TimeoutSec 3 -ErrorAction Stop } catch { return $null } $modelList = [System.Collections.Generic.List[string]]::new() if ($tagsResponse.models) { foreach ($m in $tagsResponse.models) { if ($m.name) { $modelList.Add([string]$m.name) } elseif ($m.model) { $modelList.Add([string]$m.model) } } } return ,$modelList.ToArray() } function Initialize-AIConfiguration { param([switch]$Interactive) if (-not (Get-Command ollama -ErrorAction SilentlyContinue)) { Write-Host "" Write-Host "[ERROR] 'ollama' command was not found on this system!" -ForegroundColor Red Write-Host "Please download and install Ollama: https://ollama.com" -ForegroundColor Yellow Write-Host "" return $false } [string[]]$models = Get-OllamaModelsSafe if ($null -eq $models) { Write-Host "" Write-Host "[ERROR] Ollama server is not running or unreachable at http://localhost:11434!" -ForegroundColor Red Write-Host "Please start the Ollama application and try again." -ForegroundColor Yellow Write-Host "" return $false } if ($models.Length -eq 0) { Write-Host "" Write-Host "[ERROR] No installed models found in Ollama!" -ForegroundColor Red Write-Host "Please pull a model first, for example:" -ForegroundColor Yellow Write-Host " ollama pull qwen2.5-coder:1.5b" -ForegroundColor Cyan Write-Host "" return $false } Update-ScriptConfigFromDisk if (-not $Interactive -and $script:Model) { if ($models -contains $script:Model -or $models -contains "$($script:Model):latest") { return $true } } Write-Host "" Write-Host "=== AI PowerShell Assistant Setup ===" -ForegroundColor Cyan Write-Host "Available models found in your Ollama:" -ForegroundColor DarkGray for ($i = 0; $i -lt $models.Length; $i++) { Write-Host " [$($i+1)] $($models[$i])" -ForegroundColor Yellow } Write-Host "" $defaultIdx = 1 for ($i = 0; $i -lt $models.Length; $i++) { if ($models[$i] -like "*qwen2.5-coder*") { $defaultIdx = $i + 1; break; } } $choice = Read-Host "Select model number [1-$($models.Length)] (Default: [$defaultIdx] $($models[$defaultIdx-1]))" if ([string]::IsNullOrWhiteSpace($choice)) { $choiceIndex = $defaultIdx - 1 } else { $parsed = 0 if ([int]::TryParse($choice, [ref]$parsed) -and $parsed -ge 1 -and $parsed -le $models.Length) { $choiceIndex = $parsed - 1 } else { Write-Host "Invalid choice. Using default model: $($models[$defaultIdx-1])" -ForegroundColor DarkYellow $choiceIndex = $defaultIdx - 1 } } $selectedModel = [string]$models[$choiceIndex] if (-not (Test-Path $script:ConfigDir)) { New-Item -ItemType Directory -Path $script:ConfigDir -Force | Out-Null } $newConfig = @{ Model = $selectedModel OllamaUrl = "http://localhost:11434/api/generate" } $newConfig | ConvertTo-Json -Depth 2 | Set-Content -Path $script:ConfigPath -Encoding UTF8 $script:Model = $newConfig.Model $script:OllamaUrl = $newConfig.OllamaUrl $script:ConfigLoadError = $null Write-Host "[SUCCESS] Selected model: '$selectedModel'. Configuration saved!" -ForegroundColor Green Write-Host "" return $true } function Show-AISetupMenu { Initialize-AIConfiguration -Interactive } Set-Alias -Name AI-Setup -Value Show-AISetupMenu # ========================================== # INVOKE AI # ========================================== function Invoke-AI { param( [string]$Command, [string]$ErrorMessage ) $prompt = @" You are an expert PowerShell assistant. A PowerShell command failed. Command: $Command Error: $ErrorMessage Your task: Return ONLY the corrected PowerShell command. Rules: - No markdown - No code blocks - No explanation - Return one command only - Do not invent unnecessary commands - Preserve the user's original intention - Fix only the actual problem - Correct spelling mistakes - Correct PowerShell syntax - Correct invalid parameters and arguments - Correct external program arguments when appropriate "@ $body = @{ model = $script:Model prompt = $prompt stream = $false options = @{ temperature = 0 } } | ConvertTo-Json -Depth 5 try { $response = Invoke-RestMethod ` -Uri $script:OllamaUrl ` -Method POST ` -ContentType "application/json" ` -Body $body ` -ErrorAction Stop if (-not $response.response) { return $null } $fix = $response.response.Trim() # REMOVE MARKDOWN $fix = $fix -replace '```powershell', '' $fix = $fix -replace '```PowerShell', '' $fix = $fix -replace '```ps1', '' $fix = $fix -replace '```PS1', '' $fix = $fix -replace '```', '' $fix = $fix.Trim() # REMOVE ACCIDENTAL PREFIX $fix = $fix -replace '^Suggested fix:\s*', '' $fix = $fix -replace '^Corrected command:\s*', '' $fix = $fix -replace '^Command:\s*', '' $fix = $fix.Trim() if ([string]::IsNullOrWhiteSpace($fix)) { return $null } return $fix } catch { if ($global:Error.Count -gt 0) { $global:Error.RemoveAt(0) } return $null } } # ========================================== # GET LAST HISTORY # ========================================== function Get-LastHistory { try { return Get-History -Count 1 } catch { return $null } } # ========================================== # FIND ERROR FOR CURRENT COMMAND # ========================================== function Get-CurrentCommandError { $currentCount = $global:Error.Count if ($currentCount -le $script:LastErrorCount) { $script:LastErrorCount = $currentCount return $null } $script:LastErrorCount = $currentCount return $global:Error[0] } # ========================================== # AI ERROR CHECK # ========================================== function Invoke-AIErrorCheck { if (-not $script:AgentModeEnabled) { return } if ($script:ConfigLoadError -or -not $script:Model) { return } $currentError = Get-CurrentCommandError if (-not $currentError) { return } $history = Get-LastHistory $command = $null if ($history -and $history.Id -ne $script:LastProcessedHistoryId) { $command = $history.CommandLine $script:LastProcessedHistoryId = $history.Id } if ([string]::IsNullOrWhiteSpace($command)) { try { $command = $currentError.InvocationInfo.Line } catch { } } if ([string]::IsNullOrWhiteSpace($command)) { $command = "(unknown command)" } $errorText = $currentError.ToString() Write-Host "" Write-Host "AI is analyzing the error..." -ForegroundColor Cyan $fix = Invoke-AI -Command $command -ErrorMessage $errorText if (-not $fix) { Write-Host "" Write-Host "Could not get a response from AI (check if Ollama is running)." -ForegroundColor Red Write-Host "" return } if ($fix.Trim() -eq $command.Trim()) { Write-Host "" Write-Host "AI could not find a better command." -ForegroundColor DarkGray Write-Host "" return } Write-Host "" Write-Host "Suggested fix:" -ForegroundColor Green Write-Host $fix -ForegroundColor Yellow Write-Host "" } # ========================================== # AGENT MODE ON / OFF / STATUS # ========================================== function Start-AIAssistant { if (-not (Initialize-AIConfiguration)) { return } $script:AgentModeEnabled = $true $script:LastErrorCount = $global:Error.Count Write-Host "" Write-Host "AI Assistant: ON (Model: $script:Model)" -ForegroundColor Green Write-Host "Errors will now be checked automatically." -ForegroundColor DarkGray Write-Host "" } Set-Alias -Name AI-On -Value Start-AIAssistant function Stop-AIAssistant { $script:AgentModeEnabled = $false Write-Host "" Write-Host "AI Assistant: OFF" -ForegroundColor Yellow Write-Host "" } Set-Alias -Name AI-Off -Value Stop-AIAssistant function Get-AIAssistantStatus { Update-ScriptConfigFromDisk Write-Host "" Write-Host "--- AI PowerShell Assistant (v$script:AssistantVersion) ---" -ForegroundColor Cyan Write-Host "Status: " -NoNewline if ($script:AgentModeEnabled) { Write-Host "ON" -ForegroundColor Green } else { Write-Host "OFF" -ForegroundColor Yellow } Write-Host "Model: $script:Model" Write-Host "OllamaUrl: $script:OllamaUrl" Write-Host "Config: $script:ConfigPath" Write-Host "----------------------------------------------" -ForegroundColor Cyan Write-Host "" } Set-Alias -Name AI-PSA -Value Get-AIAssistantStatus # ========================================== # POWERSHELL PROMPT # ========================================== function global:prompt { try { Invoke-AIErrorCheck } catch { } return "PS $($executionContext.SessionState.Path.CurrentLocation)> " } # ========================================== # STARTUP # ========================================== if ($script:ConfigLoadError) { Write-Host "" Write-Host "AI Assistant: config error ($script:ConfigLoadError)" -ForegroundColor Red Write-Host "" } else { Write-Host "" Write-Host "AI Assistant: OFF (type AI-On to enable, AI-PSA for status, AI-Setup to reconfigure)" -ForegroundColor DarkGray Write-Host "" } |