支持 o1 模型

taoso · Dec 8, 2024 · 623f34f · 623f34f
1 parent 384efe9
commit 623f34f
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/chat.go b/chat.go
@@ -107,7 +107,7 @@ type chatmsg struct {
 	Model     string    `json:"model"`
 	Stream    bool      `json:"stream"`
 	User      string    `json:"user"`
-	MaxTokens int       `json:"max_tokens"`
+	MaxTokens int       `json:"max_completion_tokens"`
 }
 
 func (m *chatmsg) CountToken(bpe *tiktoken.BPE) (int, error) {
@@ -336,6 +336,14 @@ func (p *Proxy) chat(w http.ResponseWriter, req *http.Request, f *FileHandler) {
 		tokenRate = 1
 		msg.Model = "gpt-4o-mini"
 		maxTokens = 4 * 1024
+	case "o1-preview":
+		tokenRate = 25
+		msg.Model = "o1-preview"
+		maxTokens = 4 * 1024
+	case "o1-mini":
+		tokenRate = 5
+		msg.Model = "o1-mini"
+		maxTokens = 4 * 1024
 	default:
 		w.WriteHeader(http.StatusBadRequest)
 		w.Write([]byte("invalid model"))