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: Refactoring plugin:LLM honeypot custom prompt (#154)

Browse Source 
refactoring LLM honeypot custom prompt
This commit is contained in:
Mario Candela
2025-01-16 08:46:13 +01:00
committed by GitHub
parent c3d2ff885d
commit 99c7287c02
2 changed files with 53 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
plugins/llm-integration_test.go
View File

@ -16,8 +16,13 @@ func TestBuildPromptEmptyHistory(t *testing.T) {
var histories []Message
command := "pwd"
honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
}
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
prompt, err := honeypot.buildPrompt(command)
//Then
assert.Nil(t, err)
@ -35,14 +40,45 @@ func TestBuildPromptWithHistory(t *testing.T) {
command := "pwd"
honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
}
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
prompt, err := honeypot.buildPrompt(command)
//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen+1, len(prompt))
}
func TestBuildPromptWithCustomPrompt(t *testing.T) {
//Given
var histories = []Message{
{
Role: "cat hello.txt",
Content: "world",
},
}
command := "pwd"
honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
CustomPrompt: "act as calculator",
}
//When
prompt, err := honeypot.buildPrompt(command)
//Then
assert.Nil(t, err)
assert.Equal(t, prompt[0].Content, "act as calculator")
assert.Equal(t, prompt[0].Role, SYSTEM.String())
}
func TestBuildExecuteModelFailValidation(t *testing.T) {
llmHoneypot := LLMHoneypot{