From 484e5cdc42db0ebec4d1544ab29e30364cb6ade4 Mon Sep 17 00:00:00 2001
From: coulsontl <149155411+coulsontl@users.noreply.github.com>
Date: Sat, 24 May 2025 09:26:20 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E7=8E=AF=E5=A2=83?=
 =?UTF-8?q?=E5=8F=98=E9=87=8F=E5=8A=A0=E8=BD=BD=E5=92=8C=E6=80=9D=E8=80=83?=
 =?UTF-8?q?=E9=85=8D=E7=BD=AE=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/domain/gemini_models.py             |  1 +
 app/main.py                             |  4 ++++
 app/service/chat/gemini_chat_service.py | 20 +++++++++++++++-----
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/app/domain/gemini_models.py b/app/domain/gemini_models.py
index d2c26e6..9528bc8 100644
--- a/app/domain/gemini_models.py
+++ b/app/domain/gemini_models.py
@@ -40,6 +40,7 @@ class GenerationConfig(BaseModel):
     frequencyPenalty: Optional[float] = None
     responseLogprobs: Optional[bool] = None
     logprobs: Optional[int] = None
+    thinkingConfig: Optional[Dict[str, Any]] = None
 
 
 class SystemInstruction(BaseModel):
diff --git a/app/main.py b/app/main.py
index 6deffa6..bbf9b50 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,4 +1,8 @@
 import uvicorn
+from dotenv import load_dotenv
+
+# 在导入应用程序配置之前加载 .env 文件到环境变量
+load_dotenv()
 
 from app.core.application import create_app
 from app.log.logger import get_main_logger
diff --git a/app/service/chat/gemini_chat_service.py b/app/service/chat/gemini_chat_service.py
index 8fa7425..cca82ae 100644
--- a/app/service/chat/gemini_chat_service.py
+++ b/app/service/chat/gemini_chat_service.py
@@ -97,11 +97,21 @@ def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]:
     if model.endswith("-image") or model.endswith("-image-generation"):
         payload.pop("systemInstruction")
         payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
-        
-    if model.endswith("-non-thinking"):
-        payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0} 
-    if model in settings.THINKING_BUDGET_MAP:
-        payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": settings.THINKING_BUDGET_MAP.get(model,1000)}
+    
+    # 处理思考配置：优先使用客户端提供的配置，否则使用默认配置
+    client_thinking_config = None
+    if request.generationConfig and request.generationConfig.thinkingConfig:
+        client_thinking_config = request.generationConfig.thinkingConfig
+    
+    if client_thinking_config is not None:
+        # 客户端提供了思考配置，直接使用
+        payload["generationConfig"]["thinkingConfig"] = client_thinking_config
+    else:
+        # 客户端没有提供思考配置，使用默认配置    
+        if model.endswith("-non-thinking"):
+            payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0} 
+        elif model in settings.THINKING_BUDGET_MAP:
+            payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": settings.THINKING_BUDGET_MAP.get(model,1000)}
 
     return payload