Skip to content

Commit aa2ea2c

Browse files
dogancanbakirdogancanbakir
authored
Add optional arg for model (#221)
* add optional arg for model * minor * fix test * change the default with a recent model * don't cachec
1 parent 6e83f3b commit aa2ea2c

File tree

6 files changed

+30
-12
lines changed

6 files changed

+30
-12
lines changed

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ func main() {
9393
| len(arg interface{}) int | Returns the length of the input | `len("Hello")` | `5` |
9494
| 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` |
9595
| 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` |
96-
| 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'}` |
96+
| 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'}` |
9797
| md5(input interface{}) string | Calculates the MD5 (Message Digest) hash of the input | `md5("Hello")` | `8b1a9953c4611296a827abf8c47804d7` |
9898
| mmh3(input interface{}) string | Calculates the MMH3 (MurmurHash3) hash of an input | `mmh3("Hello")` | `316307400` |
9999
| 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` |

‎dsl.go‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ import (
4949
maputils "github.com/projectdiscovery/utils/maps"
5050
randint "github.com/projectdiscovery/utils/rand"
5151
stringsutil "github.com/projectdiscovery/utils/strings"
52+
"github.com/sashabaranov/go-openai"
5253
"github.com/spaolacci/murmur3"
5354
"golang.org/x/text/cases"
5455
"golang.org/x/text/language"
@@ -1150,13 +1151,28 @@ func init() {
11501151
}
11511152
return formattedIps[0], nil
11521153
}))
1153-
MustAddFunction(NewWithPositionalArgs("llm_prompt", 1, true, func(args ...interface{}) (interface{}, error) {
1154-
prompt, ok := args[0].(string)
1155-
if !ok {
1156-
return nil, errors.New("invalid prompt")
1157-
}
1158-
return llm.Query(prompt)
1159-
}))
1154+
MustAddFunction(NewWithSingleSignature("llm_prompt",
1155+
"(prompt string, optionalModel string) string",
1156+
false,
1157+
func(args ...interface{}) (interface{}, error) {
1158+
if len(args) < 1 {
1159+
return nil, ErrInvalidDslFunction
1160+
}
1161+
1162+
prompt, ok := args[0].(string)
1163+
if !ok {
1164+
return nil, errors.New("invalid prompt")
1165+
}
1166+
1167+
model := openai.GPT4oMini // default model
1168+
if len(args) == 2 {
1169+
if model, ok = args[1].(string); !ok {
1170+
return nil, errors.New("invalid model")
1171+
}
1172+
}
1173+
1174+
return llm.Query(prompt, model)
1175+
}))
11601176
MustAddFunction(NewWithPositionalArgs("unpack", 2, false, func(args ...interface{}) (interface{}, error) {
11611177
// format as string (ref: https://docs.python.org/3/library/struct.html#format-characters)
11621178
format, ok := args[0].(string)

‎dsl_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
137137
len(arg1 interface{}) interface{}
138138
line_ends_with(str string, suffix ...string) bool
139139
line_starts_with(str string, prefix ...string) bool
140-
llm_prompt(arg1 interface{}) interface{}
140+
llm_prompt(prompt string, optionalModel string) string
141141
md5(arg1 interface{}) interface{}
142142
mmh3(arg1 interface{}) interface{}
143143
oct_to_dec(arg1 interface{}) interface{}

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ require (
1414
github.com/projectdiscovery/gostruct v0.0.2
1515
github.com/projectdiscovery/mapcidr v1.1.34
1616
github.com/projectdiscovery/utils v0.4.11
17-
github.com/sashabaranov/go-openai v1.14.2
17+
github.com/sashabaranov/go-openai v1.37.0
1818
github.com/spaolacci/murmur3 v1.1.0
1919
github.com/stretchr/testify v1.9.0
2020
golang.org/x/exp v0.0.0-20230315142452-642cacee5cc0

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7
7676
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU=
7777
github.com/sashabaranov/go-openai v1.14.2 h1:5DPTtR9JBjKPJS008/A409I5ntFhUPPGCmaAihcPRyo=
7878
github.com/sashabaranov/go-openai v1.14.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
79+
github.com/sashabaranov/go-openai v1.37.0 h1:hQQowgYm4OXJ1Z/wTrE+XZaO20BYsL0R3uRPSpfNZkY=
80+
github.com/sashabaranov/go-openai v1.37.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
7981
github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
8082
github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI=
8183
github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=

‎llm/llm.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,15 @@ func init() {
1717
}
1818
}
1919

20-
func Query(prompt string) (string, error) {
20+
func Query(prompt, model string) (string, error) {
2121
if client == nil {
2222
return "", errors.New("no token defined")
2323
}
2424

2525
resp, err := client.CreateChatCompletion(
2626
context.Background(),
2727
openai.ChatCompletionRequest{
28-
Model: openai.GPT3Dot5Turbo,
28+
Model: model,
2929
Messages: []openai.ChatCompletionMessage{
3030
{
3131
Role: openai.ChatMessageRoleUser,

0 commit comments

Comments
 (0)