Feat: Refactoring LLM Plugin, update docs. (#165)

Refactoring LLM Plugin, update docs.

plugins/llm-integration.go

@@ -8,13 +8,14 @@ import (
 	"github.com/mariocandela/beelzebub/v3/tracer"
 	log "github.com/sirupsen/logrus"
 	"regexp"
+	"strings"
 )
 
 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 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"
+	openAIEndpoint                      = "https://api.openai.com/v1/chat/completions"
 	ollamaEndpoint                      = "http://localhost:11434/api/chat"
 )
 
@@ -23,7 +24,8 @@ type LLMHoneypot struct {
 	OpenAIKey    string
 	client       *resty.Client
 	Protocol     tracer.Protocol
-	Model        LLMModel
+	Provider     LLMProvider
+	Model        string
 	Host         string
 	CustomPrompt string
 }
@@ -71,21 +73,21 @@ func (role Role) String() string {
 	return [...]string{"system", "user", "assistant"}[role]
 }
 
-type LLMModel int
+type LLMProvider int
 
 const (
-	LLAMA3 LLMModel = iota
-	GPT4O
+	Ollama LLMProvider = iota
+	OpenAI
 )
 
-func FromStringToLLMModel(llmModel string) (LLMModel, error) {
-	switch llmModel {
-	case "llama3":
-		return LLAMA3, nil
-	case "gpt4-o":
-		return GPT4O, nil
+func FromStringToLLMProvider(llmProvider string) (LLMProvider, error) {
+	switch strings.ToLower(llmProvider) {
+	case "ollama":
+		return Ollama, nil
+	case "openai":
+		return OpenAI, nil
 	default:
-		return -1, fmt.Errorf("model %s not found", llmModel)
+		return -1, fmt.Errorf("provider %s not found, valid providers: ollama, openai", llmProvider)
 	}
 }
 
@@ -153,7 +155,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
 	var err error
 
 	requestJson, err := json.Marshal(Request{
-		Model:    "gpt-4o",
+		Model:    llmHoneypot.Model,
 		Messages: messages,
 		Stream:   false,
 	})
@@ -166,7 +168,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
 	}
 
 	if llmHoneypot.Host == "" {
-		llmHoneypot.Host = openAIGPTEndpoint
+		llmHoneypot.Host = openAIEndpoint
 	}
 
 	log.Debug(string(requestJson))
@@ -192,7 +194,7 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)
 	var err error
 
 	requestJson, err := json.Marshal(Request{
-		Model:    "llama3",
+		Model:    llmHoneypot.Model,
 		Messages: messages,
 		Stream:   false,
 	})
@@ -229,13 +231,13 @@ func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
 		return "", err
 	}
 
-	switch llmHoneypot.Model {
-	case LLAMA3:
+	switch llmHoneypot.Provider {
+	case Ollama:
 		return llmHoneypot.ollamaCaller(prompt)
-	case GPT4O:
+	case OpenAI:
 		return llmHoneypot.openAICaller(prompt)
 	default:
-		return "", errors.New("no model selected")
+		return "", fmt.Errorf("provider %d not found, valid providers: ollama, openai", llmHoneypot.Provider)
 	}
 }