add model input & list models flags

README.md

@@ -67,8 +67,10 @@ INPUT:
    -p, -prompt string[]  prompt to query (input: stdin,string,file)
 
 MODEL:
-   -g3, -gpt3  use GPT-3.5 model (default true)
-   -g4, -gpt4  use GPT-4.0 model
+   -g3, -gpt3         use GPT-3.5 model (default true)
+   -g4, -gpt4         use GPT-4.0 model
+   -m, -model string  specify model to use (ex: gpt-4-0314)
+   -lm, -list-models  list available models
 
 CONFIG:
    -ak, -openai-api-key string    openai api key token (input: string,file,env)

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/projectdiscovery/goflags v0.1.8
 	github.com/projectdiscovery/gologger v1.1.10
 	github.com/projectdiscovery/utils v0.0.32
-	github.com/sashabaranov/go-openai v1.9.4
+	github.com/sashabaranov/go-openai v1.17.11
 )
 
 require (

go.sum

@@ -113,8 +113,8 @@ github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZV
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
 github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca h1:NugYot0LIVPxTvN8n+Kvkn6TrbMyxQiuvKdEwFdR9vI=
 github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca/go.mod h1:uugorj2VCxiV1x+LzaIdVa9b4S4qGAcH6cbhh4qVxOU=
-github.com/sashabaranov/go-openai v1.9.4 h1:KanoCEoowAI45jVXlenMCckutSRr39qOmSi9MyPBfZM=
-github.com/sashabaranov/go-openai v1.9.4/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/sashabaranov/go-openai v1.17.11 h1:XVr00J8JymJVx8Hjbh/5mG0V4PQHRarBU3v7k2x6MR0=
+github.com/sashabaranov/go-openai v1.17.11/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

internal/runner/options.go

@@ -34,6 +34,8 @@ type Options struct {
 	Prompt             string              `yaml:"prompt"`
 	Gpt3               bool                `yaml:"gpt3"`
 	Gpt4               bool                `yaml:"gpt4"`
+	Model              string              `yaml:"model"`
+	ListModels         bool                `yaml:"list_models"`
 	Update             bool                `yaml:"update"`
 	DisableUpdateCheck bool                `yaml:"disable_update_check"`
 	Output             string              `yaml:"output"`
@@ -68,6 +70,8 @@ func ParseOptions() *Options {
 	flagSet.CreateGroup("model", "Model",
 		flagSet.BoolVarP(&options.Gpt3, "gpt3", "g3", true, "use GPT-3.5 model"),
 		flagSet.BoolVarP(&options.Gpt4, "gpt4", "g4", false, "use GPT-4.0 model"),
+		flagSet.StringVarP(&options.Model, "model", "m", "", "specify model to use (ex: gpt-4-0314)"),
+		flagSet.BoolVarP(&options.ListModels, "list-models", "lm", false, "list available models"),
 	)
 
 	flagSet.CreateGroup("config", "Config",

internal/runner/runner.go

@@ -1,6 +1,7 @@
 package runner
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -40,7 +41,39 @@ func (r *Runner) Run() (*Result, error) {
 		model = openai.GPT3Dot5Turbo
 	}
 	if r.options.Gpt4 {
-		model = openai.GPT4
+		// use turbo preview by default
+		model = openai.GPT4TurboPreview
+	}
+	if r.options.Model != "" {
+		model = r.options.Model
+	}
+
+	if r.options.ListModels {
+		models, err := client.ListModels(context.Background())
+		if err != nil {
+			return &Result{}, err
+		}
+		var buff bytes.Buffer
+		for _, model := range models.Models {
+			buff.WriteString(fmt.Sprintf("%s\n", model.ID))
+		}
+
+		result := &Result{
+			Timestamp: time.Now().String(),
+			Model:     model,
+			Prompt:    r.options.Prompt,
+		}
+
+		if r.options.Stream {
+			result.SetupStreaming()
+			go func(res *Result) {
+				defer res.CloseCompletionStream()
+				res.WriteCompletionStreamResponse(buff.String())
+			}(result)
+		} else {
+			result.Completion = buff.String()
+		}
+		return result, nil
 	}
 
 	chatReq := openai.ChatCompletionRequest{