Merge pull request #146 from dfinke/Enable-Disable-Chat-Persistence

Make it so you can enable and disable persisting chat messages during a session

Author: dfinke

PowerShellAI.psd1

@@ -1,6 +1,6 @@
 @{
     RootModule        = 'PowerShellAI.psm1'
-    ModuleVersion     = '0.7.5'
+    ModuleVersion     = '0.7.6'
     GUID              = '081ce7b4-6e63-41ca-92a7-2bf72dbad018'
     Author            = 'Douglas Finke'
     CompanyName       = 'Doug Finke'
@@ -19,9 +19,13 @@ The PowerShell AI module integrates with the OpenAI API and let's you easily acc
         'ConvertFrom-GPTMarkdownTable'
         'copilot'
         'Disable-AIShortCutKey'
+        'Disable-ChatPersistence'
         'Enable-AIShortCutKey'
+        'Enable-ChatPersistence'
+        'Get-ChatCompletion'        
+        'Get-ChatPersistence'
+        'Get-CompletionFromMessages'
         'Get-DalleImage'
-        'Get-ChatCompletion'
         'Get-GPT3Completion'
         'Get-GPT4Completion'
         'Get-GPT4Response'

PowerShellAI.psm1

@@ -1,6 +1,37 @@
-$Script:OpenAIKey = $null
+# Set the OpenAI key to null
+$Script:OpenAIKey = $null
 
+# Set the chat API provider to OpenAI
+$Script:ChatAPIProvider = 'OpenAI'
 
+# Set the chat in progress flag to false
+$Script:ChatInProgress = $false
+
+# Create an array list to store chat messages
+[System.Collections.ArrayList]$Script:ChatMessages = @()
+
+# Enable chat persistence
+$Script:ChatPersistence = $true
+
+# Set the options for the chat session
+$Script:ChatSessionOptions = @{
+    'model'             = 'gpt-4'
+    'temperature'       = 0.0
+    'max_tokens'        = 256
+    'top_p'             = 1.0
+    'frequency_penalty' = 0
+    'presence_penalty'  = 0
+    'stop'              = $null
+}
+
+# Set the options for the Azure OpenAI API
+$Script:AzureOpenAIOptions = @{
+    Endpoint       = 'not set'
+    DeploymentName = 'not set'
+    ApiVersion     = 'not set'
+}
+
+# Load all PowerShell scripts in the Public and Private directories
 foreach ($directory in @('Public', 'Private')) {
     Get-ChildItem -Path "$PSScriptRoot\$directory\*.ps1" | ForEach-Object { . $_.FullName }
 }

Private/Invoke-RestMethodWithProgress.ps1

@@ -37,7 +37,8 @@ function Set-APIResponseTime {
 
 function Invoke-RestMethodWithProgress {
     param (
-        [hashtable] $Params
+        [hashtable] $Params,
+        $ProgressActivity = "Thinking..."
     )
 
     # Some hosts can't support background jobs. It's best to opt-in to this feature by using a list of supported hosts
@@ -70,10 +71,10 @@ function Invoke-RestMethodWithProgress {
             if($logPercent -eq 100) {
                 $status = "API is taking longer than expected"
             }
-            Write-Progress -Id 1 -Activity "Invoking AI" -Status $status -PercentComplete $logPercent
+            Write-Progress -Id 1 -Activity $ProgressActivity -Status $status -PercentComplete $logPercent
             Start-Sleep -Milliseconds 50
         }
-        Write-Progress -Id 1 -Activity "Invoking AI" -Completed
+        Write-Progress -Id 1 -Activity $ProgressActivity -Completed
 
         # If Invoke-RestMethod failed in the job rethrow this up to the caller so it's like a normal web error
         if($job.State -eq "Failed") {

Public/Disable-ChatPersistence.ps1

@@ -0,0 +1,13 @@
+function Disable-ChatPersistence {
+    <#
+        .SYNOPSIS
+        Disables chat persistence.
+
+        .DESCRIPTION
+        This function disables chat persistence by setting the $ChatPersistence variable to $false.
+
+        .EXAMPLE
+        Disable-ChatPersistence
+    #>
+    $Script:ChatPersistence = $false
+}

Public/Enable-ChatPersistence.ps1

@@ -0,0 +1,14 @@
+function Enable-ChatPersistence {
+    <#
+        .SYNOPSIS
+        Enables chat persistence.
+
+        .DESCRIPTION
+        The Enable-ChatPersistence function sets the $Script:ChatPersistence variable to $true, which enables chat persistence.
+
+        .EXAMPLE
+        Enable-ChatPersistence
+    #>
+
+    $Script:ChatPersistence = $true
+}

Public/Get-ChatPersistence.ps1

@@ -0,0 +1,7 @@
+function Get-ChatPersistence {
+    <#
+        .SYNOPSIS
+        Retrieves the chat persistence flag.
+    #>
+    $Script:ChatPersistence
+}

Public/Get-CompletionFromMessages.ps1

@@ -0,0 +1,40 @@
+function Get-CompletionFromMessages {
+    <#
+        .SYNOPSIS
+        Gets completion suggestions based on the array of messages.
+
+        .DESCRIPTION
+        The Get-CompletionFromMessages function returns completion suggestion based on the messages.
+
+        .PARAMETER Messages
+        Specifies the chat messages to use for generating completion suggestions.
+
+        .EXAMPLE
+        Get-CompletionFromMessages $(
+            New-ChatMessageTemplate -Role system 'You are a PowerShell expert'
+            New-ChatMessageTemplate -Role user 'List even numbers between 1 and 10'
+        )
+    #>
+    [CmdletBinding()]
+    param(
+        [Parameter(Mandatory)]
+        $Messages
+    )
+
+    $payload = (Get-ChatSessionOptions).Clone()
+
+    $payload.messages = $messages
+    $payload = $payload | ConvertTo-Json -Depth 10
+    
+    $body = [System.Text.Encoding]::UTF8.GetBytes($payload)
+
+    if ((Get-ChatAPIProvider) -eq 'OpenAI') {
+        $uri = Get-OpenAIChatCompletionUri
+    }
+    elseif ((Get-ChatAPIProvider) -eq 'AzureOpenAI') {
+        $uri = Get-ChatAzureOpenAIURI
+    }
+
+    $result = Invoke-OpenAIAPI -Uri $uri -Method 'Post' -Body $body 
+    $result.choices.message
+}

Public/Get-GPT4Completion.ps1

@@ -1,24 +1,3 @@
-$Script:ChatSessionOptions = @{
-    'model'             = 'gpt-4'
-    'temperature'       = 0.0
-    'max_tokens'        = 256
-    'top_p'             = 1.0
-    'frequency_penalty' = 0
-    'presence_penalty'  = 0
-    'stop'              = $null
-}
-
-$Script:AzureOpenAIOptions = @{
-    Endpoint       = 'not set'
-    DeploymentName = 'not set'
-    ApiVersion     = 'not set'
-}
-
-$Script:ChatAPIProvider = 'OpenAI'
-$Script:ChatInProgress = $false
-
-[System.Collections.ArrayList]$Script:ChatMessages = @()
-
 function Get-AzureOpenAIOptions {
     [CmdletBinding()]
     param()
@@ -197,6 +176,8 @@ function Reset-ChatSessionOptions {
         'presence_penalty'  = 0
         'stop'              = $null
     }
+
+    Enable-ChatPersistence
 }
 
 function Clear-ChatMessages {

Public/SessionManagement.ps1

@@ -181,6 +181,8 @@ function Export-ChatSession {
     [CmdletBinding()]
     param ()
 
+    if ((Get-ChatPersistence) -eq $false) { return }
+
     $sessionPath = Get-ChatSessionPath
     if (-not (Test-Path $sessionPath)) {
         New-Item -ItemType Directory -Path $sessionPath -Force | Out-Null

__tests__/Disable-ChatPersistence.tests.ps1

@@ -0,0 +1,20 @@
+Import-Module "$PSScriptRoot\..\PowerShellAI.psd1" -Force
+
+Describe 'Disable-ChatPersistence' -Tag ChatPersistence {
+
+    BeforeEach {
+        Reset-ChatSessionOptions
+    }
+
+    It 'tests the function Disable-ChatPersistence exists' {
+        $actual = Get-Command Disable-ChatPersistence -ErrorAction SilentlyContinue
+        $actual | Should -Not -BeNullOrEmpty
+    }
+
+    It 'tests that it disables chat persistence' {
+        Disable-ChatPersistence
+        $actual = Get-ChatPersistence
+
+        $actual | Should -Be $false
+    }
+}

__tests__/Enable-ChatPersistence.tests.ps1

@@ -0,0 +1,20 @@
+Import-Module "$PSScriptRoot\..\PowerShellAI.psd1" -Force
+
+Describe 'Enable-ChatPersistence' -Tag ChatPersistence {
+
+    AfterEach {
+        Reset-ChatSessionOptions
+    }
+
+    It 'tests the function Enable-ChatPersistence exists' {
+        $actual = Get-Command Enable-ChatPersistence -ErrorAction SilentlyContinue
+        $actual | Should -Not -BeNullOrEmpty
+    }
+
+    It 'tests that it enables chat persistence' {
+        Enable-ChatPersistence
+        $actual = Get-ChatPersistence
+
+        $actual | Should -Be $true
+    }
+}

__tests__/Get-ChatPersistence.tests.ps1

@@ -0,0 +1,13 @@
+Import-Module "$PSScriptRoot\..\PowerShellAI.psd1" -Force
+
+Describe 'Get-ChatPersistence' -Tag ChatPersistence {
+    It 'tests the function Get-ChatPersistence exists' {
+        $actual = Get-Command Get-ChatPersistence -ErrorAction SilentlyContinue
+        $actual | Should -Not -BeNullOrEmpty
+    }
+
+    It 'tests the function Get-ChatPersistence returns a boolean' {
+        $actual = Get-ChatPersistence
+        $actual.GetType().Name | Should -Be 'Boolean'
+    }
+}

__tests__/Get-CompletionFromMessages.tests.ps1

@@ -0,0 +1,47 @@
+Import-Module $PSScriptRoot\..\PowerShellAI.psd1 -Force
+
+Describe "Get-CompletionFromMessages" -Tag "Get-CompletionFromMessages" {
+    BeforeAll {
+        $script:savedKey = $env:OpenAIKey
+        $env:OpenAIKey = 'sk-1234567890'
+    }
+    
+    AfterAll {
+        $env:OpenAIKey = $savedKey
+    }
+
+    It "tests the function Get-CompletionFromMessages exists" {
+        $actual = Get-Command Get-CompletionFromMessages -ErrorAction SilentlyContinue
+        $actual | Should -Not -BeNullOrEmpty
+    }
+
+    It "tests Get-CompletionFromMessages has a parameter named Messages" {
+        $actual = Get-Command Get-CompletionFromMessages -ErrorAction SilentlyContinue
+        $actual.Parameters.Keys | Should -Contain Messages
+    }
+
+    It "tests Get-CompletionFromMessages returns a response" {
+        Mock Invoke-RestMethodWithProgress -ModuleName PowerShellAI -ParameterFilter { 
+            $Params.Method -eq 'Post' -and $Params.Uri -eq (Get-OpenAIChatCompletionUri) 
+        } -MockWith {
+            [PSCustomObject]@{
+                choices = @(
+                    [PSCustomObject]@{
+                        message = [PSCustomObject]@{
+                            content = 'Mocked Get-GPT4Completion call'
+                        }
+                    }
+                )
+            }
+        } 
+        
+        $messages = $(
+            New-ChatMessageTemplate -Role system "I am a bot"
+            New-ChatMessageTemplate -Role user "Hello"
+        )
+
+        $actual = Get-CompletionFromMessages -Messages $messages
+
+        $actual.content | Should -BeExactly "Mocked Get-GPT4Completion call"
+    }
+}

__tests__/SessionManagement.tests.ps1

@@ -12,14 +12,10 @@ Describe "Session Management" -Tag SessionManagement {
         }
     }
 
-    BeforeEach {
-        # Set-ChatSessionPath "TestDrive:\PowerShell\ChatGPT"
-    }
-
     AfterEach {
         Reset-ChatSessionPath
+        Reset-ChatSessionOptions
         Clear-ChatMessages
-        # Get-ChatSessionPath | Remove-Item -Recurse -force
     }
 
     AfterAll {
@@ -384,4 +380,43 @@ Describe "Session Management" -Tag SessionManagement {
         $result[8].role | Should -BeExactly 'assistant'
         $result[8].content | Should -BeExactly 'assistant test 3'
     }
+
+    It "tests Export-ChatSession respects ChatPersistence flag" {
+        Set-ChatSessionPath "TestDrive:\PowerShell\ChatGPT"
+        Get-ChatSessionPath | Get-ChildItem | Remove-Item -Recurse -Force -ErrorAction SilentlyContinue        
+
+        Add-ChatMessage -Message ([PSCustomObject]@{
+            role    = 'assistant'
+            content = 'assistant test 2'
+        })
+        
+        Export-ChatSession
+        
+        (Get-ChatSessionPath | Get-ChildItem ).Count | Should -Be 1
+        
+        Get-ChatSessionPath | Get-ChildItem | Remove-Item -Recurse -Force -ErrorAction SilentlyContinue
+
+        Disable-ChatPersistence
+
+        Add-ChatMessage -Message ([PSCustomObject]@{
+            role    = 'assistant'
+            content = 'assistant test 2'
+        })
+        
+        Export-ChatSession
+        
+        (Get-ChatSessionPath | Get-ChildItem ).Count | Should -Be 0        
+
+        Enable-ChatPersistence
+        Get-ChatSessionPath | Get-ChildItem | Remove-Item -Recurse -Force -ErrorAction SilentlyContinue        
+
+        Add-ChatMessage -Message ([PSCustomObject]@{
+            role    = 'assistant'
+            content = 'assistant test 2'
+        })
+        
+        Export-ChatSession
+        
+        (Get-ChatSessionPath | Get-ChildItem ).Count | Should -Be 1
+    }
 }

changelog.md

@@ -1,3 +1,8 @@
+# v0.7.6
+
+- Implement mechanism to turn on/off persistence of chat message
+- Add `Get-CompletionFromMessages` so you can get messages can be passes standalone to get the completions
+
 # v0.7.5
 
 - Fixed spaces in property names [#141](https://github.com/dfinke/PowerShellAI/issues/141)