Add --set-context-window

README.md

@@ -41,16 +41,15 @@ Azure, featuring streaming capabilities and extensive configuration options.
   each thread, much like your experience on the OpenAI website. Each unique thread has its own history, ensuring
   relevant and coherent responses across different chat instances.
 * **Sliding window history**: To stay within token limits, the chat history automatically trims while still preserving
-  the necessary context.
+  the necessary context. The size of this window can be adjusted through the `context-window` setting. 
 * **Custom context from any source**: You can provide the GPT model with a custom context during conversation. This
   context can be piped in from any source, such as local files, standard input, or even another program. This
   flexibility allows the model to adapt to a wide range of conversational scenarios.
 * **Model listing**: Access a list of available models using the `-l` or `--list-models` flag.
 * **Thread listing**: Display a list of active threads using the `--list-threads` flag.
 * **Advanced configuration options**: The CLI supports a layered configuration system where settings can be specified
-  through default values, a `config.yaml` file, and environment variables. For quick adjustments, use the `--set-model`,
-  `--set-thread` and `--set-max-tokens` flags. To verify your current settings, use the `--config` or `-c` flag. The
-  newly added `omit_history` configuration option adds another layer of customization to your user experience.
+  through default values, a `config.yaml` file, and environment variables. For quick adjustments,
+  various `--set-<value>` flags are provided. To verify your current settings, use the `--config` or `-c` flag. 
 * **Availability Note**: This CLI supports both gpt-4 and gpt-3.5-turbo models. However, the specific ChatGPT model used
   on chat.openai.com may not be available via the OpenAI API.
 
@@ -229,8 +228,9 @@ chatgpt --config
 Executing this command will display the active configuration, including any overrides instituted by environment
 variables or the user configuration file.
 
-To facilitate convenient adjustments, the ChatGPT CLI provides two flags for swiftly modifying the `model`
-and `max_tokens` parameters in your user configured `config.yaml`. These flags are `--set-model` and `--set-max-tokens`.
+To facilitate convenient adjustments, the ChatGPT CLI provides flags for swiftly modifying the `model`, `thread`
+, `context-window` and `max_tokens` parameters in your user configured `config.yaml`. These flags are `--set-model`
+, `--set-thread`, `--set-context-window` and `--set-max-tokens`.
 
 For instance, to update the model, use the following command:
 

cmd/chatgpt/main.go

@@ -29,6 +29,7 @@ var (
 	modelName       string
 	threadName      string
 	maxTokens       int
+	contextWindow   int
 	GitCommit       string
 	GitVersion      string
 	ServiceURL      string
@@ -57,8 +58,8 @@ func main() {
 	rootCmd.PersistentFlags().StringVar(&modelName, "set-model", "", "Set a new default GPT model by specifying the model name")
 	rootCmd.PersistentFlags().StringVar(&threadName, "set-thread", "", "Set a new active thread by specifying the thread name")
 	rootCmd.PersistentFlags().StringVar(&shell, "set-completions", "", "Generate autocompletion script for your current shell")
-
 	rootCmd.PersistentFlags().IntVar(&maxTokens, "set-max-tokens", 0, "Set a new default max token size by specifying the max tokens")
+	rootCmd.PersistentFlags().IntVar(&contextWindow, "set-context-window", 0, "Set a new default context window size")
 
 	viper.AutomaticEnv()
 
@@ -99,6 +100,16 @@ func run(cmd *cobra.Command, args []string) error {
 		return nil
 	}
 
+	if cmd.Flag("set-context-window").Changed {
+		cm := configmanager.New(config.New())
+
+		if err := cm.WriteContextWindow(contextWindow); err != nil {
+			return err
+		}
+		fmt.Println("Context window successfully updated to", contextWindow)
+		return nil
+	}
+
 	if cmd.Flag("set-thread").Changed {
 		cm := configmanager.New(config.New())
 

configmanager/configmanager.go

@@ -78,6 +78,14 @@ func (c *ConfigManager) WriteMaxTokens(tokens int) error {
 	return c.configStore.Write(c.Config)
 }
 
+// WriteContextWindow updates context window in the current configuration.
+// It writes the updated configuration to the config store and returns an error if the write fails.
+func (c *ConfigManager) WriteContextWindow(window int) error {
+	c.Config.ContextWindow = window
+
+	return c.configStore.Write(c.Config)
+}
+
 // WriteModel updates the model in the current configuration.
 // It writes the updated configuration to the config store and returns an error if the write fails.
 func (c *ConfigManager) WriteModel(model string) error {

configmanager/configmanager_test.go

@@ -193,6 +193,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) {
 		os.Setenv(envPrefix+"API_KEY", "env-api-key")
 		os.Setenv(envPrefix+"MODEL", "env-model")
 		os.Setenv(envPrefix+"MAX_TOKENS", "15")
+		os.Setenv(envPrefix+"CONTEXT_WINDOW", "25")
 		os.Setenv(envPrefix+"URL", "env-url")
 		os.Setenv(envPrefix+"COMPLETIONS_PATH", "env-completions-path")
 		os.Setenv(envPrefix+"MODELS_PATH", "env-models-path")
@@ -210,6 +211,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) {
 			APIKey:           "user-api-key",
 			Model:            "user-model",
 			MaxTokens:        20,
+			ContextWindow:    30,
 			URL:              "user-url",
 			CompletionsPath:  "user-completions-path",
 			ModelsPath:       "user-models-path",
@@ -232,6 +234,7 @@ func testConfig(t *testing.T, when spec.G, it spec.S) {
 		Expect(subject.Config.APIKey).To(Equal("env-api-key"))
 		Expect(subject.Config.Model).To(Equal("env-model"))
 		Expect(subject.Config.MaxTokens).To(Equal(15))
+		Expect(subject.Config.ContextWindow).To(Equal(25))
 		Expect(subject.Config.URL).To(Equal("env-url"))
 		Expect(subject.Config.CompletionsPath).To(Equal("env-completions-path"))
 		Expect(subject.Config.ModelsPath).To(Equal("env-models-path"))
@@ -254,6 +257,14 @@ func testConfig(t *testing.T, when spec.G, it spec.S) {
 		})
 	})
 
+	it("should write the context window as expected", func() {
+		window := 12345
+		performWriteTest(mockConfigStore, defaultConfig, window, "ContextWindow", func() {
+			subject := configmanager.New(mockConfigStore).WithEnvironment()
+			subject.WriteContextWindow(window)
+		})
+	})
+
 	it("should write the model as expected", func() {
 		model := "user-model"
 		performWriteTest(mockConfigStore, defaultConfig, model, "Model", func() {

integration/integration_test.go

@@ -356,6 +356,17 @@ func testIntegration(t *testing.T, when spec.G, it spec.S) {
 			Expect(output).To(ContainSubstring("flag needs an argument: --set-max-tokens"))
 		})
 
+		it("should require an argument for the --set-context-window flag", func() {
+			command := exec.Command(binaryPath, "--set-context-window")
+			session, err := gexec.Start(command, io.Discard, io.Discard)
+			Expect(err).NotTo(HaveOccurred())
+
+			Eventually(session).Should(gexec.Exit(exitFailure))
+
+			output := string(session.Out.Contents())
+			Expect(output).To(ContainSubstring("flag needs an argument: --set-context-window"))
+		})
+
 		it("should require the chatgpt-cli folder but not an API key for the --set-model flag", func() {
 			Expect(os.Unsetenv(apiKeyEnvVar)).To(Succeed())
 
@@ -398,6 +409,20 @@ func testIntegration(t *testing.T, when spec.G, it spec.S) {
 			Expect(output).NotTo(ContainSubstring(apiKeyEnvVar))
 		})
 
+		it("should require the chatgpt-cli folder but not an API key for the --set-context-window flag", func() {
+			Expect(os.Unsetenv(apiKeyEnvVar)).To(Succeed())
+
+			command := exec.Command(binaryPath, "--set-context-window", "789")
+			session, err := gexec.Start(command, io.Discard, io.Discard)
+			Expect(err).NotTo(HaveOccurred())
+
+			Eventually(session).Should(gexec.Exit(exitFailure))
+
+			output := string(session.Out.Contents())
+			Expect(output).To(ContainSubstring(".chatgpt-cli/config.yaml: no such file or directory"))
+			Expect(output).NotTo(ContainSubstring(apiKeyEnvVar))
+		})
+
 		it("should return the expected result for the --version flag", func() {
 			output := runCommand("--version")