From 97887035da52e3d170e35235d474a1eb4b560d3c Mon Sep 17 00:00:00 2001
From: Sakurasan <26715255+Sakurasan@users.noreply.github.com>
Date: Thu, 12 Dec 2024 16:13:56 +0800
Subject: [PATCH] update gemini model

---
 pkg/tokenizer/tokenizer.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pkg/tokenizer/tokenizer.go b/pkg/tokenizer/tokenizer.go
index afd660b..bdcc6a6 100644
--- a/pkg/tokenizer/tokenizer.go
+++ b/pkg/tokenizer/tokenizer.go
@@ -179,7 +179,9 @@ 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":
+	case "gemini-2.0-flash-exp":
+		cost = (0.00035/1000)*float64(prompt) + (0.00053/1000)*float64(completion)
+	case "learnlm-1.5-pro-experimental", " gemini-exp-1114", "gemini-exp-1121", "gemini-exp-1206":
 		cost = (0.00035/1000)*float64(prompt) + (0.00053/1000)*float64(completion)
 
 	// Mistral AI