pwshai.psm1
|
<# .SYNOPSIS Retrieve from OpenAI API to get a PowerShell command for a given query. .DESCRIPTION This module provides a function to retrieve a PowerShell command for a given query from OpenAI API. The provided function `ai` takes a query as input and returns the PowerShell command for the query. .PARAMETER query The query for which the PowerShell command is to be retrieved. .OUTPUTS System.String The PowerShell command for the given query. .EXAMPLE ai "list all files in current directory that ends with .txt" Retrieves the PowerShell command to list all files in the current directory that ends with .txt. .NOTES Author: Nicolo' Avanzini #> function ai { [CmdletBinding()] param ( [Parameter(Position=0, Mandatory=$true)] [string]$query ) $config = Update-Configurations $response = Get-ProviderResponse -config $config -query $query Write-Host "`r`n$response`r`n" -ForegroundColor Blue Get-UserChoice -message "Do you want to run the command? [Y/n]" -choices @("Y", "n") -default "Y" | ForEach-Object { if ($_ -eq "Y") { Invoke-Expression $response } } } enum Provider { OpenAI } function Update-Configurations { # Check if the configuration file exists in home directory $configFile = "$env:USERPROFILE\.aipwsh\config.json" if (-not (Test-Path $configFile)) { New-Item -Path $configFile -ItemType File -Force | Out-Null # Ask for API key $apiKey = Read-Host "Enter your OpenAI API key" $config = [Config]::new() $config.provider = "OpenAI" $config.apiKey = $apiKey $config.Save($configFile) return $config } # Load the configuration file $config = [Config]::new() $config.Load($configFile) return $config } function Get-ProviderResponse { param ( [Config]$config, [string]$query ) if ($config.apiKey -eq "") { Write-Error "API key is missing. Please run 'ai' command to update the configuration" Exit 1 } if ($config.prompt -eq "") { Write-Error "Prompt not provided." Exit 1 } switch ($config.provider) { OpenAI { # Call OpenAI API to retrieve the command asked by the user $request = $openAIPrompt.Replace("{user_query}", $query).Replace("{system_query}", $config.prompt) try { return Set-Loading -function { param($config, $request) $response = Invoke-RestMethod -Method Post -Uri "https://api.openai.com/v1/chat/completions" -Headers @{ "Authorization" = "Bearer $($config.apiKey)"; "Content-Type" = "application/json" } -Body $request return $response.choices[0].message.content | ConvertFrom-Json | Select-Object -ExpandProperty command } -arg $config, $request } catch { Write-Error "Failed to get response from OpenAI API: $($_.Exception.Message)" Exit 1 } } Default { Write-Error "Provider not supported" Exit 1 } } } function Set-Loading { param( [scriptblock]$function, [Object[]]$arg, [string]$Label) $job = Start-Job -ScriptBlock $function -ArgumentList $arg $symbols = @("⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷") $i = 0; while ($job.State -eq "Running") { $symbol = $symbols[$i] Write-Host -NoNewLine "`r$symbol $Label" -ForegroundColor Green Start-Sleep -Milliseconds 100 $i++ if ($i -eq $symbols.Count){ $i = 0; } } Write-Host -NoNewLine "`r " return Receive-Job $job } function Get-UserChoice { param ( [string]$message, [string[]]$choices, [string]$default ) do { $choice = Read-Host $message if ($choice -eq "" -and $default -ne "") { $choice = $default } } while ( $choices -notcontains $choice ) return $choice } class Config { [Provider]$provider [string]$apiKey [string]$prompt Config() { $this.provider = [Provider]::OpenAI $this.apiKey = "" $this.prompt = "You are a powershell expert. You are asked how to perform an action in powershell and you should provide a command. Try to be as concise as possible and try to always use one line command." } Load($configFile) { $cfg = Get-Content $configFile | ConvertFrom-Json $this.provider = $cfg.provider $this.apiKey = $cfg.apiKey $this.prompt = $cfg.prompt } Save($configFile) { $this | ConvertTo-Json | Set-Content $configFile } } [string]$openAIPrompt = @' { "model": "gpt-4o-2024-08-06", "messages": [ { "role": "system", "content": "{system_query}" }, { "role": "user", "content": "{user_query}" } ], "response_format": { "type": "json_schema", "json_schema": { "name": "powershell_command", "schema": { "type": "object", "properties": { "command": { "type": "string" } }, "required": ["command"], "additionalProperties": false }, "strict": true } } } '@ Export-ModuleMember -Function ai |