From a50375c9605a87546a0fcbcacabe5482fdfa1c2c Mon Sep 17 00:00:00 2001
From: Samir Tahir <30797145+samirtahir91@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:16:29 +0000
Subject: [PATCH] fix: add maxTokens to serve mode (#1280)

Signed-off-by: samir-tahir <samirtahir91@gmail.com>
Co-authored-by: Alex Jones <alexsimonjones@gmail.com>
---
 cmd/serve/serve.go | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/cmd/serve/serve.go b/cmd/serve/serve.go
index be49588..f362893 100644
--- a/cmd/serve/serve.go
+++ b/cmd/serve/serve.go
@@ -30,6 +30,7 @@ const (
 	defaultTemperature float32 = 0.7
 	defaultTopP        float32 = 1.0
 	defaultTopK        int32   = 50
+	defaultMaxTokens   int     = 2048
 )
 
 var (
@@ -102,6 +103,18 @@ var ServeCmd = &cobra.Command{
 				}
 				return int32(topK)
 			}
+			maxTokens := func() int {
+				env := os.Getenv("K8SGPT_MAX_TOKENS")
+				if env == "" {
+					return defaultMaxTokens
+				}
+				maxTokens, err := strconv.ParseInt(env, 10, 32)
+				if err != nil {
+					color.Red("Unable to convert maxTokens value: %v", err)
+					os.Exit(1)
+				}
+				return int(maxTokens)
+			}
 			// Check for env injection
 			backend = os.Getenv("K8SGPT_BACKEND")
 			password := os.Getenv("K8SGPT_PASSWORD")
@@ -125,6 +138,7 @@ var ServeCmd = &cobra.Command{
 					Temperature:   temperature(),
 					TopP:          topP(),
 					TopK:          topK(),
+					MaxTokens:     maxTokens(),
 				}
 
 				configAI.Providers = append(configAI.Providers, *aiProvider)