PowerShellAssistant.psm1
|
using namespace OpenAI using namespace System.Net.Http using namespace System.Net.Http.Headers using namespace System.Collections.Generic using namespace System.Management.Automation using namespace System.Reflection using namespace Microsoft.PowerShell $ErrorActionPreference = 'Stop' #TODO: This should be better $debugBinPath = Join-Path $PSScriptRoot '/bin/Debug/net7.0' if (Test-Path $debugBinPath) { Write-Warning "Debug build detected. Using assemblies at $debugBinPath" Add-Type -Path $debugBinPath/*.dll } else { Add-Type -Path $PSScriptRoot/*.dll } #These are the cheapest models for testing, opt into more powerful models $SCRIPT:aiDefaultModel = 'ada' $SCRIPT:aiDefaultChatModel = 'gpt-3.5-turbo' $SCRIPT:aiDefaultCodeModel = 'code-davinci-002' #region Public function Connect-AI { [CmdletBinding()] param( # Provide your API Key as the password, and optionally your organization ID as the username [string]$APIKey, # By default, this uses the OpenAI API. Specify this if you want to use GitHub Copilot (UNSUPPORTED) [switch]$GitHubCopilot, # Don't set this client as the default client. You can pass the client to the various commands instead. Implies -PassThru [switch]$NoDefault, # Return the client for use in other commands [switch]$PassThru, #Replace the existing default client if it exists [switch]$Force ) if ($SCRIPT:aiClient -and (-not $NoDefault -and -not $Force)) { Write-Warning 'Already connected to an AI engine. You can use -NoDefault to not set this client as the default client, or -Force to replace the existing default client.' return } if (-not $APIKey -and $env:OPENAI_API_KEY) { Write-Verbose 'Using API key from environment variable OPENAI_API_KEY' $APIKey = $env:OPENAI_API_KEY } $client = New-AIClient @newAIClientParams -APIKey $APIKey -GithubCopilot:$GitHubCopilot if ($NoDefault) { $PassThru = $true } else { $SCRIPT:aiClient = $client } if ($PassThru) { return $client } } filter Get-AIModel { [OutputType([OpenAI.Model])] [CmdletBinding()] param( # The ID of the model to get. If not specified, returns all models. [Parameter(ValueFromPipeline)][string]$Id, [ValidateNotNullOrEmpty()][OpenAI.Client]$Client = $SCRIPT:aiClient ) if (-not $Client) { Assert-Connected $Client = $SCRIPT:aiClient } if ($Id) { return $Client.RetrieveModel($Id) } $Client.ListModels().Data } function Get-AIEngine { [OutputType([OpenAI.Engine])] [CmdletBinding()] param( [ValidateNotNullOrEmpty()][OpenAI.Client]$Client = $SCRIPT:aiClient ) Write-Warning 'Engines are deprecated. Use Get-AIModel instead.' if (-not $Client) { Assert-Connected $Client = $SCRIPT:aiClient } $Client.ListEngines() | ConvertFrom-ListResponse } function Get-AICompletion { [CmdletBinding()] [OutputType([OpenAI.CreateCompletionResponse])] param( [Parameter(Mandatory)]$Prompt, #The name of the model to use. [ValidateSet([AvailableModels])][String]$Model = $SCRIPT:aiDefaultModel, [ValidateNotNullOrEmpty()][OpenAI.Client]$Client = $SCRIPT:aiClient, [ValidateNotNullOrEmpty()][uint]$MaxTokens = 1000, [ValidateNotNullOrEmpty()][uint]$Temperature = 0 ) if (-not $Client) { Assert-Connected $Client = $SCRIPT:aiClient } $request = [CreateCompletionRequest]@{ Prompt = $Prompt Stream = $false Model = $Model Max_tokens = $MaxTokens Temperature = $Temperature } $Client.CreateCompletion($request) } function Get-AICode { <# .SYNOPSIS Utilizes the Codex models to fetch a code completion given a prompt. .LINK https://platform.openai.com/docs/guides/code/introduction #> [OutputType([OpenAI.CreateCompletionResponse])] [CmdletBinding()] param( [string[]]$Prompt, #The name of the model to use. $Language = 'PowerShell 7', [ValidateSet([AvailableModels])][String]$Model = $SCRIPT:aiDefaultCodeModel, [ValidateNotNullOrEmpty()][OpenAI.Client]$Client = $SCRIPT:aiClient, [ValidateNotNullOrEmpty()][uint]$MaxTokens = 1000, [ValidateNotNullOrEmpty()][uint]$Temperature = 0 ) if (-not $Client) { Assert-Connected $Client = $SCRIPT:aiClient } #Add a language specifier to the prompt $Prompt.Insert(0, "#$Language") Get-AICompletion -Prompt $Prompt -Model $Model -MaxTokens $MaxTokens -Temperature $Temperature } function Get-AIChat { [OutputType([OpenAI.ChatConversation])] [CmdletBinding(DefaultParameterSetName = 'Prompt')] param( #Include one or more prompts to start the conversation [Parameter(Mandatory, Position = 0, ValueFromPipeline, ParameterSetName = 'Prompt')] [Parameter(ParameterSetName = 'ChatSession')] [OpenAI.ChatCompletionRequestMessage[]]$Prompt, #Supply a previous chat session to add new responses to it [Parameter(Mandatory, ValueFromPipeline, ParameterSetName = 'ChatSession')] [Parameter(ParameterSetName = 'Prompt')] [OpenAI.ChatConversation]$ChatSession, #The name of the model to use. [ValidateSet([AvailableModels])] [String]$Model = $SCRIPT:aiDefaultChatModel, [ValidateNotNullOrEmpty()] [OpenAI.Client]$Client = $SCRIPT:aiClient, [ValidateNotNullOrEmpty()] [uint]$MaxTokens = 1000, [ValidateNotNullOrEmpty()] [uint]$Temperature = 0, #Stream the response. You will lose syntax highlighting and usage info. [switch]$Stream ) if (-not $Client) { Assert-Connected $Client = $SCRIPT:aiClient } $ChatSession ??= [ChatConversation]@{ Request = @{ Messages = [List[ChatCompletionRequestMessage]]@() Stream = $false Model = $Model Max_tokens = $MaxTokens Temperature = $Temperature } } #Append any response to the initial request. This is the continuation of a chat. $responseChoices = $ChatSession.Response.Choices $requestMessages = $ChatSession.Request.Messages if ($responseChoices.Count -gt 0) { if ($responseChoices.count -gt 1) { Write-Error 'The previous chat response contained more than one choice. Continuing a conversation with multiple choices is not supported.' -Category 'NotImplemented' return } $requestMessages.Add($responseChoices[0].Message) } foreach ($PromptItem in $Prompt) { $requestMessages.Add( $PromptItem ) } if ($Stream) { $Client.CreateChatCompletionAsStream($ChatSession.Request) | ForEach-Object { $PSItem } Write-Host return } $chatResponse = $Client.CreateChatCompletion($ChatSession.Request) $chatSession.Response = $chatResponse $price = Get-UsagePrice -Model $chatResponse.Model -Total $chatResponse.Usage.Total_tokens Write-Verbose "Chat usage - $($chatResponse.Usage) $($price ? "$price " : $null)for Id $($chatResponse.Id)" return $chatSession #Stream the response } #endregion Public #Region Private function New-AIClient { [OutputType([OpenAI.Client])] param( [string]$ApiKey, [Switch]$GithubCopilot ) if (-not $APIKey) { Write-Error 'You must supply an OpenAI API key via the -APIKey parameter or by setting the OPENAI_API_KEY variable' return } if ($SCRIPT:client -and -not $Force) { Write-Warning 'Assistant is already connected. Please use -Force to reset the client.' return } $httpClient = [HttpClient]::new() $httpClient.DefaultRequestHeaders.Authorization = [AuthenticationHeaderValue]::new('Bearer', $APIKey) $aiClient = [Client]::new($httpClient) if ($GitHubCopilot) { $aiClient.BaseUrl = 'https://copilot-proxy.githubusercontent.com' } return $aiClient } function Assert-Connected { if (-not $SCRIPT:aiClient) { Connect-AI } } #If the returned result was a list, return the actual data filter ConvertFrom-ListResponse { if ($PSItem.Object -ne 'list') { return } return $PSItem.Data } #endregion Private # function Connect-Copilot { # [CmdletBinding()] # param( # # Provide your Copilot API Key as the password, and optionally your organization ID as the username # [string]$Token, # #Reset if a client already exists # [Switch]$Force # ) # $ErrorActionPreference = 'Stop' # if ($SCRIPT:GHClient -and -not $Force) { # Write-Warning 'Copilot is already connected. Please use -Force to reset the client.' # return # } # if ($SCRIPT:GHCopilotToken -and -not $Force) { # Write-Warning 'GitHub Copilot is already connected. Please use -Force to reset the client.' # return # } # $SCRIPT:GHCopilotToken = if (-not $Token) { # #Try to autodiscover it from GitHub Copilot CLI # if (-not (Test-Path $HOME/.copilot-cli-access-token)) { # Write-Error "To use PowerShell Assistant with GitHub Copilot, you must install GitHub Copilot CLI and run 'github-copilot-cli auth' at least once to generate a Copilot Personal Access Token (PAT)" # return # } # Get-Content $HOME/.copilot-cli-access-token # } else { # $Token # } # $config = [OpenAIOptions]@{ # ApiKey = Update-GitHubCopilotToken $SCRIPT:GHCopilotToken # BaseDomain = 'https://copilot-proxy.githubusercontent.com' # DefaultEngineId = 'copilot-labs-codex' # } # $SCRIPT:GHClient = [OpenAIService]::new($config) # } # function Get-CopilotSuggestion { # [CmdletBinding()] # param( # [Parameter(Mandatory)][string]$prompt, # [ValidateNotNullOrEmpty()]$client = $SCRIPT:GHClient # ) # if (-not $SCRIPT:GHClient) { Connect-Copilot } # $request = [CompletionCreateRequest]@{ # N = 1 # StopAsList = [string[]]@('---', '\n') # MaxTokens = 256 # Temperature = 0 # TopP = 1 # Prompt = $prompt # Stream = $true # } # $resultStream = $client.Completions.CreateCompletionAsStream($request).GetAwaiter.GetResult() # foreach ($resultItem in $resultStream) { # Write-Host -NoNewline 'NEW TOKEN' # #This gives us intellisense in vscode # [CompletionCreateResponse]$result = $resultItem # if ($result.Error) { # Write-Error $result.Error # return # } # $token = $result.Choices[0].Text # Write-Host -NoNewline -fore DarkGray $token # } # Write-Host 'DONE' # } # function Assert-Connected { # if (-not $SCRIPT:client) { Connect-Assistant } # } function Update-GitHubCopilotToken { <# .SYNOPSIS Fetches the latest token for GitHub Copilot #> param( [ValidateNotNullOrEmpty()] $GitHubToken = $SCRIPT:GHCopilotToken ) $ErrorActionPreference = 'Stop' $response = Invoke-RestMethod 'https://api.github.com/copilot_internal/v2/token' -Headers @{ Authorization = "token $($GitHubToken.trim())" } return $response.token } function Get-Chat { <# .SYNOPSIS Provides an interactive assistant for PowerShell. Mostly a frontend to Get-AIChat #> [CmdletBinding()] param( #Provide a chat prompt to initiate the conversation [string[]]$chatPrompt, [ValidateNotNullOrEmpty()] #Specify a prompt that guides Chat how to behave. By default, it is told to prefer PowerShell as a language. [string]$SystemPrompt = 'Prefer PowerShell syntax and always use fenced code blocks with the language specified.', #Maximum tokens to generate. Defaults to 500 to minimize accidental API billing [ValidateNotNullOrEmpty()] [uint]$MaxTokens = 500, [ValidateSet([AvailableModels])] [string]$Model, #Specify this to add the latest recommended codeblock to your psreadline history [switch]$AddToHistory, #By default, the latest code recommended is added to your clipboard for easy pasting. [switch]$NoClipboard, #If you just want the result and don't want to be prompted for further replies, specify this [Switch]$NoReply, #Code blocks in these languages will be copied to the clipboard or added to history. By default, this is just PowerShell. This only affects code copying, you will need to specify a new -SystemPrompt if you want non-powershell recommendations. [string[]]$languages = 'powershell', #By default, chat operates inline. Specify this to use the alternate screen buffer which will operate chat in a separate "screen" and then return you to your original prompt once the chat is complete. [switch]$UseAltScreenBuffer ) begin { $ErrorActionPreference = 'Stop' Assert-Connected [List[ChatCompletionRequestMessage]]$chatHistory = @( [ChatCompletionRequestMessage]@{ Role = [ChatCompletionRequestMessageRole]::System Content = $SystemPrompt } ) } process { if ($UseAltScreenBuffer) { $inChatSession = $true Write-Host -NoNewline "`e[?1049h" } do { $chatPrompt ??= Read-Host -Prompt 'You' foreach ($promptItem in $chatPrompt) { $chatHistory.Add( ([ChatCompletionRequestMessage]$promptItem) ) } $chatParams = @{ Prompt = $chatHistory MaxTokens = $MaxTokens Stream = $true } if ($Model) { $chatParams.Model = $Model } [List[CreateChatCompletionChunkedResponse]]$streamedResponse = @() [Text.StringBuilder]$chatStream = '' Get-AIChat @chatParams | ForEach-Object { [CreateChatCompletionChunkedResponse]$response = $PSItem $streamedResponse.Add($response) [DeltaChoice]$firstChoice = $response.Choices[0] [string]$firstChoiceContent = $firstChoice.Delta.Content [void]$chatStream.Append($firstChoiceContent) $markdownCodeFenceRegex = '```(?<lang>\w+)?\s*(?<code>[\s\S]*?)```' #Start recording if a code block occurs, and if it does, reformat it and copy it to clipboard #TODO: This could maybe be faster by watching the stream for the starting and trailing backticks if ($chatStream -match $markdownCodeFenceRegex) { $codeblock = $matches[0] $code = $matches.code $lang = $matches.lang $codeBlockLineCount = ($codeBlock -split '\r?\n').Count - 1 #Use ANSI Codes Move the cursor up to the start of the code block to overwrite it Write-Host -NoNewline "`e[${codeBlockLineCount}F" Write-Host -NoNewline "`e[0J" #Two newlines replace where the previous code was $formattedCodeBlock = [Environment]::NewLine + (Show-Markdown -InputObject $codeBlock) + [Environment]::NewLine if ($lang -in $languages) { if ($AddToHistory) { [PSConsoleReadline]::AddToHistory($code.Trim()) } if (-not $NoClipboard) { Set-Clipboard -Value ". {$($code.Trim())}" } } Write-Host -NoNewline ($formattedCodeBlock) #Update the stringbuilder #TODO: Add the start index which generally should not be necessary but probably smart [void]$chatStream.Replace($codeBlock, $formattedCodeBlock) } else { Write-Host -NoNewline ($PSStyle.Foreground.BrightBlack + $firstChoice.Delta.Content) } #Matches single code blocks $markdownCodeRegex = '(?<!`)(`)(?!`)[^\n]+?\1' if ($chatStream -match $markdownCodeRegex) { $code = $matches[0] $formattedCode = (Show-Markdown -InputObject $code).Trim() #Move cursor back to the start of the code block and clear text after it Write-Host -NoNewline "`e[$($code.Length)D" Write-Host -NoNewline "`e[0J" Write-Host -NoNewline $formattedCode [void]$chatStream.Replace($code, $formattedCode) } if ($firstChoice.Finish_reason -eq 'length') { Write-Host -ForegroundColor $PSStyle.Formatting.Warning '[END]' Write-Warning "Response truncated due to length. Consider setting -MaxTokens greater than $MaxTokens" } } $message = [ChatCompletionRequestMessage]::new( [string]::Concat($streamedResponse.Choices.Delta.Content), [ChatCompletionRequestMessageRole]::Assistant ) $chatHistory.Add($message) #TODO: Move this into the formatter # switch ($aiResponse.FinishReason) { # 'stop' {} #This is the normal response # 'length' { # Write-Warning "$MaxTokens tokens reached. Consider increasing the value of -MaxTokens for longer responses." # } # $null { # Write-Debug 'Null FinishReason received. This seems to occur on occasion and may or may not be a bug.' # } # default { # Write-Warning "Chat response finished abruply due to: $($aiResponse.FinishReason)" # } # } $chatPrompt = $null if (-not $NoReply) { Write-Host ($PSStyle.Foreground.Cyan + '<Ctrl-C to exit>' + $PSStyle.Reset) } } while ( -not $NoReply ) } clean { if ($inChatSession) { #Closes the alternate screen buffer Write-Host "`e[?1049l" } } } filter Convert-ChatCodeToClipboard { <# .SYNOPSIS Given a string, take the last occurance of text surrounded by a fenced code block, and copy it to the clipboard. It will also pass through the string for further filtering #> $fencedCodeBlockRegex = '(?s)```[\r|\n|powershell]+(.+?)```' $matchResult = $PSItem -match $fencedCodeBlockRegex $savedMatches = $matches $cbMatch = $savedMatches.($savedMatches.Keys | Sort-Object | Select-Object -Last 1) if (-not $matchResult) { Write-Debug 'No code block detected, skipping this step' return $PSItem } Write-Verbose "Copying last suggested code block to clipboard:`n$cbMatch" Set-Clipboard -Value $cbMatch return $PSItem } class AvailableModels : IValidateSetValuesGenerator { [String[]] GetValidValues() { trap { Write-Host ''; Write-Host -NoNewline -ForegroundColor Red "Validation Error: $PSItem" } $models = Get-AIModel return $models.Id } } filter Format-ChatCode { <# .SYNOPSIS Given a string, for any occurance of text surrounded by backticks, replace the backticks with ANSI escape codes #> $codeBlockRegex = '(?s)```[\r|\n|powershell]+(.+?)```' $codeSnippetRegex = '(?s)`(.+?)`' $boldSelectedText = ($PSStyle.Italic + '$1' + $PSStyle.ItalicOff) $PSItem -replace $codeBlockRegex, $boldSelectedText -replace $codeSnippetRegex, $boldSelectedText } filter Format-ChatMessage { param( [Parameter(ValueFromPipeline)]$message, #Notes that the content should be streamed rather than returned line by line [switch]$Stream ) $role = $message.Role $content = $message.Content $roleColor = switch ($role) { 'System' { 'DarkYellow' } 'Assistant' { 'Green' } 'User' { 'DarkCyan' } default { 'DarkGray' } } if ($Stream) { if ($role) { return "$($PSStyle.Foreground.$roleColor)$role`:$($PSStyle.Reset) " } elseif ($content) { return "$($PSStyle.Foreground.BrightBlack)$content$($PSStyle.Reset)" } else { #Blank entry, we might want to throw here just in case tho it is technically allowed. return } } $formattedMessage = $content.Trim() | Format-ChatCode return "$($PSStyle.Foreground.$roleColor)$role`:$($PSStyle.Reset) $($PSStyle.ForeGround.BrightBlack)$formattedMessage" } function Format-CreateChatCompletionChunkedResponse { param( [Parameter(ValueFromPipeline)][CreateChatCompletionChunkedResponse]$response ) Format-ChatMessage -Stream $response.Choices[0].Delta } function Format-Choices2 { [AssemblyMetadata('Format-Custom', 'Choices2')] param( [Choices2]$choice ) $PSStyle.Foreground.BrightCyan + "Choice $([int]$choice.Index + 1): " + (Format-ChatMessage $choice.Message) } filter Format-CreateChatCompletionRequest { [AssemblyMetadata('Format-Custom', 'OpenAI.CreateChatCompletionRequest')] param( [Parameter(ValueFromPipeline)][CreateChatCompletionRequest]$request ) $request.messages | Format-ChatMessage } filter Format-CreateChatCompletionResponse { [AssemblyMetadata('Format-Custom', 'OpenAI.CreateChatCompletionResponse')] param( [Parameter(ValueFromPipeline)][CreateChatCompletionResponse]$response ) if ($response.Choices.Count -eq 1) { Format-ChatMessage $response.Choices[0].Message } else { $Response.Choices } } function Format-ChatConversation { param( [ChatConversation]$conversation ) $messages = @() $messages += $conversation.Request | Format-CreateChatCompletionRequest $messages += $conversation.Response | Format-CreateChatCompletionResponse return $messages -join ($PSStyle.Reset + [Environment]::NewLine) } function Get-UsagePrice { param( [string]$Model, [int]$Total ) #Taken from: https://openai.com/pricing $pricePerToken = @{ 'code' = 0 'gpt-3.5-turbo' = .002 / 1000 'ada' = .0004 / 1000 'babbage' = .0005 / 1000 'curie' = .002 / 1000 'davinci' = .002 / 1000 } foreach ($priceItem in $pricePerToken.GetEnumerator()) { if ($Model.Contains($priceItem.key)) { #Will return the first match $totalPrice = $total * $priceItem.Value #Formats as currency ($3.2629) and strips trailing zeroes return $totalPrice.ToString('C15').TrimEnd('0') } } #Return an empty string if no pricing engine found. return [string]::Empty } |