add grok

radiantspace · Nov 18, 2024 · fd6816a · fd6816a
1 parent 8c25c45
commit fd6816a
Show file tree

Hide file tree

Showing 13 changed files with 68 additions and 4 deletions.
diff --git a/.github/workflows/deploy-to-doks.yaml b/.github/workflows/deploy-to-doks.yaml
@@ -71,6 +71,7 @@ jobs:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
           CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
+          GROK_API_KEY: ${{ secrets.GROK_API_KEY }}
           SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
           SLACK_SIGNING_SECRET: ${{ secrets.SLACK_SIGNING_SECRET }}
           STRIPE_ENDPOINT_SECRET: ${{ secrets.STRIPE_ENDPOINT_SECRET }}
@@ -85,6 +86,7 @@ jobs:
           echo OPENAI_API_KEY=$OPENAI_API_KEY >> environment-properties.env
           echo FIREWORKS_API_KEY=$FIREWORKS_API_KEY >> environment-properties.env
           echo CLAUDE_API_KEY=$CLAUDE_API_KEY >> environment-properties.env
+          echo GROK_API_KEY=$GROK_API_KEY >> environment-properties.env
           echo SLACK_BOT_TOKEN=$SLACK_BOT_TOKEN >> environment-properties.env
           echo SLACK_SIGNING_SECRET=$SLACK_SIGNING_SECRET >> environment-properties.env
           echo STRIPE_ENDPOINT_SECRET=$STRIPE_ENDPOINT_SECRET >> environment-properties.env

diff --git a/backend/app/ai/ai.go b/backend/app/ai/ai.go
@@ -67,3 +67,7 @@ func IsClaudeAI(model models.Engine) bool {
 
 	return false
 }
+
+func IsGrok(model models.Engine) bool {
+	return model == models.Grok
+}
diff --git a/backend/app/ai/chat.go b/backend/app/ai/chat.go
@@ -61,6 +61,10 @@ const (
 	OPUS_INPUT_PRICE  = 15.0 / 1000000
 	OPUS_OUTPUT_PRICE = 75.0 / 1000000
 
+	// grok-beta
+	GROK_INPUT_PRICE  = 0.5 / 1000000
+	GROK_OUTPUT_PRICE = 1.5 / 1000000
+
 	CHARS_PER_TOKEN = 2.0 // average number of characters per token, must be tuned or moved to tiktoken
 )
 
@@ -236,6 +240,9 @@ func urlFromModel(model models.Engine) string {
 	if IsClaudeAI(model) {
 		return "https://api.anthropic.com/v1/messages"
 	}
+	if IsGrok(model) {
+		return "https://api.x.ai/v1/chat/completions"
+	}
 	return "https://api.openai.com/v1/chat/completions"
 }
 
@@ -246,6 +253,9 @@ func authTokenFromModel(model models.Engine) string {
 	if IsClaudeAI(model) {
 		return config.CONFIG.ClaudeAPIKey
 	}
+	if IsGrok(model) {
+		return config.CONFIG.GrokAPIKey
+	}
 
 	return config.CONFIG.OpenAIAPIKey
 }
@@ -274,6 +284,8 @@ func PricePerInputToken(model models.Engine) float64 {
 		return SONET_INPUT_PRICE
 	case models.Haiku3:
 		return HAIKU_INPUT_PRICE
+	case models.Grok:
+		return GROK_INPUT_PRICE
 	default:
 		return CHAT_INPUT_PRICE
 	}
@@ -297,6 +309,8 @@ func PricePerOutputToken(model models.Engine) float64 {
 		return SONET_OUTPUT_PRICE
 	case models.Haiku3:
 		return HAIKU_OUTPUT_PRICE
+	case models.Grok:
+		return GROK_OUTPUT_PRICE
 	default:
 		return CHAT_OUTPUT_PRICE
 	}

diff --git a/backend/app/config/config.go b/backend/app/config/config.go
@@ -48,6 +48,7 @@ type Config struct {
 	DataDogClient          *statsd.Client
 	Environment            string
 	FireworksAPIKey        string
+	GrokAPIKey             string
 	MongoDBName            string
 	MongoDBConnection      string
 	OpenAIAPIKey           string

diff --git a/backend/app/models/grok.go b/backend/app/models/grok.go
@@ -0,0 +1,5 @@
+package models
+
+const (
+	Grok Engine = "grok-beta"
+)
diff --git a/backend/app/telegram/message_processors.go b/backend/app/telegram/message_processors.go
@@ -79,7 +79,7 @@ func ProcessThreadedStreamingMessage(
 	chatIDString := util.GetChatIDString(message)
 	topicID := util.GetTopicID(message)
 
-	if ai.IsFireworksAI(engineModel) || ai.IsClaudeAI(engineModel) {
+	if ai.IsFireworksAI(engineModel) || ai.IsClaudeAI(engineModel) || ai.IsGrok(engineModel) {
 		ProcessStreamingMessageWithLocalThreads(ctx, bot, message, []models.Message{}, "", mode, engineModel, cancelContext)
 		return
 	}
@@ -140,7 +140,7 @@ func ProcessThreadedNonStreamingMessage(
 	chatIDString := util.GetChatIDString(message)
 	topicID := util.GetTopicID(message)
 
-	if ai.IsFireworksAI(engineModel) {
+	if ai.IsFireworksAI(engineModel) || ai.IsClaudeAI(engineModel) || ai.IsGrok(engineModel) {
 		ProcessChatCompleteNonStreamingMessage(ctx, bot, message, []models.Message{}, "", mode, engineModel)
 		return
 	}

diff --git a/backend/app/telegram/telegram.go b/backend/app/telegram/telegram.go
@@ -370,7 +370,7 @@ func handleCallbackQuery(bot *telego.Bot, callbackQuery telego.CallbackQuery) {
 			MessageID:   messageId,
 			ReplyMarkup: GetStatusKeyboard(ctx),
 		})
-	case string(models.ChatGpt35Turbo), string(models.ChatGpt4), string(models.ChatGpt4o), string(models.ChatGpt4oMini), string(models.ChatGpt4Turbo), string(models.ChatGpt4TurboVision), string(models.LlamaV3_8b), string(models.LlamaV3_70b), string(models.Sonet35), string(models.Haiku3), string(models.Opus3), string(models.Sonet35_241022):
+	case string(models.ChatGpt35Turbo), string(models.ChatGpt4), string(models.ChatGpt4o), string(models.ChatGpt4oMini), string(models.ChatGpt4Turbo), string(models.ChatGpt4TurboVision), string(models.LlamaV3_8b), string(models.LlamaV3_70b), string(models.Sonet35), string(models.Haiku3), string(models.Opus3), string(models.Sonet35_241022), string(models.Grok):
 		handleEngineSwitchCallbackQuery(callbackQuery, topicString)
 	case string(models.DallE3), string(models.Midjourney6), string(models.StableDiffusion3), string(models.Playground25):
 		handleImageModelSwitchCallbackQuery(callbackQuery, topicString)
@@ -594,6 +594,23 @@ func handleEngineSwitchCallbackQuery(callbackQuery telego.CallbackQuery, topicSt
 		}
 		return
 	}
+	if callbackQuery.Data == string(models.Grok) {
+		go redis.SaveModel(chatIDString, models.Grok)
+		notification := "Switched to Grok model, intelligent and fun with Web Search enabled!"
+		notification = lib.AddBotSuffixToGroupCommands(ctx, notification)
+		_, err := BOT.SendMessage(tu.Message(tu.ID(chatID), notification).WithMessageThreadID(topicID))
+		if err != nil {
+			log.Errorf("handleEngineSwitchCallbackQuery failed to send Grok message: %v", err)
+		}
+		err = BOT.AnswerCallbackQuery(&telego.AnswerCallbackQueryParams{
+			CallbackQueryID: callbackQuery.ID,
+			Text:            "Switched to Grok engine!",
+		})
+		if err != nil {
+			log.Errorf("handleEngineSwitchCallbackQuery failed to answer callback query: %v", err)
+		}
+		return
+	}
 
 	log.Errorf("Unknown engine switch callback query: %s, chat id: %s", callbackQuery.Data, chatIDString)
 }

diff --git a/backend/app/telegram/user_status.go b/backend/app/telegram/user_status.go
@@ -149,6 +149,7 @@ func GetModelsKeyboard(ctx context.Context) *telego.InlineKeyboardMarkup {
 	haiku3Active := ""
 	bigLlama3Active := ""
 	smallLlama3Active := ""
+	grokActive := ""
 	switch model {
 	case models.ChatGpt4o:
 		gpt4oActive = "✅ "
@@ -162,13 +163,21 @@ func GetModelsKeyboard(ctx context.Context) *telego.InlineKeyboardMarkup {
 		bigLlama3Active = "✅ "
 	case models.LlamaV3_8b:
 		smallLlama3Active = "✅ "
+	case models.Grok:
+		grokActive = "✅ "
 	}
 
 	return &telego.InlineKeyboardMarkup{
 		InlineKeyboard: [][]telego.InlineKeyboardButton{
 			{
 				{
-					Text:         gpt4oActive + "GPT 4o (best) 💰💰💰🏃🏃🧠🧠🧠🧠",
+					Text:         grokActive + "Grok + Web 🌍 💰💰🏃🏃🧠🧠🧠",
+					CallbackData: string(models.Grok) + ":" + topicString,
+				},
+			},
+			{
+				{
+					Text:         gpt4oActive + "GPT 4o 💰💰💰🏃🏃🧠🧠🧠🧠",
 					CallbackData: string(models.ChatGpt4o) + ":" + topicString,
 				},
 			},

diff --git a/backend/main.go b/backend/main.go
@@ -47,6 +47,7 @@ func main() {
 		OpenAIAPIKey:    util.Env("OPENAI_API_KEY"),
 		FireworksAPIKey: util.Env("FIREWORKS_API_KEY"),
 		ClaudeAPIKey:    util.Env("CLAUDE_API_KEY"),
+		GrokAPIKey:      util.Env("GROK_API_KEY"),
 		Redis: config.Redis{
 			Host:     util.Env("REDIS_HOST"),
 			Port:     "6379",

diff --git a/docker-compose.cs.yml b/docker-compose.cs.yml
@@ -17,6 +17,7 @@ services:
       BACKEND_LISTEN_URL: http://0.0.0.0:8080
       CLAUDE_API_KEY:
       OPENAI_API_KEY:
+      GROK_API_KEY:
       FIREWORKS_API_KEY:
       STRIPE_TOKEN: setup-for-testing
       TELEGRAM_BOT_TOKEN:

diff --git a/docker-compose.local.yml b/docker-compose.local.yml
@@ -18,6 +18,7 @@ services:
       BACKEND_LISTEN_ADDRESS: 0.0.0.0:8080
       BACKEND_LISTEN_URL: http://0.0.0.0:8080
       CLAUDE_API_KEY:
+      GROK_API_KEY:
       OPENAI_API_KEY:
       FIREWORKS_API_KEY:
       STRIPE_TOKEN: setup-for-testing

diff --git a/infra/kustomize/base/backend.yaml b/infra/kustomize/base/backend.yaml
@@ -85,6 +85,8 @@ spec:
               value: $(FIREWORKS_API_KEY)
             - name: CLAUDE_API_KEY
               value: $(CLAUDE_API_KEY)
+            - name: GROK_API_KEY
+              value: $(GROK_API_KEY)
             - name: BACKEND_BASE_URL
               value: $(BACKEND_BASE_URL)
             - name: PROMETHEUS_LISTEN_ADDRESS

diff --git a/infra/kustomize/kustomization.yaml b/infra/kustomize/kustomization.yaml
@@ -92,6 +92,13 @@ vars:
 - fieldref:
     fieldPath: data.CLAUDE_API_KEY
   name: CLAUDE_API_KEY
+  objref:
+    apiVersion: v1
+    kind: ConfigMap
+    name: environment-variables
+- fieldref:
+    fieldPath: data.GROK_API_KEY
+  name: GROK_API_KEY
   objref:
     apiVersion: v1
     kind: ConfigMap