Commit 158faa1
Changed files (1)
src
llm
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":