Adds a key handler to PSReadLine that when pressed, asks the user what they're trying to do, then uses Meta Llama 3 405B (via NVIDIA) to get a short PowerShell snippet (other CLIs work as well) and injects it into the input buffer WITHOUT executing it, so that the user can review and make edits as needed.

PSReadLine-AskAI.ps1

# To use:
# - Add an environment variable called NV_LLAMA_API_KEY and set it to a valid API key
# - Add the following script somewhere to your profile
# - At a *blank* PowerShell prompt, press Ctrl+` (backtick)
# - Type what you want to do at the prompt in plain english
# - The code generated by the LLM will be inserted into the input buffer without executing it
# - Make whatever changes are necessary to the generated command(s) and proceed as normal
# 
# > Ask AI: recursively get all files in the current directory with the archive bit set
# > Get-ChildItem -Recurse -File | Where-Object { $_.Attributes -match 'Archive' }
#
# > Ask AI: use pandoc to convert myfile.docx to markdown
# > pandoc myfile.docx -f docx -t markdown -o myfile.md
#
# > Ask AI: pull the latest git commit without losing my pending changes!
# > git pull --autostash

    Set-PSReadLineKeyHandler -Key 'Ctrl+`' -ScriptBlock {

        [String]$API_URL   = 'https://integrate.api.nvidia.com/v1/chat/completions'
        [String]$API_MODEL = 'meta/llama-3.1-405b-instruct'

        # Probably should switch this to PowerShell's secret management module.
        # For now this needs to be set as a system or user environment variable.
        # Get an API key at https://build.nvidia.com
        [String]$API_KEY   = $ENV:NV_LLAMA_API_KEY

        [String]$InputText = ''
        [Int32]$Cursor = 0

        [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$InputText, [ref]$Cursor)

        # Abort if there's anything in the input buffer because I haven't really
        # thought about how to handle asking what to do mid-command yet.
        if ($InputText.Length) { return }

        if (!$API_KEY) {
            Write-Warning 'API_KEY not set.'
            return
        }

        if ($InputText = Read-Host -Prompt 'Ask AI') {

            $Headers = @{
                'Content-Type' = 'application/json'
                'Authorization' = "Bearer $($API_KEY)"
            }   

            $RequestBody = @{
                model = $API_MODEL
                max_tokens = 200
                temperature = 0.1
                messages = @(
                    @{
                        role='system';
                        content=
                        'You will provide assistance to a user with PowerShell or command line tasks. ' +
                        'Your response should consist of ONLY the short command line to carry out the desired task. ' +
                        'If you dont know, or the user asked a truly stupid or unrelated question, just return a snarky/dismissive reply in a PowerShell comment block.' +
                        'Your input will be injected directly into the input buffer, so do not include any markdown or commentary. ' +
                        'Your output *must* be able to be interpreted by PowerShell. ' +
                        'It is perfectly acceptable (and preferable even) to use well-known, default aliases like ? for Where-Object or % for ForEach-Object. ' +
                        'Prefer concise output where possible. For example, `? LinkType -eq SymbolicLink` over `Where-Object { $_.LinkType -eq "SymbolicLink" }`. ' +
                        'You may use CLI tools, but be sure to use PowerShell syntax if chaining commands.' +
                        'DO NOT USE BASH SYNTAX OR KITTENS WILL DIE. '
                    }
                    # This synthetic chat history acts as few-shot training examples.
                    # Seems to guide the model better by making it think it provided these
                    # responses instead of jamming them into the system prompt.
                    @{role='user';content='How do I create a symbolic link?'}
                    @{role='assistant';content='New-Item -ItemType SymbolicLink -Path $LinkPath -Target $TargetPath'}
                    @{role='user';content='pull the latest git commit without losing my pending changes'}
                    @{role='assistant';content='git pull --autostash'}
                    @{role='user';content='perform a long running task on each line in a file with concurrency'}
                    @{role='assistant';content="cat input.txt | % -Parallel -ThrottleLimit 4 {`n    # do long-running task`n}"}

                    # User input
                    @{role='user';content=$InputText}
                )
            }
        
            # Absolutely no error handling here
            $Response = Invoke-RestMethod -Method POST -Uri $API_URL -Headers $Headers -Body (ConvertTo-Json $RequestBody)
            $OutputText = $Response.choices[0].message.content

            # No matter how hard I try to guide it, occasionally it will still wrap a multi-line
            # response in a markdown code fence. This completely un-optimized and fragile regex
            # will attempt to strip off that markdown wrapper.
            $OutputText = $OutputText -replace '(?s)^(\s|\r|\n)*```powershell(\s|\r|\n)+(.*)(\s|\r|\n)+```(\s|\r|\n)*$','$3'

            # Goes back to PSReadLine prompt editing mode with the generated code sitting in the
            # input buffer waiting to be edited or executed. Afterwards, it will be in your command
            # history just as if you had typed it yourself.
            [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt()
            [Microsoft.PowerShell.PSConsoleReadLine]::Insert($OutputText)

        }

    }