Commit 158faa1

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-22 09:51:24
feat(gemini): add support for disabling thinking and including thoughts
1 parent c18bdaf
Changed files (1)
src
src/llm/gemini.py
@@ -44,7 +44,17 @@ HELP = f"""🌠**AI生图**
 """
 
 
-async def gemini_response(client: Client, message: Message, conversations: list[Message], modality: str = "image", *, append_grounding: bool = True, **kwargs):
+async def gemini_response(
+    client: Client,
+    message: Message,
+    conversations: list[Message],
+    modality: str = "image",
+    *,
+    append_grounding: bool = True,
+    disable_thinking: bool = False,
+    include_thoughts: bool = True,
+    **kwargs,
+):
     r"""Get Gemini response.
 
     Args:
@@ -53,6 +63,8 @@ async def gemini_response(client: Client, message: Message, conversations: list[
         conversations (list[Message]): list of chat conversations.
         modality (str): response modality
         append_grounding (bool, optional): Whether to append grounding to the response. Defaults to True.
+        disable_thinking (bool, optional): Whether to disable thinking. Defaults to False.
+        include_thoughts (bool, optional): Whether to include thoughts. Defaults to True.
     """
     info = parse_msg(message)
     model = GEMINI.TEXT_MODEL if modality == "text" else GEMINI.IMG_MODEL
@@ -78,9 +90,9 @@ async def gemini_response(client: Client, message: Message, conversations: list[
             genconfig |= {"tools": tools}
         if GEMINI.PREFER_LANG and modality == "text":
             genconfig |= {"system_instruction": f"请优先使用{GEMINI.PREFER_LANG}思考和回复"}
-        if thinking_budget is not None:
+        if thinking_budget is not None and not disable_thinking:
             thinking_budget = min(round(float(thinking_budget)), GEMINI.MAX_THINKING_BUDGET)
-            genconfig |= {"thinking_config": ThinkingConfig(include_thoughts=True, thinking_budget=thinking_budget)}
+            genconfig |= {"thinking_config": ThinkingConfig(include_thoughts=include_thoughts, thinking_budget=thinking_budget)}
         params = {"model": model, "contents": contexts, "config": GenerateContentConfig(**genconfig)}
         logger.trace(params)
         if modality == "image":