Commit 4116ced

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-05-05 12:57:50
fix(gemini): collapse long text in single message
1 parent 7808832
Changed files (1)
src
src/llm/gemini.py
@@ -14,7 +14,7 @@ from PIL import Image
 from pyrogram.client import Client
 from pyrogram.types import Message, ReplyParameters
 
-from config import DOWNLOAD_DIR, GEMINI, GPT, PREFIX, TEXT_LENGTH
+from config import CAPTION_LENGTH, DOWNLOAD_DIR, GEMINI, GPT, PREFIX, TEXT_LENGTH
 from llm.contexts import get_conversation_contexts
 from llm.utils import BOT_TIPS, beautify_llm_response, clean_cmd_prefix, clean_gemini_sourcemarks, clean_source_marks
 from messages.parser import parse_msg
@@ -145,10 +145,6 @@ async def gemini_nonstream(
     retry: int = 0,
     **kwargs,
 ):
-    """Currently, the non-stream mode is for image generation.
-
-    For other uses, please use stream mode.
-    """
     try:
         clean_gemini_sourcemarks(params["contents"])
         api_keys = [x.strip() for x in GEMINI.API_KEYS.split(",") if x.strip()]
@@ -157,8 +153,14 @@ async def gemini_nonstream(
         app = genai.Client(api_key=api_keys[retry], http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
         response = await app.aio.models.generate_content(**params)
         prefix = f"🤖**{model_name}**:{BOT_TIPS}\n"
-        res = parse_response(response.model_dump(), prefix=prefix)
-        await send2tg(client, message, caption_above=True, **res, **kwargs)
+        res = parse_response(response.model_dump())
+        texts = res.get("texts", "")
+        media = res.get("media", [])
+        length = await count_without_entities(prefix + texts)
+        single_msg_length = CAPTION_LENGTH if media else TEXT_LENGTH
+        if length <= single_msg_length:  # short answer in single msg
+            texts = f"{prefix}{blockquote(texts)}" if length > GPT.COLLAPSE_LENGTH else f"{prefix}{texts}"
+        await send2tg(client, message, caption_above=True, texts=texts, media=media, **kwargs)
         await modify_progress(del_status=True, **kwargs)
     except Exception as e:
         logger.error(e)
@@ -171,7 +173,7 @@ async def gemini_nonstream(
         return await gemini_nonstream(client, message, model_name, params, retry + 1, **kwargs)  # type: ignore
 
 
-def parse_response(data: dict, prefix: str = "") -> dict:
+def parse_response(data: dict) -> dict:
     """Parse gemini response, includes texts, image and websearch."""
     logger.trace(data)
     parts = glom(data, "candidates.0.content.parts", default=[]) or []
@@ -195,7 +197,7 @@ def parse_response(data: dict, prefix: str = "") -> dict:
         title = glom(grounding, "web.title", default="Web")
         url = glom(grounding, "web.uri", default="https://www.google.com")
         texts += f"\n{number_to_emoji(idx + 1)}[{title}]({url})"
-    return {"texts": prefix + beautify_llm_response(texts, newline_level=2), "media": media}
+    return {"texts": beautify_llm_response(texts, newline_level=2), "media": media}
 
 
 def gemini_logging(contexts: list):