feat(cli.go): add OpenAIChunkTimeout option for configurable token ch…

…unk timeout
modernice · Apr 26, 2024 · b75d957 · b75d957
1 parent 5c8349d
commit b75d957
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/internal/cli/cli.go b/internal/cli/cli.go
@@ -49,6 +49,7 @@ type cliOptions struct {
 	OpenAITemperature    float32 `name:"temperature" help:"OpenAI temperature" env:"OPENAI_TEMPERATURE" default:"0.3"`
 	OpenAITopP           float32 `name:"top-p" help:"OpenAI top_p" env:"OPENAI_TOP_P" default:"0.3"`
 	OpenAIResponseFormat string  `name:"format" help:"OpenAI response format ('text' or 'json_object')" env:"OPENAI_RESPONSE_FORMAT" default:"text"`
+	OpenAIChunkTimeout   string  `name:"chunk-timeout" help:"Timeout for each token chunk" env:"OPENAI_CHUNK_TIMEOUT"`
 
 	Timeout time.Duration `short:"T" help:"Timeout for API requests" env:"DRAGOMAN_TIMEOUT" default:"3m"`
 	Verbose bool          `short:"v" help:"Verbose output"`
@@ -126,6 +127,14 @@ func (app *App) translate() {
 		opts = append(opts, openai.Stream(os.Stdout))
 	}
 
+	if options.OpenAIChunkTimeout != "" {
+		chunkTimeout, err := time.ParseDuration(options.OpenAIChunkTimeout)
+		if err != nil {
+			app.kong.Fatalf("invalid chunk timeout: %v", err)
+		}
+		opts = append(opts, openai.ChunkTimeout(chunkTimeout))
+	}
+
 	model := openai.New(options.OpenAIKey, opts...)
 	translator := dragoman.NewTranslator(model)