Commit 4bd1da4

@@ -13,6 +13,7 @@ from pyrogram.types import Message
 
 from config import ASR, GEMINI, TEXT_LENGTH
 from llm.gemini import parse_response
+from llm.hooks import hook_gemini_httpoptions
 from llm.utils import beautify_llm_response
 from messages.progress import modify_progress
 from messages.utils import blockquote, count_without_entities, smart_split
@@ -46,7 +47,9 @@ async def gemini_stream_asr(client: Client, message: Message, path: str | Path,
         status = None
     try:
         logger.debug(f"ASR via {ASR.GEMINI_MODEL}: {path.as_posix()} , proxy={ASR.GEMINI_PROXY}")
-        app = genai.Client(api_key=random.choice(api_keys), http_options=HttpOptions(base_url=ASR.GEMINI_BASR_URL, async_client_args={"proxy": ASR.GEMINI_PROXY}))
+        http_options = HttpOptions(base_url=ASR.GEMINI_BASR_URL, async_client_args={"proxy": ASR.GEMINI_PROXY})
+        http_options = hook_gemini_httpoptions(http_options, message)
+        app = genai.Client(api_key=random.choice(api_keys), http_options=http_options)
         uploaded_audio = await app.aio.files.upload(file=path, config=UploadFileConfig(mime_type=f"audio/{voice_format}"))
         logger.debug(uploaded_audio)
         genconfig = {}

@@ -16,6 +16,7 @@ from pyrogram.types import Message, ReplyParameters
 
 from config import CAPTION_LENGTH, DOWNLOAD_DIR, GEMINI, GPT, PREFIX, TEXT_LENGTH
 from llm.contexts import get_conversation_contexts
+from llm.hooks import hook_gemini_httpoptions
 from llm.utils import BOT_TIPS, beautify_llm_response, clean_cmd_prefix, clean_gemini_sourcemarks, clean_source_marks
 from messages.parser import parse_msg
 from messages.progress import modify_progress
@@ -97,7 +98,9 @@ async def gemini_stream(
         if retry > len(api_keys) - 1:
             return None
         api_key = kwargs.get("gemini_api_key", api_keys[retry])
-        app = genai.Client(api_key=api_key, http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
+        http_options = HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY})
+        http_options = hook_gemini_httpoptions(http_options, message)
+        app = genai.Client(api_key=api_key, http_options=http_options)
         runtime_texts = ""
         async for chunk in await app.aio.models.generate_content_stream(**params):
             resp = parse_response(chunk.model_dump())
@@ -156,7 +159,9 @@ async def gemini_nonstream(
         if retry > len(api_keys) - 1:
             return None
         api_key = kwargs.get("gemini_api_key", api_keys[retry])
-        app = genai.Client(api_key=api_key, http_options=HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY}))
+        http_options = HttpOptions(base_url=GEMINI.BASR_URL, async_client_args={"proxy": GEMINI.PROXY})
+        http_options = hook_gemini_httpoptions(http_options, message)
+        app = genai.Client(api_key=api_key, http_options=http_options)
         response = await app.aio.models.generate_content(**params)
         prefix = f"🤖**{model_name}**:{BOT_TIPS}\n"
         res = parse_response(response.model_dump())

@@ -1,7 +1,11 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+from google.genai.types import HttpOptions
+from pyrogram.types import Message
+
 from config import GEMINI, GPT
 from llm.prompts import modify_prompts, refine_prompts
+from messages.parser import parse_msg
 from utils import unicode_to_ascii
 
 
@@ -35,3 +39,15 @@ def pre_helicone_hook(client: dict, message_info: dict | None) -> None:
     if user_name := message_info.get("full_name"):
         headers |= {"Helicone-User-Id": unicode_to_ascii(user_name), "Helicone-Property-User": str(message_info["uid"])}
     client |= {"default_headers": headers}
+
+
+def hook_gemini_httpoptions(http_options: HttpOptions, message: Message) -> HttpOptions:
+    if http_options.base_url == "https://gateway.helicone.ai" and GPT.HELICONE_API_KEY:
+        info = parse_msg(message, silent=True)
+        headers = {"helicone-auth": f"Bearer {GPT.HELICONE_API_KEY}", "helicone-target-url": "https://generativelanguage.googleapis.com"}
+        if chat_title := info["ctitle"]:
+            headers |= {"Helicone-Property-Chat": unicode_to_ascii(chat_title), "Helicone-Property-ChatID": str(info["cid"])}
+        if user_name := info["full_name"]:
+            headers |= {"Helicone-User-Id": unicode_to_ascii(user_name), "Helicone-Property-User": str(info["uid"])}
+        http_options.headers = headers
+    return http_options