Commit 4116ced
Changed files (1)
src
llm
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):