Commit 5653ad3

benny-dou <60535774+benny-dou@users.noreply.github.com>
2026-05-22 06:50:18
chore(ai): add `add_sender` param to control whether add sender info in contexts
1 parent 86bb25f
src/ai/texts/claude.py
@@ -38,6 +38,7 @@ async def anthropic_responses(
     anthropic_append_citation: bool = True,
     skills: str = "",
     hide_thinking: bool = False,
+    add_sender: bool | None = None,
     silent: bool = False,
     max_retries: int = 3,
     **kwargs,
@@ -82,6 +83,7 @@ async def anthropic_responses(
                     anthropic=anthropic,
                     cache_hour=cache_hour,
                     media_send_as=anthropic_media_send_as,
+                    add_sender=add_sender,
                 ),
             }
             if literal_eval(anthropic_responses_config):
src/ai/texts/contexts.py
@@ -19,7 +19,7 @@ from pyrogram.types import Message
 
 from ai.utils import BOT_TIPS, clean_context
 from asr.utils import GEMINI_AUDIO_EXT, downsampe_audio
-from config import AI, DOWNLOAD_DIR
+from config import AI, DOWNLOAD_DIR, TID
 from database.r2 import head_cf_r2, set_cf_r2
 from messages.parser import get_thread_id, parse_msg
 from utils import convert_md, read_text
@@ -48,17 +48,19 @@ async def base64_media(client: Client, message: Message) -> dict:
     }
 
 
-async def get_openai_completion_contexts(client: Client, message: Message) -> list[dict]:
+async def get_openai_completion_contexts(client: Client, message: Message, *, add_sender: bool | None = None) -> list[dict]:
     """Generate OpenAI chat completion contexts."""
     messages = [message]
     while message.reply_to_message:
         message = message.reply_to_message
         messages.append(message)
+    if add_sender is None:
+        add_sender = is_multi_user_chat(messages)
     messages = messages[: int(AI.MAX_CONTEXTS_NUM)][::-1]  # old to new
-    return [ctx for msg in messages if (ctx := await single_openai_chat_context(client, msg))]
+    return [ctx for msg in messages if (ctx := await single_openai_chat_context(client, msg, add_sender=add_sender))]
 
 
-async def single_openai_chat_context(client: Client, message: Message) -> dict:
+async def single_openai_chat_context(client: Client, message: Message, *, add_sender: bool) -> dict:
     """Generate OpenAI chat completion contexts for a single message.
 
     Returns:
@@ -107,11 +109,12 @@ async def single_openai_chat_context(client: Client, message: Message) -> dict:
                         }
                     )
             # user message has entity urls, use full html
-            clean_texts = clean_context(info["html"] or info["text"]) if role == "user" and info["entity_urls"] else clean_context(info["text"])
+            texts = info["html"] or info["text"] if role == "user" and info["entity_urls"] else info["text"]
+            clean_texts = clean_context(texts)
             if not clean_texts:
                 continue
-            if role == "user" and sender:  # noqa: SIM108
-                texts = f"<quote>{info['quote_text']}</quote>\n[username]: {sender}\n[message]:\n{clean_texts}"
+            if role == "user" and add_sender and sender:
+                texts = f"<quote>{info['quote_text']}</quote>\n{sender} ({info['time']})\n{clean_texts}"
             else:
                 texts = f"<quote>{info['quote_text']}</quote>\n{clean_texts}"
             texts = texts.removeprefix("<quote></quote>\n")  # remove quote mark if no quote_text
@@ -122,7 +125,7 @@ async def single_openai_chat_context(client: Client, message: Message) -> dict:
     return {"role": role, "content": contexts} if contexts else {}
 
 
-async def get_openai_response_contexts(client: Client, message: Message, openai_params: dict) -> tuple[str, list[dict]]:
+async def get_openai_response_contexts(client: Client, message: Message, params: dict) -> tuple[str, list[dict]]:
     """Generate OpenAI response contexts.
 
     Returns:
@@ -135,11 +138,11 @@ async def get_openai_response_contexts(client: Client, message: Message, openai_
         Returns:
             previous_response_id: str
         """
-        cache_day = openai_params["cache_day"]
+        cache_day = params["cache_day"]
         if cache_day == 0:
             return ""
-        api_key = openai_params["api_key"]
-        model_id = openai_params["model_id"]
+        api_key = params["api_key"]
+        model_id = params["model_id"]
         key_hash = hashlib.sha256(api_key.encode()).hexdigest()
         tid = get_thread_id(msg)
         resp = await head_cf_r2(f"TTL/{cache_day}d/OpenAI/{msg.chat.id}/{msg.id}{'/' + str(tid) if tid else ''}/{model_id}/{key_hash}")
@@ -154,10 +157,12 @@ async def get_openai_response_contexts(client: Client, message: Message, openai_
             break
         messages.append(message)
     messages.reverse()  # old to new
-    return previous_response_id, [ctx for msg in messages if (ctx := await single_openai_response_context(client, msg, openai_params))]
+    if params.get("add_sender") is None:
+        params["add_sender"] = is_multi_user_chat(messages)
+    return previous_response_id, [ctx for msg in messages if (ctx := await single_openai_response_context(client, msg, params))]
 
 
-async def single_openai_response_context(client: Client, message: Message, openai_params: dict) -> dict:
+async def single_openai_response_context(client: Client, message: Message, params: dict) -> dict:
     """Generate OpenAI response contexts for a single message.
 
     Returns:
@@ -177,11 +182,11 @@ async def single_openai_response_context(client: Client, message: Message, opena
     extra_markdown_extensions = [".html", ".docx", ".pptx", ".xls", ".xlsx"]  # convert to markdown
 
     messages = await client.get_media_group(message.chat.id, message.id) if message.media_group_id else [message]
-    media_send_as = openai_params.get("openai_media_send_as", "base64")
-    allow_image = bool(openai_params.get("allow_image"))
-    allow_video = bool(openai_params.get("allow_video"))
-    allow_audio = bool(openai_params.get("allow_audio"))
-    allow_file = bool(openai_params.get("allow_file"))
+    media_send_as = params.get("openai_media_send_as", "base64")
+    allow_image = bool(params.get("allow_image"))
+    allow_video = bool(params.get("allow_video"))
+    allow_audio = bool(params.get("allow_audio"))
+    allow_file = bool(params.get("allow_file"))
     contexts = []
     for msg in messages:
         info = parse_msg(msg, silent=True)
@@ -190,19 +195,19 @@ async def single_openai_response_context(client: Client, message: Message, opena
         file_id = ""
         try:
             if info["mtype"] == "photo" and allow_image:
-                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, openai_params)):
+                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, params)):
                     contexts.append({"type": "input_image", "file_id": file_id})
                 if not file_id:
                     res = await base64_media(client, msg)
                     contexts.append({"type": "input_image", "image_url": f"data:image/{res['ext']};base64,{res['base64']}"})
             elif info["mtype"] == "video" and allow_video:
-                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, openai_params)):
+                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, params)):
                     contexts.append({"type": "input_video", "file_id": file_id})
                 if not file_id:
                     res = await base64_media(client, msg)
                     contexts.append({"type": "input_video", "video_url": f"data:video/{res['ext']};base64,{res['base64']}"})
             elif info["mtype"] == "audio" and allow_audio:
-                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, openai_params)):
+                if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, params)):
                     contexts.append({"type": "input_audio", "file_id": file_id})
                 if not file_id:
                     res = await base64_media(client, msg)
@@ -210,7 +215,7 @@ async def single_openai_response_context(client: Client, message: Message, opena
             elif info["mtype"] == "document" and allow_file:
                 guessed_mime, _ = mimetypes.guess_type(info["file_name"])
                 if info["mime_type"] == "application/pdf" or guessed_mime == "application/pdf":
-                    if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, openai_params)):
+                    if media_send_as == "file_id" and (file_id := await get_openai_file_id(client, msg, params)):
                         contexts.append({"type": "input_file", "file_id": file_id})
                     if not file_id:
                         res = await base64_media(client, msg)
@@ -235,11 +240,12 @@ async def single_openai_response_context(client: Client, message: Message, opena
                         }
                     )
             # user message has entity urls, use full html
-            clean_texts = clean_context(info["html"] or info["text"]) if role == "user" and info["entity_urls"] else clean_context(info["text"])
+            texts = info["html"] or info["text"] if role == "user" and info["entity_urls"] else info["text"]
+            clean_texts = clean_context(texts)
             if not clean_texts:
                 continue
-            if role == "user" and sender:  # noqa: SIM108
-                texts = f"<quote>{info['quote_text']}</quote>\n[username]: {sender}\n[message]:\n{clean_texts}"
+            if role == "user" and params.get("add_sender") and sender:
+                texts = f"<quote>{info['quote_text']}</quote>\n{sender} ({info['time']})\n{clean_texts}"
             else:
                 texts = f"<quote>{info['quote_text']}</quote>\n{clean_texts}"
             texts = texts.removeprefix("<quote></quote>\n")  # remove quote mark if no quote_text
@@ -253,10 +259,10 @@ async def single_openai_response_context(client: Client, message: Message, opena
     return {"role": role, "type": "message", "content": contexts, **extra}
 
 
-async def get_openai_file_id(client: Client, message: Message, openai_params: dict) -> str:
+async def get_openai_file_id(client: Client, message: Message, params: dict) -> str:
     def get_real_baseurl() -> str:
-        base_url = str(openai_params["base_url"]) or ""
-        default_headers = openai_params.get("default_headers", {})
+        base_url = str(params["base_url"]) or ""
+        default_headers = params.get("default_headers", {})
         default_headers = {k.lower(): v for k, v in default_headers.items()}
         if base_url.startswith("https://gateway.helicone.ai"):
             helicone_target_url = default_headers.get("helicone-target-url") or ""
@@ -265,12 +271,12 @@ async def get_openai_file_id(client: Client, message: Message, openai_params: di
             return default_headers.get("x-portkey-custom-host") or ""
         return base_url
 
-    if openai_params.get("max_upload_size") and message_bytes(message) > int(openai_params["max_upload_size"]):
-        logger.warning(f"Message-{message.id} size {message_bytes(message)} bytes exceeds max_upload_size {openai_params['max_upload_size']}")
+    if params.get("max_upload_size") and message_bytes(message) > int(params["max_upload_size"]):
+        logger.warning(f"Message-{message.id} size {message_bytes(message)} bytes exceeds max_upload_size {params['max_upload_size']}")
         return ""
-    cache_day = openai_params.get("cache_day", 30)
-    api_key = openai_params["api_key"]
-    model_id = openai_params["model_id"]
+    cache_day = params.get("cache_day", 30)
+    api_key = params["api_key"]
+    model_id = params["model_id"]
     key_hash = hashlib.sha256(api_key.encode()).hexdigest()
     tid = get_thread_id(message)
     r2_key = f"TTL/{cache_day}d/OpenAI/{message.chat.id}/{message.id}{'/' + str(tid) if tid else ''}/{model_id}/{key_hash}-file_id"
@@ -281,7 +287,7 @@ async def get_openai_file_id(client: Client, message: Message, openai_params: di
     openai = AsyncOpenAI(
         base_url=get_real_baseurl(),
         api_key=api_key,
-        http_client=DefaultAsyncHttpxClient(proxy=openai_params["proxy"]) if openai_params.get("proxy") else None,
+        http_client=DefaultAsyncHttpxClient(proxy=params["proxy"]) if params.get("proxy") else None,
     )
     fpath: str | Path = await client.download_media(message)  # ty:ignore[invalid-assignment]
     try:
@@ -303,7 +309,7 @@ async def get_openai_file_id(client: Client, message: Message, openai_params: di
     return ""
 
 
-async def get_gemini_contexts(client: Client, message: Message, gemini: genai.Client) -> list[dict]:
+async def get_gemini_contexts(client: Client, message: Message, gemini: genai.Client, *, add_sender: bool | None = None) -> list[dict]:
     """Generate Gemini contexts from old to new.
 
     Returns:
@@ -313,6 +319,8 @@ async def get_gemini_contexts(client: Client, message: Message, gemini: genai.Cl
     while message.reply_to_message:
         message = message.reply_to_message
         ctx_messages.append(message)
+    if add_sender is None:
+        add_sender = is_multi_user_chat(ctx_messages)
     ctx_messages = ctx_messages[: int(AI.MAX_CONTEXTS_NUM)][::-1]  # old to new
     contexts = []
     for m in ctx_messages:
@@ -338,7 +346,7 @@ async def get_gemini_contexts(client: Client, message: Message, gemini: genai.Cl
                     if info["mtype"] in ["audio", "voice"] and Path(fpath).suffix not in GEMINI_AUDIO_EXT:
                         audio_path = await downsampe_audio(fpath)
                         fpath = audio_path.as_posix()
-                    upload = await gemini.aio.files.upload(file=fpath, config=UploadFileConfig(display_name=info["file_name"] or f"send from {sender}"))
+                    upload = await gemini.aio.files.upload(file=fpath, config=UploadFileConfig(display_name=info["file_name"] or f"send_from_{sender}"))
                     while upload.state == FileState.PROCESSING:
                         logger.trace("Waiting for upload to complete...")
                         await asyncio.sleep(1)
@@ -356,11 +364,12 @@ async def get_gemini_contexts(client: Client, message: Message, gemini: genai.Cl
                         text = convert_md(fpath)
                         Path(fpath).unlink(missing_ok=True)
                         parts.append(Part.from_text(text=f"[filename]: {info['file_name']}\n[file content]:\n{text.strip()}"))
-                clean_texts = clean_context(info["html"] or info["text"])
+                texts = info["html"] or info["text"] if role == "user" and info["entity_urls"] else info["text"]
+                clean_texts = clean_context(texts)
                 if not clean_texts:
                     continue
-                if role == "user" and sender:  # noqa: SIM108
-                    texts = f"<quote>{info['quote_text']}</quote>\n[username]: {sender}\n[message]:\n{clean_texts}"
+                if role == "user" and add_sender and sender:
+                    texts = f"<quote>{info['quote_text']}</quote>\n{sender} ({info['time']})\n{clean_texts}"
                 else:
                     texts = f"<quote>{info['quote_text']}</quote>\n{clean_texts}"
                 texts = texts.removeprefix("<quote></quote>\n")  # remove quote mark if no quote_text
@@ -379,6 +388,8 @@ async def get_anthropic_contexts(client: Client, message: Message, **kwargs) ->
     while message.reply_to_message:
         message = message.reply_to_message
         messages.append(message)
+    if kwargs.get("add_sender") is None:
+        kwargs["add_sender"] = is_multi_user_chat(messages)
     messages = messages[: int(AI.MAX_CONTEXTS_NUM)][::-1]  # old to new
     return [ctx for msg in messages if (ctx := await single_anthropic_context(client, msg, **kwargs))]
 
@@ -389,6 +400,8 @@ async def single_anthropic_context(
     anthropic: AsyncAnthropic,
     cache_hour: int = 0,
     media_send_as: Literal["base64", "file_id"] = "file_id",
+    *,
+    add_sender: bool = True,
 ) -> dict:
     """Generate Anthropic contexts for a single message.
 
@@ -441,11 +454,12 @@ async def single_anthropic_context(
                     Path(fpath).unlink(missing_ok=True)
                     contexts.append({"type": "text", "text": f"[filename]: {info['file_name']}\n[file content]:\n{text.strip()}"})
             # user message has entity urls, use full html
-            clean_texts = clean_context(info["html"] or info["text"]) if role == "user" and info["entity_urls"] else clean_context(info["text"])
+            texts = info["html"] or info["text"] if role == "user" and info["entity_urls"] else info["text"]
+            clean_texts = clean_context(texts)
             if not clean_texts:
                 continue
-            if role == "user" and sender:  # noqa: SIM108
-                texts = f"<quote>{info['quote_text']}</quote>\n[username]: {sender}\n[message]:\n{clean_texts}"
+            if role == "user" and add_sender and sender:
+                texts = f"<quote>{info['quote_text']}</quote>\n{sender} ({info['time']})\n{clean_texts}"
             else:
                 texts = f"<quote>{info['quote_text']}</quote>\n{clean_texts}"
             texts = texts.removeprefix("<quote></quote>\n")  # remove quote mark if no quote_text
@@ -490,3 +504,10 @@ async def context_bytes(client: Client, message: Message) -> int:
 
 def message_bytes(message: Message) -> int:
     return glom(message, Coalesce("photo.sizes.-1.file_size", "video.file_size", "document.file_size"), default=0)
+
+
+def is_multi_user_chat(messages: list[Message]) -> bool:
+    """Check if this chat history group has multiple users."""
+    uids = {glom(x, "from_user.id", default=0) for x in messages}
+    uids.discard(TID.ME)
+    return len(uids) > 1
src/ai/texts/gemini.py
@@ -38,6 +38,7 @@ async def gemini_chat_completion(
     gemini_append_grounding: bool = True,
     skills: str = "",
     hide_thinking: bool = False,
+    add_sender: bool | None = None,
     silent: bool = False,
     max_retries: int = 3,
     **kwargs,
@@ -60,7 +61,7 @@ async def gemini_chat_completion(
         try:
             http_options = types.HttpOptions(base_url=gemini_base_url, async_client_args={"proxy": gemini_proxy}, headers=literal_eval(gemini_default_headers))
             gemini = genai.Client(api_key=api_key, http_options=http_options)
-            params: dict = {"model": model_id, "contents": await get_gemini_contexts(client, message, gemini)}
+            params: dict = {"model": model_id, "contents": await get_gemini_contexts(client, message, gemini, add_sender=add_sender)}
             if skills:
                 gemini_generate_content_config = literal_eval(gemini_generate_content_config) | {"system_instruction": await load_skills(skills)}
             if conf := literal_eval(gemini_generate_content_config):
src/ai/texts/openai_chat.py
@@ -35,6 +35,7 @@ async def openai_chat_completions(
     openai_tools: list[dict] | None = None,
     skills: str = "",
     hide_thinking: bool = False,
+    add_sender: bool | None = None,
     silent: bool = False,
     max_retries: int = 3,
     **kwargs,
@@ -61,7 +62,7 @@ async def openai_chat_completions(
             openai_client |= {"default_headers": literal_eval(openai_default_headers)}
         if openai_proxy:
             openai_client |= {"http_client": DefaultAsyncHttpxClient(proxy=openai_proxy)}
-        contexts = openai_contexts or await get_openai_completion_contexts(client, message)
+        contexts = openai_contexts or await get_openai_completion_contexts(client, message, add_sender=add_sender)
         if openai_system_prompt and glom(contexts, "0.role", default="") != "system":
             contexts.insert(0, {"role": "system", "content": openai_system_prompt})
         if skills:
src/ai/texts/openai_response.py
@@ -46,6 +46,7 @@ async def openai_responses_api(
     skills: str = "",
     openai_append_tool_results: bool = True,
     hide_thinking: bool = False,
+    add_sender: bool | None = None,
     silent: bool = False,
     max_retries: int = 3,
     **kwargs,
@@ -85,7 +86,7 @@ async def openai_responses_api(
             previous_response_id, contexts = await get_openai_response_contexts(
                 client,
                 message,
-                openai_params=openai_client
+                params=openai_client
                 | {
                     "proxy": openai_proxy,
                     "model_id": model_id,
@@ -95,6 +96,7 @@ async def openai_responses_api(
                     "allow_audio": openai_allow_audio,
                     "allow_file": openai_allow_file,
                     "openai_media_send_as": openai_media_send_as,
+                    "add_sender": add_sender,
                 },
             )
             params = {}
src/config.py
@@ -184,6 +184,7 @@ class COOKIE:  # See: https://github.com/easychen/CookieCloud
 
 
 class TID:  # see more TID usecase in `src/permission.py`
+    ME = int(os.getenv("TID_ME", "0"))  # a temperary chat for some tasks
     ADMIN = os.getenv("TID_ADMIN", "")  # comma separated userid or @username
     TEMP = os.getenv("TID_TEMP", "me")  # a temperary chat for some tasks
     HISTORY_ADMIN = os.getenv("TID_HISTORY_ADMIN", "")  # comma separated userid (@username is NOT supported!)