chrissewell/beelzebub
1
0
Fork 0
mirror of https://github.com/mariocandela/beelzebub.git synced 2025-07-01 18:47:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

Fix: llm plugin OpenAI generates random plaintext (hallucinations) (#163)

Browse Source 
* Add tests to adopt TDD.

* Fix bug, LLM hallucinations
This commit is contained in:
Mario Candela
2025-02-16 16:27:10 +01:00
committed by GitHub
parent db804474d3
commit 8703d1afda
2 changed files with 99 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
plugins/llm-integration.go
View File

@ -6,12 +6,12 @@ import (
"fmt"
"github.com/go-resty/resty/v2"
"github.com/mariocandela/beelzebub/v3/tracer"
log "github.com/sirupsen/logrus"
"regexp"
)
const (
systemPromptVirtualizeLinuxTerminal = "You will act as an Ubuntu Linux terminal. The user will type commands, and you are to reply with what the terminal should show. Your responses must be contained within a single code block. Do not provide explanations or type commands unless explicitly instructed by the user. Your entire response/output is going to consist of a simple text with \n for new line, and you will NOT wrap it within string md markers"
systemPromptVirtualizeLinuxTerminal = "You will act as an Ubuntu Linux terminal. The user will type commands, and you are to reply with what the terminal should show. Your responses must be contained within a single code block. Do not provide note. Do not provide explanations or type commands unless explicitly instructed by the user. Your entire response/output is going to consist of a simple text with \n for new line, and you will NOT wrap it within string md markers"
systemPromptVirtualizeHTTPServer = "You will act as an unsecure HTTP Server with multiple vulnerability like aws and git credentials stored into root http directory. The user will send HTTP requests, and you are to reply with what the server should show. Do not provide explanations or type commands unless explicitly instructed by the user."
LLMPluginName = "LLMHoneypot"
openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions"
@ -185,7 +185,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
return "", errors.New("no choices")
}
return response.Result().(*Response).Choices[0].Message.Content, nil
return removeQuotes(response.Result().(*Response).Choices[0].Message.Content), nil
}
func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error) {
@ -216,7 +216,7 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)
}
log.Debug(response)
return response.Result().(*Response).Message.Content, nil
return removeQuotes(response.Result().(*Response).Message.Content), nil
}
func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
@ -238,3 +238,8 @@ func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
return "", errors.New("no model selected")
}
}
func removeQuotes(content string) string {
regex := regexp.MustCompile("(```( *)?([a-z]*)?(\\n)?)")
return regex.ReplaceAllString(content, "")
}