From fde0e46cbd45c64cff850d60eaf78de45537b320 Mon Sep 17 00:00:00 2001
From: Sakurasan <26715255+Sakurasan@users.noreply.github.com>
Date: Tue, 3 Dec 2024 16:28:21 +0800
Subject: [PATCH] update gemini model

---
 pkg/tokenizer/tokenizer.go | 2 ++
 router/chat.go             | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/pkg/tokenizer/tokenizer.go b/pkg/tokenizer/tokenizer.go
index 3151bd3..afd660b 100644
--- a/pkg/tokenizer/tokenizer.go
+++ b/pkg/tokenizer/tokenizer.go
@@ -179,6 +179,8 @@ func Cost(model string, promptCount, completionCount int) float64 {
 		cost = (0.0035/1000)*float64(prompt) + (0.0105/1000)*float64(completion)
 	case "gemini-1.5-flash-latest":
 		cost = (0.00035/1000)*float64(prompt) + (0.00053/1000)*float64(completion)
+	case "learnlm-1.5-pro-experimental", " gemini-exp-1114", "gemini-exp-1121":
+		cost = (0.00035/1000)*float64(prompt) + (0.00053/1000)*float64(completion)
 
 	// Mistral AI
 	// https://docs.mistral.ai/platform/pricing/
diff --git a/router/chat.go b/router/chat.go
index 5292a9f..fa3d2b7 100644
--- a/router/chat.go
+++ b/router/chat.go
@@ -28,7 +28,7 @@ func ChatHandler(c *gin.Context) {
 		return
 	}
 
-	if strings.HasPrefix(chatreq.Model, "gemini") {
+	if strings.HasPrefix(chatreq.Model, "gemini") || strings.HasPrefix(chatreq.Model, "learnlm") {
 		google.ChatProxy(c, &chatreq)
 		return
 	}