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
This commit is contained in:
Mario Candela
2022-12-16 23:02:16 +01:00
committed by GitHub
parent 9ddb076621
commit d062435818
9 changed files with 260 additions and 11 deletions

113
plugin/OpenAiGPT.go Normal file
View File

@ -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
}

82
plugin/OpenAiGPT_test.go Normal file
View File

@ -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)
}