Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add optional arg for model #221

Merged
merged 5 commits into from
Feb 20, 2025
Merged

Add optional arg for model #221

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
25e568e
add optional arg for model
dogancanbakir Feb 18, 2025
009a336
minor
dogancanbakir Feb 18, 2025
50ba0b8
fix test
dogancanbakir Feb 18, 2025
6d12b83
change the default with a recent model
dogancanbakir Feb 19, 2025
69581d1
don't cachec
dogancanbakir Feb 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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` |
Expand Down
30 changes: 23 additions & 7 deletions dsl.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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",
true,
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.GPT3Dot5Turbo // 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)
Expand Down
2 changes: 1 addition & 1 deletion dsl_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down
2 changes: 1 addition & 1 deletion go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
4 changes: 2 additions & 2 deletions llm/llm.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading