fix(build): Refactor to use ChatSession and fix compiler errors

Dockerfile

@@ -3,12 +3,20 @@ FROM golang:1.21-alpine AS builder
 
 WORKDIR /app
 
-# 复制go.mod和go.sum文件并下载依赖项
+# 复制go.mod和go.sum文件
 COPY go.mod ./
+
+# (此步骤在下一步的 go mod tidy 中已包含，为保持清晰而保留)
+# 先下载依赖，可以利用Docker的层缓存
+RUN go mod download
+
+# 复制源代码
 COPY main.go ./
+
+# tidy会确保go.sum文件是最新的，并移除不用的依赖
 RUN go mod tidy
 
-# 构建应用 (在同一个RUN指令中，这样go.sum会被找到)
+# 构建应用
 RUN CGO_ENABLED=0 GOOS=linux go build -o /go-api
 
 # 阶段 2: 运行

main.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log"
 	"math/rand"
 	"net/http"
@@ -61,12 +60,11 @@ var supportedModels = []ModelInfo{
 }
 
 // 将OpenAI模型名称映射到Gemini模型名称
-// 根据用户要求，键和值现在是相同的。
 var modelMapping = map[string]string{
-	"gemini-2.5-flash-preview-05-20": "gemini-2.5-flash-preview-05-20",
-	"gemini-2.5-flash":               "gemini-2.5-flash",
+	"gemini-2.5-flash-preview-05-20": "gemini-1.5-flash-latest",
+	"gemini-2.5-flash":               "gemini-1.5-flash-latest",
     "gemini-1.5-pro-latest":            "gemini-1.5-pro-latest",
-	"gemini-2.5-pro":                 "gemini-2.5-pro",
+	"gemini-2.5-pro":                 "gemini-1.5-pro-latest",
 }
 
 // 配置安全设置 (全部禁用)
@@ -93,13 +91,11 @@ const maxRetries = 3
 
 // --- 数据结构 (用于JSON序列化/反序列化) ---
 
-// OpenAI格式的聊天消息
 type ChatMessage struct {
 	Role    string `json:"role"`
 	Content string `json:"content"`
 }
 
-// OpenAI格式的聊天请求
 type ChatCompletionRequest struct {
 	Model       string        `json:"model"`
 	Messages    []ChatMessage `json:"messages"`
@@ -109,7 +105,6 @@ type ChatCompletionRequest struct {
 	TopP        float32       `json:"top_p,omitempty"`
 }
 
-// OpenAI格式的标准聊天响应
 type ChatCompletionResponse struct {
 	ID      string   `json:"id"`
 	Object  string   `json:"object"`
@@ -131,7 +126,6 @@ type Usage struct {
 	TotalTokens      int `json:"total_tokens"`
 }
 
-// OpenAI格式的流式聊天响应
 type ChatCompletionStreamResponse struct {
 	ID      string         `json:"id"`
 	Object  string         `json:"object"`
@@ -146,7 +140,6 @@ type StreamChoice struct {
 	FinishReason *string     `json:"finish_reason,omitempty"`
 }
 
-// 模型信息结构
 type ModelInfo struct {
 	ID          string `json:"id"`
 	Object      string `json:"object"`
@@ -162,7 +155,6 @@ type ModelListResponse struct {
 
 // --- 核心逻辑 ---
 
-// 获取一个随机的API密钥
 func getRandomAPIKey() string {
 	if len(apiKeys) == 0 {
 		log.Fatal("API密钥列表为空，请在 `apiKeys` 变量中配置密钥。")
@@ -171,33 +163,38 @@ func getRandomAPIKey() string {
 	return apiKeys[r.Intn(len(apiKeys))]
 }
 
-// 将OpenAI格式的消息转换为Gemini格式
-func convertMessages(messages []ChatMessage) ([]*genai.Content, *genai.Content) {
-	var geminiContents []*genai.Content
-	var systemInstruction *genai.Content
+// convertMessages 将OpenAI格式的消息转换为Gemini格式的历史记录和最后一个用户的提示
+func convertMessages(messages []ChatMessage) (history []*genai.Content, lastPrompt []genai.Part, systemInstruction *genai.Content) {
+	if len(messages) == 0 {
+		return nil, nil, nil
+	}
 
-	for _, msg := range messages {
+	for i, msg := range messages {
 		var role string
-		if msg.Role == "user" {
-			role = "user"
-		} else if msg.Role == "assistant" {
-			role = "model"
-		} else if msg.Role == "system" {
-			// 将系统指令分开处理
+		if msg.Role == "system" {
 			systemInstruction = &genai.Content{Parts: []genai.Part{genai.Text(msg.Content)}}
-			continue // 系统指令不包含在主要内容中
+			continue
+		}
+
+		if i == len(messages)-1 && msg.Role == "user" {
+			lastPrompt = append(lastPrompt, genai.Text(msg.Content))
+			continue
+		}
+
+		if msg.Role == "assistant" {
+			role = "model"
 		} else {
-			role = "user" // 默认为用户
+			role = "user"
 		}
-		geminiContents = append(geminiContents, &genai.Content{
+
+		history = append(history, &genai.Content{
 			Role:  role,
 			Parts: []genai.Part{genai.Text(msg.Content)},
 		})
 	}
-	return geminiContents, systemInstruction
+	return history, lastPrompt, systemInstruction
 }
 
-// chatCompletionsHandler 处理聊天请求
 func chatCompletionsHandler(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPost {
 		http.Error(w, "仅支持POST方法", http.StatusMethodNotAllowed)
@@ -210,18 +207,12 @@ func chatCompletionsHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// 映射模型名称
 	modelName, ok := modelMapping[req.Model]
 	if !ok {
-		// 如果在映射中找不到，则直接使用请求的模型名称，
-		// 并选择一个默认的最新模型作为备用。
-		modelName = req.Model
-		log.Printf("警告: 模型 '%s' 不在预定义的映射中。将直接使用该名称。", req.Model)
+		modelName = "gemini-1.5-flash-latest" // 默认模型
 	}
 
-
-	// 转换消息格式
-	contents, systemInstruction := convertMessages(req.Messages)
+	history, lastPrompt, systemInstruction := convertMessages(req.Messages)
 
 	var lastErr error
 	usedKeys := make(map[string]bool)
@@ -230,7 +221,6 @@ func chatCompletionsHandler(w http.ResponseWriter, r *http.Request) {
 		ctx := context.Background()
 		apiKey := getRandomAPIKey()
 
-		// 确保在一次重试中不使用重复的密钥
 		if len(usedKeys) < len(apiKeys) {
 			for usedKeys[apiKey] {
 				apiKey = getRandomAPIKey()
@@ -256,32 +246,34 @@ func chatCompletionsHandler(w http.ResponseWriter, r *http.Request) {
 		if req.MaxTokens > 0 {
 			model.SetMaxOutputTokens(req.MaxTokens)
 		}
+		
+		chat := model.StartChat()
+		chat.History = history
 
 		if req.Stream {
-			err = handleStream(w, ctx, model, contents, req.Model)
+			err = handleStream(w, ctx, chat, lastPrompt, req.Model)
 		} else {
-			err = handleNonStream(w, ctx, model, contents, req.Model)
+			err = handleNonStream(w, ctx, model, chat, lastPrompt, req.Model)
 		}
 
 		if err == nil {
-			return // 成功处理
+			return 
 		}
 
 		lastErr = err
 		log.Printf("第 %d 次尝试失败: %v", i+1, err)
-		time.Sleep(1 * time.Second) // 等待1秒后重试
+		time.Sleep(1 * time.Second)
 	}
 
 	http.Error(w, fmt.Sprintf("所有重试均失败: %v", lastErr), http.StatusInternalServerError)
 }
 
-// handleStream 处理流式响应
-func handleStream(w http.ResponseWriter, ctx context.Context, model *genai.GenerativeModel, contents []*genai.Content, modelID string) error {
+func handleStream(w http.ResponseWriter, ctx context.Context, chat *genai.ChatSession, prompt []genai.Part, modelID string) error {
 	w.Header().Set("Content-Type", "text/event-stream")
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Connection", "keep-alive")
 
-	iter := model.GenerateContentStream(ctx, contents...)
+	iter := chat.SendMessageStream(ctx, prompt...)
 	for {
 		resp, err := iter.Next()
 		if err == iterator.Done {
@@ -297,7 +289,7 @@ func handleStream(w http.ResponseWriter, ctx context.Context, model *genai.Gener
 				contentBuilder.WriteString(string(txt))
 			}
 		}
-
+		
 		chunk := ChatCompletionStreamResponse{
 			ID:      fmt.Sprintf("chatcmpl-%d", time.Now().Unix()),
 			Object:  "chat.completion.chunk",
@@ -325,7 +317,6 @@ func handleStream(w http.ResponseWriter, ctx context.Context, model *genai.Gener
 		}
 	}
 
-	// 发送结束标志
 	finishReason := "stop"
 	doneChunk := ChatCompletionStreamResponse{
 		ID:      fmt.Sprintf("chatcmpl-%d-done", time.Now().Unix()),
@@ -350,23 +341,38 @@ func handleStream(w http.ResponseWriter, ctx context.Context, model *genai.Gener
 	return nil
 }
 
-// handleNonStream 处理非流式响应
-func handleNonStream(w http.ResponseWriter, ctx context.Context, model *genai.GenerativeModel, contents []*genai.Content, modelID string) error {
-	resp, err := model.GenerateContent(ctx, contents...)
+func handleNonStream(w http.ResponseWriter, ctx context.Context, model *genai.GenerativeModel, chat *genai.ChatSession, prompt []genai.Part, modelID string) error {
+	resp, err := chat.SendMessage(ctx, prompt...)
 	if err != nil {
 		return fmt.Errorf("生成内容失败: %v", err)
 	}
 
 	var contentBuilder strings.Builder
-	for _, part := range resp.Candidates[0].Content.Parts {
-		if txt, ok := part.(genai.Text); ok {
-			contentBuilder.WriteString(string(txt))
+	if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
+		for _, part := range resp.Candidates[0].Content.Parts {
+			if txt, ok := part.(genai.Text); ok {
+				contentBuilder.WriteString(string(txt))
+			}
 		}
 	}
+	
+	// 计算Token
+	var promptParts []genai.Part
+	for _, c := range chat.History {
+		promptParts = append(promptParts, c.Parts...)
+	}
+	promptParts = append(promptParts, prompt...)
 
-	promptTokens := int(model.CountTokens(ctx, contents...).TotalTokens)
-	completionTokens := int(model.CountTokens(ctx, resp.Candidates[0].Content).TotalTokens)
+	promptTokenCount, err := model.CountTokens(ctx, promptParts...)
+	if err != nil {
+		return fmt.Errorf("计算prompt tokens失败: %v", err)
+	}
 
+	completionTokenCount, err := model.CountTokens(ctx, resp.Candidates[0].Content.Parts...)
+	if err != nil {
+		return fmt.Errorf("计算completion tokens失败: %v", err)
+	}
+	
 	response := ChatCompletionResponse{
 		ID:      fmt.Sprintf("chatcmpl-%d", time.Now().Unix()),
 		Object:  "chat.completion",
@@ -383,9 +389,9 @@ func handleNonStream(w http.ResponseWriter, ctx context.Context, model *genai.Ge
 			},
 		},
 		Usage: Usage{
-			PromptTokens:     promptTokens,
-			CompletionTokens: completionTokens,
-			TotalTokens:      promptTokens + completionTokens,
+			PromptTokens:     int(promptTokenCount.TotalTokens),
+			CompletionTokens: int(completionTokenCount.TotalTokens),
+			TotalTokens:      int(promptTokenCount.TotalTokens) + int(completionTokenCount.TotalTokens),
 		},
 	}
 
@@ -393,7 +399,6 @@ func handleNonStream(w http.ResponseWriter, ctx context.Context, model *genai.Ge
 	return json.NewEncoder(w).Encode(response)
 }
 
-// --- 辅助端点 ---
 
 func modelsHandler(w http.ResponseWriter, r *http.Request) {
 	resp := ModelListResponse{
@@ -435,8 +440,6 @@ func healthHandler(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(health)
 }
 
-// --- Main函数 ---
-
 func main() {
 	mux := http.NewServeMux()
 
@@ -444,10 +447,8 @@ func main() {
 	mux.HandleFunc("/health", healthHandler)
 	mux.HandleFunc("/v1/models", modelsHandler)
 	mux.HandleFunc("/v1/chat/completions", chatCompletionsHandler)
-	// 添加兼容路径
 	mux.HandleFunc("/v1/chat/completions/v1/models", modelsHandler)
 
-	// 配置CORS
 	c := cors.New(cors.Options{
 		AllowedOrigins:   []string{"*"},
 		AllowedMethods:   []string{"GET", "POST", "OPTIONS"},
@@ -469,14 +470,12 @@ func main() {
 	log.Println("🔄 支持自动重试和密钥轮换")
 	log.Printf("🔗 服务器正在监听 http://0.0.0.0:%s", port)
 
-	// 从环境变量中读取密钥
 	envKey := os.Getenv("GEMINI_API_KEY")
 	if envKey != "" {
 		apiKeys = strings.Split(envKey, ",")
 		log.Printf("从环境变量 GEMINI_API_KEY 加载了 %d 个密钥", len(apiKeys))
 	}
 
-
 	if err := http.ListenAndServe(":"+port, handler); err != nil {
 		log.Fatalf("启动服务器失败: %v", err)
 	}