From eceddf120e91780431166fa6bebf5926ca1b8407 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Fri, 11 Sep 2026 17:11:23 +0300 Subject: [PATCH 1/3] route llm_prompt through shared provider layer --- dsl.go | 18 ++++++++--- dsl_test.go | 2 +- go.mod | 8 ++--- go.sum | 16 +++------ llm/llm.go | 93 +++++++++++++++++++++++++++++++++++++---------------- 5 files changed, 85 insertions(+), 52 deletions(-) diff --git a/dsl.go b/dsl.go index d8be979..12940e5 100644 --- a/dsl.go +++ b/dsl.go @@ -57,7 +57,6 @@ 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" @@ -1389,7 +1388,7 @@ func init() { return formattedIps[0], nil })) MustAddFunction(NewWithSingleSignature("llm_prompt", - "(prompt string, optionalModel string) string", + "(prompt string, optionalModel string, optionalJSON bool) string", false, func(args ...interface{}) (interface{}, error) { if len(args) < 1 { @@ -1401,14 +1400,23 @@ func init() { return nil, errors.New("invalid prompt") } - model := openai.GPT4oMini // default model - if len(args) == 2 { + // model and json are optional; an empty model lets the llm layer + // apply its default (or the LLM_MODEL env override) + var model string + if len(args) >= 2 { if model, ok = args[1].(string); !ok { return nil, errors.New("invalid model") } } - return llm.Query(prompt, model) + var asJSON bool + if len(args) >= 3 { + if asJSON, ok = args[2].(bool); !ok { + return nil, errors.New("invalid json flag") + } + } + + return llm.QueryJSON(prompt, model, asJSON) })) MustAddFunction(NewWithPositionalArgs("unpack", 2, true, func(args ...interface{}) (interface{}, error) { // format as string (ref: https://docs.python.org/3/library/struct.html#format-characters) diff --git a/dsl_test.go b/dsl_test.go index 2403628..a2c0f71 100644 --- a/dsl_test.go +++ b/dsl_test.go @@ -278,7 +278,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(prompt string, optionalModel string) string + llm_prompt(prompt string, optionalModel string, optionalJSON bool) string md5(arg1 interface{}) interface{} mmh3(arg1 interface{}) interface{} oct_to_dec(arg1 interface{}) interface{} diff --git a/go.mod b/go.mod index e6ec819..10306fb 100644 --- a/go.mod +++ b/go.mod @@ -15,8 +15,8 @@ require ( github.com/projectdiscovery/gologger v1.1.72 github.com/projectdiscovery/gostruct v0.0.2 github.com/projectdiscovery/mapcidr v1.1.97 - github.com/projectdiscovery/utils v0.11.1 - github.com/sashabaranov/go-openai v1.37.0 + github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f + github.com/sashabaranov/go-openai v1.42.0 // indirect github.com/spaolacci/murmur3 v1.1.0 github.com/spf13/cast v1.10.0 github.com/stretchr/testify v1.11.1 @@ -42,7 +42,6 @@ require ( github.com/json-iterator/go v1.1.12 // indirect github.com/logrusorgru/aurora/v4 v4.0.0 // indirect github.com/microcosm-cc/bluemonday v1.0.27 // indirect - github.com/miekg/dns v1.1.62 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect @@ -54,10 +53,7 @@ require ( github.com/tidwall/pretty v1.2.1 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect go.uber.org/multierr v1.11.0 // indirect - golang.org/x/mod v0.37.0 // indirect golang.org/x/net v0.57.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect - golang.org/x/tools v0.47.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index fdc93ad..72ee474 100644 --- a/go.sum +++ b/go.sum @@ -46,8 +46,6 @@ github.com/logrusorgru/aurora/v4 v4.0.0/go.mod h1:lP0iIa2nrnT/qoFXcOZSrZQpJ1o6n2 github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= github.com/miekg/dns v1.1.35/go.mod h1:KNUDUusw/aVsxyTYZM1oqvCicbwhgbNgztCETuNZ7xM= -github.com/miekg/dns v1.1.62 h1:cN8OuEF1/x5Rq6Np+h1epln8OiyPWV+lROx9LxcGgIQ= -github.com/miekg/dns v1.1.62/go.mod h1:mvDlcItzm+br7MToIKqkglaGhlFMHJ9DTNNWONWXbNQ= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -68,16 +66,16 @@ github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e h1:o+ul github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e/go.mod h1:xH7bPwHxUlz1yx9UlVeTF+UVCUaKhTnZgaxHb5z362E= github.com/projectdiscovery/mapcidr v1.1.97 h1:7FkxNNVXp+m1rIu5Nv/2SrF9k4+LwP8QuWs2puwy+2w= github.com/projectdiscovery/mapcidr v1.1.97/go.mod h1:9dgTJh1SP02gYZdpzMjm6vtYFkEHQHoTyaVNvaeJ7lA= -github.com/projectdiscovery/utils v0.11.1 h1:PWj1KjIASxt8icxommH72C0TQqNOvGkcSODRkiq0SQw= -github.com/projectdiscovery/utils v0.11.1/go.mod h1:yktGrHGk2CTjNiccXovnvGrLHX9sV2bqz9nSnbA3V8M= +github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f h1:KW4F+sCsu66wRQ4O28S5N5I3hnpQTgQeJPJJ5oJ+a7A= +github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f/go.mod h1:ue+Ucoa8kw9lSH+Xc2VIBjbwbwSFZ5lGXYCxzt7zRnM= github.com/projectdiscovery/wappalyzergo v0.2.94 h1:uHyNIb5OFLzyvEd6v64JLoc7Dlg3dBSjlPyjCIz0+T8= github.com/projectdiscovery/wappalyzergo v0.2.94/go.mod h1:E2p8L90ysTTUkU5FUOgGoCQ7ucEUgqJ/tvYgiGGAlEM= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7AwssoOcM/tq5JjjG2yYOc8odClEiXA= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU= -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/sashabaranov/go-openai v1.42.0 h1:fgeZx7/D8dRT//PwXAGe9ylOMtj6vrs999uWF71K+f8= +github.com/sashabaranov/go-openai v1.42.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= @@ -107,8 +105,6 @@ golang.org/x/crypto v0.0.0-20200510223506-06a226fb4e37/go.mod h1:LzIPMQfyMNhhGPh golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b h1:M2rDM6z3Fhozi9O7NWsxAkg/yqS/lQJ6PmkyIV3YP+o= golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b/go.mod h1:3//PLf8L/X+8b4vuAfHzxeRUl04Adcb341+IGKfnqS8= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= -golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= -golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20190923162816-aa69164e4478/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= @@ -117,8 +113,6 @@ golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= -golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -133,8 +127,6 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= golang.org/x/tools v0.0.0-20191216052735-49a3e744a425/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= -golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= -golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= diff --git a/llm/llm.go b/llm/llm.go index a617336..ccf4e42 100644 --- a/llm/llm.go +++ b/llm/llm.go @@ -1,48 +1,85 @@ +// Package llm backs the llm_prompt DSL helper. It delegates to the shared +// provider layer (github.com/projectdiscovery/utils/llm) so the helper gains +// provider choice, local-model support, caching, and a timeout without +// embedding its own client. package llm import ( "context" - "errors" "os" + "sync" + "time" - "github.com/sashabaranov/go-openai" + utilsllm "github.com/projectdiscovery/utils/llm" ) -var client *openai.Client +// Configuration comes from the environment, since a DSL helper has no other +// place to carry it. Defaults target OpenAI with gpt-4o-mini so existing use is +// unchanged; setting LLM_BASE_URL points the helper at a local or self-hosted +// model. +const ( + envBaseURL = "LLM_BASE_URL" + envModel = "LLM_MODEL" -func init() { - openaiToken := os.Getenv("OPENAI_API_KEY") - if openaiToken != "" { - client = openai.NewClient(openaiToken) + defaultModel = "gpt-4o-mini" + defaultTimeout = 30 * time.Second +) + +// clients are memoized per (base-url, model) so caching persists across calls +// within a process rather than resetting on every llm_prompt invocation. +var ( + clientsMu sync.Mutex + clients = map[string]*utilsllm.Client{} +) + +func client(model string) (*utilsllm.Client, error) { + baseURL := os.Getenv(envBaseURL) + key := baseURL + "\x00" + model + + clientsMu.Lock() + defer clientsMu.Unlock() + + if c, ok := clients[key]; ok { + return c, nil + } + + c, err := utilsllm.New(utilsllm.Config{ + BaseURL: baseURL, + Model: model, + Cache: true, + Timeout: defaultTimeout, + }) + if err != nil { + return nil, err } + + clients[key] = c + + return c, nil } +// Query runs a prompt and returns the completion. The model falls back to +// LLM_MODEL, then to a default, so callers that pass no model still work. func Query(prompt, model string) (string, error) { - if client == nil { - return "", errors.New("no token defined") + return QueryJSON(prompt, model, false) +} + +// QueryJSON is Query with an optional request for a JSON response, used when a +// DSL payload needs parseable output rather than prose. +func QueryJSON(prompt, model string, asJSON bool) (string, error) { + if model == "" { + if model = os.Getenv(envModel); model == "" { + model = defaultModel + } } - resp, err := client.CreateChatCompletion( - context.Background(), - openai.ChatCompletionRequest{ - Model: model, - Messages: []openai.ChatCompletionMessage{ - { - Role: openai.ChatMessageRoleUser, - Content: prompt, - }, - }, - }, - ) + c, err := client(model) if err != nil { return "", err } - if len(resp.Choices) == 0 { - return "", errors.New("no data") - } - - data := resp.Choices[0].Message.Content - - return data, nil + return c.Complete(context.Background(), utilsllm.Request{ + Prompt: prompt, + Format: utilsllm.Format{JSON: asJSON}, + }) } From 15b81e078ac408e83a3daebe6a930a01eded589b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Fri, 11 Sep 2026 17:17:35 +0300 Subject: [PATCH 2/3] fall back to OPENAI_API_KEY for backward compatibility --- go.mod | 2 +- go.sum | 4 ++-- llm/llm.go | 16 ++++++++++++++++ 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 10306fb..4dce8ab 100644 --- a/go.mod +++ b/go.mod @@ -15,7 +15,7 @@ require ( github.com/projectdiscovery/gologger v1.1.72 github.com/projectdiscovery/gostruct v0.0.2 github.com/projectdiscovery/mapcidr v1.1.97 - github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f + github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304 github.com/sashabaranov/go-openai v1.42.0 // indirect github.com/spaolacci/murmur3 v1.1.0 github.com/spf13/cast v1.10.0 diff --git a/go.sum b/go.sum index 72ee474..0d30efc 100644 --- a/go.sum +++ b/go.sum @@ -66,8 +66,8 @@ github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e h1:o+ul github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e/go.mod h1:xH7bPwHxUlz1yx9UlVeTF+UVCUaKhTnZgaxHb5z362E= github.com/projectdiscovery/mapcidr v1.1.97 h1:7FkxNNVXp+m1rIu5Nv/2SrF9k4+LwP8QuWs2puwy+2w= github.com/projectdiscovery/mapcidr v1.1.97/go.mod h1:9dgTJh1SP02gYZdpzMjm6vtYFkEHQHoTyaVNvaeJ7lA= -github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f h1:KW4F+sCsu66wRQ4O28S5N5I3hnpQTgQeJPJJ5oJ+a7A= -github.com/projectdiscovery/utils v0.11.4-0.20260911140400-c49f0d35e94f/go.mod h1:ue+Ucoa8kw9lSH+Xc2VIBjbwbwSFZ5lGXYCxzt7zRnM= +github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304 h1:WRbLaej0c5ZDlF5LWNRE1gdREm8Te+kT2nahqItEMTs= +github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304/go.mod h1:ue+Ucoa8kw9lSH+Xc2VIBjbwbwSFZ5lGXYCxzt7zRnM= github.com/projectdiscovery/wappalyzergo v0.2.94 h1:uHyNIb5OFLzyvEd6v64JLoc7Dlg3dBSjlPyjCIz0+T8= github.com/projectdiscovery/wappalyzergo v0.2.94/go.mod h1:E2p8L90ysTTUkU5FUOgGoCQ7ucEUgqJ/tvYgiGGAlEM= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= diff --git a/llm/llm.go b/llm/llm.go index ccf4e42..6345b0a 100644 --- a/llm/llm.go +++ b/llm/llm.go @@ -21,6 +21,11 @@ const ( envBaseURL = "LLM_BASE_URL" envModel = "LLM_MODEL" + // envAPIKey is the shared layer's key var; envLegacyAPIKey keeps the old + // llm_prompt behaviour working for callers that still set OPENAI_API_KEY. + envAPIKey = "LLM_API_KEY" + envLegacyAPIKey = "OPENAI_API_KEY" + defaultModel = "gpt-4o-mini" defaultTimeout = 30 * time.Second ) @@ -46,6 +51,7 @@ func client(model string) (*utilsllm.Client, error) { c, err := utilsllm.New(utilsllm.Config{ BaseURL: baseURL, Model: model, + APIKey: apiKey(), Cache: true, Timeout: defaultTimeout, }) @@ -58,6 +64,16 @@ func client(model string) (*utilsllm.Client, error) { return c, nil } +// apiKey resolves the provider key, preferring the shared LLM_API_KEY and +// falling back to OPENAI_API_KEY so existing llm_prompt setups keep working. +func apiKey() string { + if key := os.Getenv(envAPIKey); key != "" { + return key + } + + return os.Getenv(envLegacyAPIKey) +} + // Query runs a prompt and returns the completion. The model falls back to // LLM_MODEL, then to a default, so callers that pass no model still work. func Query(prompt, model string) (string, error) { From c4a48a5ef83a3736aaf759bcd1dc6541059cb478 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Fri, 11 Sep 2026 18:02:05 +0200 Subject: [PATCH 3/3] harden llm --- go.mod | 3 +- go.sum | 6 +-- llm/llm.go | 85 ++++++++++++++++++++--------- llm/llm_test.go | 140 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 204 insertions(+), 30 deletions(-) create mode 100644 llm/llm_test.go diff --git a/go.mod b/go.mod index 4dce8ab..0aed91d 100644 --- a/go.mod +++ b/go.mod @@ -15,8 +15,7 @@ require ( github.com/projectdiscovery/gologger v1.1.72 github.com/projectdiscovery/gostruct v0.0.2 github.com/projectdiscovery/mapcidr v1.1.97 - github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304 - github.com/sashabaranov/go-openai v1.42.0 // indirect + github.com/projectdiscovery/utils v0.11.4-0.20260911153819-02811729e1f0 github.com/spaolacci/murmur3 v1.1.0 github.com/spf13/cast v1.10.0 github.com/stretchr/testify v1.11.1 diff --git a/go.sum b/go.sum index 0d30efc..4ad9919 100644 --- a/go.sum +++ b/go.sum @@ -66,16 +66,14 @@ github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e h1:o+ul github.com/projectdiscovery/govaluate v0.0.0-20260504230327-80320480bb6e/go.mod h1:xH7bPwHxUlz1yx9UlVeTF+UVCUaKhTnZgaxHb5z362E= github.com/projectdiscovery/mapcidr v1.1.97 h1:7FkxNNVXp+m1rIu5Nv/2SrF9k4+LwP8QuWs2puwy+2w= github.com/projectdiscovery/mapcidr v1.1.97/go.mod h1:9dgTJh1SP02gYZdpzMjm6vtYFkEHQHoTyaVNvaeJ7lA= -github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304 h1:WRbLaej0c5ZDlF5LWNRE1gdREm8Te+kT2nahqItEMTs= -github.com/projectdiscovery/utils v0.11.4-0.20260911141643-50ff53a9f304/go.mod h1:ue+Ucoa8kw9lSH+Xc2VIBjbwbwSFZ5lGXYCxzt7zRnM= +github.com/projectdiscovery/utils v0.11.4-0.20260911153819-02811729e1f0 h1:yS19Rl7iJAwcGdYD4I+DSyEGgRVUr+VfM1FYHfFiWo0= +github.com/projectdiscovery/utils v0.11.4-0.20260911153819-02811729e1f0/go.mod h1:HMxhxLigsAr+M9Oa8n9Z0ROZcs8wAJHGsgR7UGb+oUE= github.com/projectdiscovery/wappalyzergo v0.2.94 h1:uHyNIb5OFLzyvEd6v64JLoc7Dlg3dBSjlPyjCIz0+T8= github.com/projectdiscovery/wappalyzergo v0.2.94/go.mod h1:E2p8L90ysTTUkU5FUOgGoCQ7ucEUgqJ/tvYgiGGAlEM= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d h1:hrujxIzL1woJ7AwssoOcM/tq5JjjG2yYOc8odClEiXA= github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU= -github.com/sashabaranov/go-openai v1.42.0 h1:fgeZx7/D8dRT//PwXAGe9ylOMtj6vrs999uWF71K+f8= -github.com/sashabaranov/go-openai v1.42.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 index 6345b0a..04772e5 100644 --- a/llm/llm.go +++ b/llm/llm.go @@ -7,6 +7,7 @@ package llm import ( "context" "os" + "strings" "sync" "time" @@ -14,32 +15,37 @@ import ( ) // Configuration comes from the environment, since a DSL helper has no other -// place to carry it. Defaults target OpenAI with gpt-4o-mini so existing use is -// unchanged; setting LLM_BASE_URL points the helper at a local or self-hosted -// model. +// place to carry it. With no LLM_BASE_URL or LLM_PROVIDER, the helper keeps the +// historical OpenAI + gpt-4o-mini default so existing OPENAI_API_KEY setups +// keep working. const ( - envBaseURL = "LLM_BASE_URL" - envModel = "LLM_MODEL" + envBaseURL = "LLM_BASE_URL" + envProvider = "LLM_PROVIDER" + envModel = "LLM_MODEL" // envAPIKey is the shared layer's key var; envLegacyAPIKey keeps the old // llm_prompt behaviour working for callers that still set OPENAI_API_KEY. envAPIKey = "LLM_API_KEY" envLegacyAPIKey = "OPENAI_API_KEY" - defaultModel = "gpt-4o-mini" - defaultTimeout = 30 * time.Second + defaultProvider = "openai" + defaultModel = "gpt-4o-mini" + defaultTimeout = 30 * time.Second + defaultMaxConcurrency = 4 + defaultMaxCalls = 1024 ) -// clients are memoized per (base-url, model) so caching persists across calls -// within a process rather than resetting on every llm_prompt invocation. +// clients are memoized per (provider, base-url, model) so caching persists +// across calls within a process rather than resetting on every llm_prompt +// invocation. var ( clientsMu sync.Mutex clients = map[string]*utilsllm.Client{} ) func client(model string) (*utilsllm.Client, error) { - baseURL := os.Getenv(envBaseURL) - key := baseURL + "\x00" + model + cfg := newConfig(model) + key := cfg.Provider + "\x00" + cfg.BaseURL + "\x00" + model clientsMu.Lock() defer clientsMu.Unlock() @@ -48,13 +54,7 @@ func client(model string) (*utilsllm.Client, error) { return c, nil } - c, err := utilsllm.New(utilsllm.Config{ - BaseURL: baseURL, - Model: model, - APIKey: apiKey(), - Cache: true, - Timeout: defaultTimeout, - }) + c, err := utilsllm.New(cfg) if err != nil { return nil, err } @@ -64,6 +64,34 @@ func client(model string) (*utilsllm.Client, error) { return c, nil } +func newConfig(model string) utilsllm.Config { + cfg := utilsllm.Config{ + Model: model, + APIKey: apiKey(), + Cache: true, + Timeout: defaultTimeout, + MaxConcurrency: defaultMaxConcurrency, + MaxCalls: defaultMaxCalls, + } + + if base := strings.TrimSpace(os.Getenv(envBaseURL)); base != "" { + cfg.BaseURL = base + } + if provider := strings.TrimSpace(os.Getenv(envProvider)); provider != "" { + cfg.Provider = provider + } else if cfg.BaseURL == "" { + cfg.Provider = defaultProvider + } + + return cfg +} + +func resetClients() { + clientsMu.Lock() + defer clientsMu.Unlock() + clients = map[string]*utilsllm.Client{} +} + // apiKey resolves the provider key, preferring the shared LLM_API_KEY and // falling back to OPENAI_API_KEY so existing llm_prompt setups keep working. func apiKey() string { @@ -74,6 +102,16 @@ func apiKey() string { return os.Getenv(envLegacyAPIKey) } +func resolveModel(model string) string { + if model != "" { + return model + } + if model = os.Getenv(envModel); model != "" { + return model + } + return defaultModel +} + // Query runs a prompt and returns the completion. The model falls back to // LLM_MODEL, then to a default, so callers that pass no model still work. func Query(prompt, model string) (string, error) { @@ -83,18 +121,17 @@ func Query(prompt, model string) (string, error) { // QueryJSON is Query with an optional request for a JSON response, used when a // DSL payload needs parseable output rather than prose. func QueryJSON(prompt, model string, asJSON bool) (string, error) { - if model == "" { - if model = os.Getenv(envModel); model == "" { - model = defaultModel - } - } + model = resolveModel(model) c, err := client(model) if err != nil { return "", err } - return c.Complete(context.Background(), utilsllm.Request{ + ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout) + defer cancel() + + return c.Complete(ctx, utilsllm.Request{ Prompt: prompt, Format: utilsllm.Format{JSON: asJSON}, }) diff --git a/llm/llm_test.go b/llm/llm_test.go new file mode 100644 index 0000000..b923153 --- /dev/null +++ b/llm/llm_test.go @@ -0,0 +1,140 @@ +package llm + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func okBody(content string) string { + return fmt.Sprintf(`{"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":%q}}]}`, content) +} + +func stubServer(t *testing.T, handler http.HandlerFunc) string { + t.Helper() + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + return server.URL + "/v1" +} + +func withFreshClients(t *testing.T) { + t.Helper() + resetClients() + t.Cleanup(resetClients) +} + +func TestDefaultProviderRequiresKey(t *testing.T) { + withFreshClients(t) + t.Setenv(envBaseURL, "") + t.Setenv(envProvider, "") + t.Setenv(envAPIKey, "") + t.Setenv(envLegacyAPIKey, "") + + _, err := Query("hi", "") + require.ErrorContains(t, err, "llm api key required") +} + +func TestLocalProviderDoesNotRequireKey(t *testing.T) { + withFreshClients(t) + t.Setenv(envBaseURL, "") + t.Setenv(envProvider, "ollama") + t.Setenv(envAPIKey, "") + t.Setenv(envLegacyAPIKey, "") + + _, err := client(defaultModel) + require.NoError(t, err) +} + +func TestLegacyAPIKeyReachesAuthorizationHeader(t *testing.T) { + withFreshClients(t) + t.Setenv(envAPIKey, "") + t.Setenv(envLegacyAPIKey, "legacy-key") + t.Setenv(envProvider, "") + + var gotAuth, gotModel string + var wantJSON bool + t.Setenv(envBaseURL, stubServer(t, func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + raw, _ := io.ReadAll(r.Body) + var payload struct { + Model string `json:"model"` + ResponseFormat *struct { + Type string `json:"type"` + } `json:"response_format"` + } + _ = json.Unmarshal(raw, &payload) + gotModel = payload.Model + wantJSON = payload.ResponseFormat != nil && payload.ResponseFormat.Type == "json_object" + _, _ = w.Write([]byte(okBody("ok"))) + })) + + got, err := QueryJSON("hi", "", true) + require.NoError(t, err) + require.Equal(t, "ok", got) + require.Equal(t, "Bearer legacy-key", gotAuth) + require.Equal(t, defaultModel, gotModel) + require.True(t, wantJSON) +} + +func TestLLMAPIKeyWinsOverLegacy(t *testing.T) { + withFreshClients(t) + t.Setenv(envAPIKey, "shared-key") + t.Setenv(envLegacyAPIKey, "legacy-key") + t.Setenv(envProvider, "") + + var gotAuth string + t.Setenv(envBaseURL, stubServer(t, func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + _, _ = w.Write([]byte(okBody("ok"))) + })) + + _, err := Query("hi", "custom-model") + require.NoError(t, err) + require.Equal(t, "Bearer shared-key", gotAuth) +} + +func TestEnvModelOverride(t *testing.T) { + withFreshClients(t) + t.Setenv(envAPIKey, "k") + t.Setenv(envModel, "env-model") + t.Setenv(envProvider, "") + + var gotModel string + t.Setenv(envBaseURL, stubServer(t, func(w http.ResponseWriter, r *http.Request) { + raw, _ := io.ReadAll(r.Body) + var payload struct { + Model string `json:"model"` + } + _ = json.Unmarshal(raw, &payload) + gotModel = payload.Model + _, _ = w.Write([]byte(okBody("ok"))) + })) + + _, err := Query("hi", "") + require.NoError(t, err) + require.Equal(t, "env-model", gotModel) +} + +func TestNewConfigDefaults(t *testing.T) { + t.Setenv(envBaseURL, "") + t.Setenv(envProvider, "") + cfg := newConfig("m") + require.Equal(t, defaultProvider, cfg.Provider) + require.Equal(t, "m", cfg.Model) + require.Equal(t, defaultMaxConcurrency, cfg.MaxConcurrency) + require.Equal(t, defaultMaxCalls, cfg.MaxCalls) + require.True(t, cfg.Cache) +} + +func TestNewConfigRespectsProviderAndBaseURL(t *testing.T) { + t.Setenv(envProvider, "ollama") + t.Setenv(envBaseURL, "http://host:9/v1") + cfg := newConfig("m") + require.Equal(t, "ollama", cfg.Provider) + require.Equal(t, "http://host:9/v1", cfg.BaseURL) +}