Add optional arg for model (#221)

* add optional arg for model * minor * fix test * change the default with a recent model * don't cachec
projectdiscovery · Feb 20, 2025 · aa2ea2c · aa2ea2c
1 parent 6e83f3b
commit aa2ea2c
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -93,7 +93,7 @@ func main() {
 | len(arg interface{}) int                                              | Returns the length of the input                                                                                     | `len("Hello")`                                                                                                                                       | `5`                                                                                                                                                                                                                                                                                                                                                                                        |
 | line_ends_with(str string, suffix ...string) bool                     | Checks if any line of the string ends with any of the provided substrings                                           | `line_ends_with("Hello\nHi", "lo")`                                                                                                                  | `true`                                                                                                                                                                                                                                                                                                                                                                                     |
 | line_starts_with(str string, prefix ...string) bool                   | Checks if any line of the string starts with any of the provided substrings                                         | `line_starts_with("Hi\nHello", "He")`                                                                                                                | `true`                                                                                                                                                                                                                                                                                                                                                                                     |
-| llm_prompt(str string) string | Query OpenAI LLM (default GPT 3.5) with the provided text prompt and result the result as string (requires api token as environment variable `OPENAI_API_KEY`) | `llm_prompt("produce a generic json")` | `{'a':'b'}` |
+| llm_prompt(str string, optionalModel string) string | Query OpenAI LLM (default GPT 3.5, supports o1-mini, o1-mini-2024-09-12, o1-preview, o1-preview-2024-09-12, o1, o1-2024-12-17, o3-mini, o3-mini-2025-01-31, gpt-4-32k-0613, gpt-4-32k-0314, gpt-4-32k, gpt-4-0613, gpt-4-0314, gpt-4o, gpt-4o-2024-05-13, gpt-4o-2024-08-06, gpt-4o-2024-11-20, chatgpt-4o-latest, gpt-4o-mini, gpt-4o-mini-2024-07-18, gpt-4-turbo, gpt-4-turbo-2024-04-09, gpt-4-0125-preview, gpt-4-1106-preview, gpt-4-turbo-preview, gpt-4-vision-preview, gpt-4, gpt-3.5-turbo-0125, gpt-3.5-turbo-1106, gpt-3.5-turbo-0613, gpt-3.5-turbo-0301, gpt-3.5-turbo-16k, gpt-3.5-turbo-16k-0613, gpt-3.5-turbo, gpt-3.5-turbo-instruct") with the provided text prompt and return the result as string (requires api token as environment variable `OPENAI_API_KEY`) | `llm_prompt("produce a generic json")` | `{'a':'b'}` |
 | md5(input interface{}) string                                         | Calculates the MD5 (Message Digest) hash of the input                                                               | `md5("Hello")`                                                                                                                                       | `8b1a9953c4611296a827abf8c47804d7`                                                                                                                                                                                                                                                                                                                                                         |
 | mmh3(input interface{}) string                                        | Calculates the MMH3 (MurmurHash3) hash of an input                                                                  | `mmh3("Hello")`                                                                                                                                      | `316307400`                                                                                                                                                                                                                                                                                                                                                                                |
 | oct_to_dec(octalNumber number &#124; string) float64                  | Transforms the input octal number into a decimal format                                                             | `oct_to_dec("0o1234567")`<br>`oct_to_dec(1234567)`                                                                                                   | `342391`                                                                                                                                                                                                                                                                                                                                                                                   |

diff --git a/dsl.go b/dsl.go
@@ -49,6 +49,7 @@ import (
 	maputils "github.com/projectdiscovery/utils/maps"
 	randint "github.com/projectdiscovery/utils/rand"
 	stringsutil "github.com/projectdiscovery/utils/strings"
+	"github.com/sashabaranov/go-openai"
 	"github.com/spaolacci/murmur3"
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
@@ -1150,13 +1151,28 @@ func init() {
 		}
 		return formattedIps[0], nil
 	}))
-	MustAddFunction(NewWithPositionalArgs("llm_prompt", 1, true, func(args ...interface{}) (interface{}, error) {
-		prompt, ok := args[0].(string)
-		if !ok {
-			return nil, errors.New("invalid prompt")
-		}
-		return llm.Query(prompt)
-	}))
+	MustAddFunction(NewWithSingleSignature("llm_prompt",
+		"(prompt string, optionalModel string) string",
+		false,
+		func(args ...interface{}) (interface{}, error) {
+			if len(args) < 1 {
+				return nil, ErrInvalidDslFunction
+			}
+
+			prompt, ok := args[0].(string)
+			if !ok {
+				return nil, errors.New("invalid prompt")
+			}
+
+			model := openai.GPT4oMini // default model
+			if len(args) == 2 {
+				if model, ok = args[1].(string); !ok {
+					return nil, errors.New("invalid model")
+				}
+			}
+
+			return llm.Query(prompt, model)
+		}))
 	MustAddFunction(NewWithPositionalArgs("unpack", 2, false, func(args ...interface{}) (interface{}, error) {
 		// format as string (ref: https://docs.python.org/3/library/struct.html#format-characters)
 		format, ok := args[0].(string)

diff --git a/dsl_test.go b/dsl_test.go
@@ -137,7 +137,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
 	len(arg1 interface{}) interface{}
 	line_ends_with(str string, suffix ...string) bool
 	line_starts_with(str string, prefix ...string) bool
-	llm_prompt(arg1 interface{}) interface{}
+	llm_prompt(prompt string, optionalModel string) string
 	md5(arg1 interface{}) interface{}
 	mmh3(arg1 interface{}) interface{}
 	oct_to_dec(arg1 interface{}) interface{}

diff --git a/go.mod b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/projectdiscovery/gostruct v0.0.2
 	github.com/projectdiscovery/mapcidr v1.1.34
 	github.com/projectdiscovery/utils v0.4.11
-	github.com/sashabaranov/go-openai v1.14.2
+	github.com/sashabaranov/go-openai v1.37.0
 	github.com/spaolacci/murmur3 v1.1.0
 	github.com/stretchr/testify v1.9.0
 	golang.org/x/exp v0.0.0-20230315142452-642cacee5cc0

diff --git a/go.sum b/go.sum
@@ -76,6 +76,8 @@ github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7
 github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU=
 github.com/sashabaranov/go-openai v1.14.2 h1:5DPTtR9JBjKPJS008/A409I5ntFhUPPGCmaAihcPRyo=
 github.com/sashabaranov/go-openai v1.14.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/sashabaranov/go-openai v1.37.0 h1:hQQowgYm4OXJ1Z/wTrE+XZaO20BYsL0R3uRPSpfNZkY=
+github.com/sashabaranov/go-openai v1.37.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
 github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=

diff --git a/llm/llm.go b/llm/llm.go
@@ -17,15 +17,15 @@ func init() {
 	}
 }
 
-func Query(prompt string) (string, error) {
+func Query(prompt, model string) (string, error) {
 	if client == nil {
 		return "", errors.New("no token defined")
 	}
 
 	resp, err := client.CreateChatCompletion(
 		context.Background(),
 		openai.ChatCompletionRequest{
-			Model: openai.GPT3Dot5Turbo,
+			Model: model,
 			Messages: []openai.ChatCompletionMessage{
 				{
 					Role:    openai.ChatMessageRoleUser,