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

Feat: LLM Honeypot allow specifying the custom prompt #152 (#153)

Browse Source 
* implement new feature, custom prompt

* Add doc for custom prompt
This commit is contained in:
Mario Candela
2025-01-14 08:45:30 +01:00
committed by GitHub
parent f1b35e9e43
commit c3d2ff885d
7 changed files with 105 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
plugins/llm-integration.go
View File

@ -19,12 +19,13 @@ const (
)
type LLMHoneypot struct {
Histories []Message
OpenAIKey string
client *resty.Client
Protocol tracer.Protocol
Model LLMModel
Host string
Histories []Message
OpenAIKey string
client *resty.Client
Protocol tracer.Protocol
Model LLMModel
Host string
CustomPrompt string
}
type Choice struct {
@ -211,8 +212,20 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)
func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
var err error
var prompt []Message
prompt, err := buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
if llmHoneypot.CustomPrompt != "" {
prompt = append(prompt, Message{
Role: SYSTEM.String(),
Content: llmHoneypot.CustomPrompt,
})
prompt = append(prompt, Message{
Role: USER.String(),
Content: command,
})
} else {
prompt, err = buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
}
if err != nil {
return "", err