mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
refactor:Added Integration test and tiny refactoring (#23)
* Refactoring name convention * Added integration test * Added Makefile * Bump golang.org/x/crypto from 0.0.0-20220826181053-bd7e27e6170d to 0.6.0 Bumps [golang.org/x/crypto](https://github.com/golang/crypto) from 0.0.0-20220826181053-bd7e27e6170d to 0.6.0. - [Release notes](https://github.com/golang/crypto/releases) - [Commits](https://github.com/golang/crypto/commits/v0.6.0) --- updated-dependencies: - dependency-name: golang.org/x/crypto dependency-type: direct:production update-type: version-update:semver-minor ... * Upgrade go from 1.16 to 1.20 * Added integration test: HTTP, TCP, SSH * Added Makefile Improve README.md * Fixed unit test CI * Fixed go-version * Added integration test into C.I. actions --------- Signed-off-by: Mario Candela <m4r10.php@gmail.com> Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
This commit is contained in:
113
plugins/openai-gpt.go
Normal file
113
plugins/openai-gpt.go
Normal file
@ -0,0 +1,113 @@
|
||||
package plugins
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/go-resty/resty/v2"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const ChatGPTPluginName = "OpenAIGPTLinuxTerminal"
|
||||
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 promptVirtualizeLinuxTerminal = "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(promptVirtualizeLinuxTerminal)
|
||||
|
||||
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
plugins/openai-gpt_test.go
Normal file
82
plugins/openai-gpt_test.go
Normal file
@ -0,0 +1,82 @@
|
||||
package plugins
|
||||
|
||||
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)
|
||||
}
|
Reference in New Issue
Block a user