route llm_prompt through shared provider layer

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)

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{}

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.1
-	github.com/sashabaranov/go-openai v1.37.0
+	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
@@ -42,7 +41,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 +52,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
 )

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,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.1 h1:PWj1KjIASxt8icxommH72C0TQqNOvGkcSODRkiq0SQw=
-github.com/projectdiscovery/utils v0.11.1/go.mod h1:yktGrHGk2CTjNiccXovnvGrLHX9sV2bqz9nSnbA3V8M=
+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.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=
@@ -107,8 +103,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 +111,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 +125,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=

llm/llm.go

@@ -1,48 +1,138 @@
+// 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"
+	"strings"
+	"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. 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"
+	envProvider = "LLM_PROVIDER"
+	envModel    = "LLM_MODEL"
 
-func init() {
-	openaiToken := os.Getenv("OPENAI_API_KEY")
-	if openaiToken != "" {
-		client = openai.NewClient(openaiToken)
+	// 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"
+
+	defaultProvider       = "openai"
+	defaultModel          = "gpt-4o-mini"
+	defaultTimeout        = 30 * time.Second
+	defaultMaxConcurrency = 4
+	defaultMaxCalls       = 1024
+)
+
+// 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) {
+	cfg := newConfig(model)
+	key := cfg.Provider + "\x00" + cfg.BaseURL + "\x00" + model
+
+	clientsMu.Lock()
+	defer clientsMu.Unlock()
+
+	if c, ok := clients[key]; ok {
+		return c, nil
+	}
+
+	c, err := utilsllm.New(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	clients[key] = c
+
+	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 {
+	if key := os.Getenv(envAPIKey); key != "" {
+		return key
+	}
+
+	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) {
-	if client == nil {
-		return "", errors.New("no token defined")
-	}
-
-	resp, err := client.CreateChatCompletion(
-		context.Background(),
-		openai.ChatCompletionRequest{
-			Model: model,
-			Messages: []openai.ChatCompletionMessage{
-				{
-					Role:    openai.ChatMessageRoleUser,
-					Content: prompt,
-				},
-			},
-		},
-	)
+	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) {
+	model = resolveModel(model)
+
+	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
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
+	defer cancel()
 
-	return data, nil
+	return c.Complete(ctx, utilsllm.Request{
+		Prompt: prompt,
+		Format: utilsllm.Format{JSON: asJSON},
+	})
 }