Commit a064c0c

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-02-20 08:36:55
feat(gpt): add helicone gateway support
1 parent dd8ce9f
Changed files (4)
src/llm/gpt.py
@@ -94,7 +94,7 @@ async def gpt_response(client: Client, message: Message, **kwargs):
         await send2tg(client, message, texts=model_type, **kwargs)
         return
     contexts = await get_conversation_contexts(client, conversations)
-    config = get_model_config_with_contexts(model_type, contexts, force_model)
+    config = get_model_config_with_contexts(model_type, contexts, force_model, info)
     msg = f"🤖{config['friendly_name']}: 思考中..."
     if kwargs.get("show_progress"):
         res = await send2tg(client, message, texts=msg, **kwargs)
src/llm/models.py
@@ -8,6 +8,7 @@ from pyrogram.types import Message
 from config import GPT, PROXY
 from llm.prompts import force_reasoning, refine_prompts
 from messages.parser import parse_msg
+from utils import unicode_to_ascii
 
 
 def get_model_type(conversations: list[Message]) -> str:
@@ -28,7 +29,7 @@ def get_model_type(conversations: list[Message]) -> str:
     return model_type
 
 
-def get_model_config_with_contexts(model_type: str, contexts: list[dict], force_model: str = "N/A") -> dict:
+def get_model_config_with_contexts(model_type: str, contexts: list[dict], force_model: str = "N/A", message_info: dict | None = None) -> dict:
     """Get GPT model config based on contexts, and return the config and adjusted contexts.
 
     contexts:
@@ -72,6 +73,8 @@ def get_model_config_with_contexts(model_type: str, contexts: list[dict], force_
         client["base_url"] = GPT.DEEPSEEK_BASE_URL
         model_name = GPT.DEEPSEEK_MODEL_NAME
 
+    client = helicone_hook(client, message_info)  # this line should be after setting `force_model``
+
     # params for `openai.chat.completions.create()`
     completions = {"model": model, "messages": contexts, "temperature": float(GPT.TEMPERATURE)}
     completions = model_hook(completions)
@@ -102,6 +105,22 @@ def openrouter_hook(base_url: str, *, for_tools: bool = False) -> dict:
     return params
 
 
+def helicone_hook(client: dict, message_info: dict | None) -> dict:
+    """Add special parameters for helicone gateway."""
+    if not GPT.HELICONE_API_KEY:
+        return client
+    headers = client.get("default_headers", {})
+    headers |= {
+        "Helicone-Auth": f"Bearer {GPT.HELICONE_API_KEY}",
+    }
+    message_info = message_info or {}
+    if chat_title := message_info.get("ctitle"):
+        headers |= {"Helicone-Property-Chat": unicode_to_ascii(chat_title), "Helicone-Property-ChatID": str(message_info["cid"])}
+    if user_name := message_info.get("full_name"):
+        headers |= {"Helicone-User-Id": unicode_to_ascii(user_name), "Helicone-Property-User": str(message_info["uid"])}
+    return client | {"default_headers": headers}
+
+
 def model_hook(params: dict) -> dict:
     """Add parameters for special models."""
     return force_reasoning(params)
src/llm/response.py
@@ -31,7 +31,7 @@ async def merge_tools_response(config: dict, **kwargs) -> dict:
     completions |= openrouter_hook(GPT.TOOLS_BASE_URL, for_tools=True)
     tools_config = {
         "friendly_name": config["friendly_name"],
-        "client": {"base_url": GPT.TOOLS_BASE_URL, "api_key": GPT.TOOLS_API_KEY, "http_client": config["client"]["http_client"]},
+        "client": config["client"] | {"base_url": GPT.TOOLS_BASE_URL, "api_key": GPT.TOOLS_API_KEY},
         "completions": add_tools(completions),
     }
     try:
src/config.py
@@ -166,6 +166,8 @@ class GPT:  # see `llm/README.md`
     TOOLS_BASE_URL = os.getenv("GPT_TOOLS_BASE_URL", "https://api.openai.com/v1")
     TOKEN_ENCODING = os.getenv("GPT_TOKEN_ENCODING", "o200k_base")  # https://github.com/openai/tiktoken
     MAX_RETRY = int(os.getenv("GPT_MAX_RETRY", "2"))
+    HELICONE_API_KEY = os.getenv("HELICONE_API_KEY", "")
+
     # comma separated reasoning models, add system prompt to the models to ensure the output format.
     REASONING_MODELS = os.getenv("GPT_REASONING_MODELS", "deepseek-r1,o1,o3")
     # /gemini command