Commit 5e5094c

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-05 10:44:32
feat(gemini): support set default config for Gemini
1 parent 8c71164
Changed files (2)
src/llm/gemini.py
@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 import contextlib
+import json
 from io import BytesIO
 from pathlib import Path
 
@@ -50,14 +51,17 @@ async def gemini_response(client: Client, message: Message, conversations: list[
     response_modalities = ["TEXT", "IMAGE"] if modality == "image" else ["TEXT"]
     thinking_budget = GEMINI.IMG_THINKING_BUDGET if modality == "image" else GEMINI.TEXT_THINKING_BUDGET
     tools = [Tool(google_search=GoogleSearch())] if modality == "text" else None
-
+    # parse config from environment variable
+    genconfig = {}
+    with contextlib.suppress(Exception):
+        extra_config_str = GEMINI.IMG_CONFIG if modality == "image" else GEMINI.TEXT_CONFIG
+        genconfig = json.loads(extra_config_str)
     try:
         msg = f"πŸ€–**{model_name}**: 思考中...\nπŸ‘€**[{info['full_name']}](tg://user?id={info['uid']})**: β€œ{clean_cmd_prefix(info['text'])}”"
         status_msg = (await send2tg(client, message, texts=msg, **kwargs))[0]
         kwargs["progress"] = status_msg
         contexts = await get_conversation_contexts(client, conversations, ctx_format="gemini")
         gemini_logging(contexts)
-        genconfig = {}
         genconfig |= {"response_modalities": response_modalities}
         if tools:
             genconfig |= {"tools": tools}
src/config.py
@@ -272,8 +272,10 @@ class GEMINI:  # Official Gemini
     TEXT_MODEL = os.getenv("GEMINI_TEXT_MODEL", "gemini-2.5-pro-exp-03-25")
     TEXT_MODEL_NAME = os.getenv("GEMINI_TEXT_MODEL_NAME", "Gemini-2.5-Pro")
     TEXT_THINKING_BUDGET = os.getenv("GEMINI_TEXT_THINKING_BUDGET", None)  # 0 to disable thinking. DO NOT set this if the model is not a thinking model
+    TEXT_CONFIG = os.getenv("GEMINI_TEXT_CONFIG", "{}")  # default config passed to GenerateContentConfig. Should be a json string: '{"key": "value"}'
 
     # response modality: image
     IMG_MODEL = os.getenv("GEMINI_IMG_MODEL", "gemini-2.0-flash-exp")
     IMG_MODEL_NAME = os.getenv("GEMINI_IMG_MODEL_NAME", "Gemini-2.0-Flash")
     IMG_THINKING_BUDGET = os.getenv("GEMINI_IMG_THINKING_BUDGET", None)  # 0 to disable thinking. DO NOT set this if the model is not a thinking model
+    IMG_CONFIG = os.getenv("GEMINI_IMG_CONFIG", "{}")  # default config passed to GenerateContentConfig. Should be a json string: '{"key": "value"}'