-
-
Notifications
You must be signed in to change notification settings - Fork 54
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add ChatBot GPT-3 integration (#16)
* Added dependency go-resty * Configured Parser to read plugin configurations * Add example ssh with ChatGPT plugin * Add client ChatBot * Improve logging * Add integration with plugin OpenAIChatGPT * Improve readme with ChatBot Example * Add contributed ChatGPT question * Refactoring * Refactoring and improve unit test
- Loading branch information
1 parent
9ddb076
commit d062435
Showing
9 changed files
with
260 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
apiVersion: "v1" | ||
protocol: "ssh" | ||
address: ":2222" | ||
description: "SSH interactive ChatGPT" | ||
commands: | ||
- regex: "^(.+)$" | ||
plugin: "OpenAIChatGPT" | ||
serverVersion: "OpenSSH" | ||
serverName: "ubuntu" | ||
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" | ||
deadlineTimeoutSeconds: 60 | ||
plugin: | ||
openAPIChatGPTSecretKey: "" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,113 @@ | ||
package plugin | ||
|
||
import ( | ||
"encoding/json" | ||
"errors" | ||
"fmt" | ||
"github.com/go-resty/resty/v2" | ||
"strings" | ||
) | ||
|
||
const ChatGPTPluginName = "OpenAIChatGPT" | ||
const openAIGPTEndpoint = "https://api.openai.com/v1/completions" | ||
|
||
type History struct { | ||
Input, Output string | ||
} | ||
|
||
type OpenAIGPTVirtualTerminal struct { | ||
Histories []History | ||
OpenAPIChatGPTSecretKey string | ||
client *resty.Client | ||
} | ||
|
||
func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) InjectDependency() { | ||
if openAIGPTVirtualTerminal.client == nil { | ||
openAIGPTVirtualTerminal.client = resty.New() | ||
} | ||
} | ||
|
||
type Choice struct { | ||
Text string `json:"text"` | ||
Index int `json:"index"` | ||
Logprobs interface{} `json:"logprobs"` | ||
FinishReason string `json:"finish_reason"` | ||
} | ||
|
||
type gptResponse struct { | ||
ID string `json:"id"` | ||
Object string `json:"object"` | ||
Created int `json:"created"` | ||
Model string `json:"model"` | ||
Choices []Choice `json:"choices"` | ||
Usage struct { | ||
PromptTokens int `json:"prompt_tokens"` | ||
CompletionTokens int `json:"completion_tokens"` | ||
TotalTokens int `json:"total_tokens"` | ||
} `json:"usage"` | ||
} | ||
|
||
type gptRequest struct { | ||
Model string `json:"model"` | ||
Prompt string `json:"prompt"` | ||
Temperature int `json:"temperature"` | ||
MaxTokens int `json:"max_tokens"` | ||
TopP int `json:"top_p"` | ||
FrequencyPenalty int `json:"frequency_penalty"` | ||
PresencePenalty int `json:"presence_penalty"` | ||
Stop []string `json:"stop"` | ||
} | ||
|
||
//Reference: https://www.engraved.blog/building-a-virtual-machine-inside/ | ||
const chatGPTFirstQuestion = "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\n" | ||
|
||
func buildPrompt(histories []History, command string) string { | ||
var sb strings.Builder | ||
|
||
sb.WriteString(chatGPTFirstQuestion) | ||
|
||
for _, history := range histories { | ||
sb.WriteString(fmt.Sprintf("A:%s\n\nQ:%s\n\n", history.Input, history.Output)) | ||
} | ||
// Append command to evaluate | ||
sb.WriteString(fmt.Sprintf("A:%s\n\nQ:", command)) | ||
|
||
return sb.String() | ||
} | ||
|
||
func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command string) (string, error) { | ||
requestJson, err := json.Marshal(gptRequest{ | ||
Model: "text-davinci-003", | ||
Prompt: buildPrompt(openAIGPTVirtualTerminal.Histories, command), | ||
Temperature: 0, | ||
MaxTokens: 100, | ||
TopP: 1, | ||
FrequencyPenalty: 0, | ||
PresencePenalty: 0, | ||
Stop: []string{"\n"}, | ||
}) | ||
if err != nil { | ||
return "", err | ||
} | ||
|
||
if openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey == "" { | ||
return "", errors.New("OpenAPIChatGPTSecretKey is empty") | ||
} | ||
|
||
response, err := openAIGPTVirtualTerminal.client.R(). | ||
SetHeader("Content-Type", "application/json"). | ||
SetBody(requestJson). | ||
SetAuthToken(openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey). | ||
SetResult(&gptResponse{}). | ||
Post(openAIGPTEndpoint) | ||
|
||
if err != nil { | ||
return "", err | ||
} | ||
|
||
if len(response.Result().(*gptResponse).Choices) == 0 { | ||
return "", errors.New("no choices") | ||
} | ||
|
||
return response.Result().(*gptResponse).Choices[0].Text, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
package plugin | ||
|
||
import ( | ||
"github.com/go-resty/resty/v2" | ||
"github.com/jarcoal/httpmock" | ||
"github.com/stretchr/testify/assert" | ||
"net/http" | ||
"testing" | ||
) | ||
|
||
func TestBuildPromptEmptyHistory(t *testing.T) { | ||
//Given | ||
var histories []History | ||
command := "pwd" | ||
|
||
//When | ||
prompt := buildPrompt(histories, command) | ||
|
||
//Then | ||
assert.Equal(t, | ||
"I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\nA:pwd\n\nQ:", | ||
prompt) | ||
} | ||
|
||
func TestBuildPromptWithHistory(t *testing.T) { | ||
//Given | ||
var histories = []History{ | ||
{ | ||
Input: "cat hello.txt", | ||
Output: "world", | ||
}, | ||
{ | ||
Input: "echo 1234", | ||
Output: "1234", | ||
}, | ||
} | ||
|
||
command := "pwd" | ||
|
||
//When | ||
prompt := buildPrompt(histories, command) | ||
|
||
//Then | ||
assert.Equal(t, | ||
"I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\nA:cat hello.txt\n\nQ:world\n\nA:echo 1234\n\nQ:1234\n\nA:pwd\n\nQ:", | ||
prompt) | ||
} | ||
|
||
func TestBuildGetCompletions(t *testing.T) { | ||
client := resty.New() | ||
httpmock.ActivateNonDefault(client.GetClient()) | ||
defer httpmock.DeactivateAndReset() | ||
|
||
// Given | ||
httpmock.RegisterResponder("POST", openAIGPTEndpoint, | ||
func(req *http.Request) (*http.Response, error) { | ||
resp, err := httpmock.NewJsonResponse(200, &gptResponse{ | ||
Choices: []Choice{ | ||
{ | ||
Text: "prova.txt", | ||
}, | ||
}, | ||
}) | ||
if err != nil { | ||
return httpmock.NewStringResponse(500, ""), nil | ||
} | ||
return resp, nil | ||
}, | ||
) | ||
|
||
openAIGPTVirtualTerminal := OpenAIGPTVirtualTerminal{ | ||
OpenAPIChatGPTSecretKey: "sdjdnklfjndslkjanfk", | ||
client: client, | ||
} | ||
|
||
//When | ||
str, err := openAIGPTVirtualTerminal.GetCompletions("ls") | ||
|
||
//Then | ||
assert.Nil(t, err) | ||
assert.Equal(t, "prova.txt", str) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.