src/utils/token_counter.py

"""
Token 计数工具模块

使用 Google Vcore AI CountTokens API 进行精确的 token 计数
"""

import json
from typing import Any, cast
from src.utils.logger import get_logger
from src.core.config import load_config
from src.api.network import NetworkClient
from src.api.model_config import ModelConfigBuilder

logger = get_logger(__name__)

class TokenCounter:
    """Token 计数器 - 使用 Google Vcore AI CountTokens API"""
    
    def __init__(self, network: NetworkClient | None = None) -> None:
        self.config = load_config()
        self.vcore_ai_anonymous_base_api = "https://cloudconsole-pa.clients6.google.com"
        self._api_key = "AIzaSyCI-zsRP85UVOi0DjtiCwWBwQ1djDy741g"
        self.network = network or NetworkClient()
        self.model_builder = ModelConfigBuilder()
        
    async def calculate_usage_metadata_async(
        self, 
        prompt_contents: list[dict[str, Any]], 
        response_parts: list[dict[str, Any]],
        model: str = "gemini-2.5-flash"
    ) -> dict[str, Any]:
        """
        异步计算完整的 usage metadata
        """
        try:
            def clean_contents(contents: list[dict[str, Any]]) -> list[dict[str, Any]]:
                cleaned = []
                for content in contents:
                    new_content = content.copy()
                    if "parts" in new_content:
                        new_parts = []
                        for part in new_content["parts"]:
                            new_part = {}
                            if "text" in part:
                                new_part["text"] = part["text"]
                            if "inlineData" in part:
                                new_part["inlineData"] = part["inlineData"]
                            if "fileData" in part:
                                new_part["fileData"] = part["fileData"]

                            # 转换为文本
                            if "functionCall" in part:
                                func_call = part["functionCall"]
                                text_rep = f"Function Call: {func_call.get('name', 'unknown')}"
                                if "args" in func_call:
                                    try: text_rep += f" Args: {json.dumps(func_call['args'])}"
                                    except: text_rep += f" Args: {str(func_call['args'])}"
                                new_part["text"] = new_part.get("text", "") + "\n" + text_rep
                                
                            if "functionResponse" in part:
                                func_resp = part["functionResponse"]
                                text_rep = f"Function Response: {func_resp.get('name', 'unknown')}"
                                if "response" in func_resp:
                                    try: text_rep += f" Result: {json.dumps(func_resp['response'])}"
                                    except: text_rep += f" Result: {str(func_resp['response'])}"
                                new_part["text"] = new_part.get("text", "") + "\n" + text_rep

                            if new_part:
                                new_parts.append(new_part)
                        
                        if new_parts:
                            new_content["parts"] = new_parts
                        else:
                            new_content["parts"] = [{"text": " "}]
                    cleaned.append(new_content)

                # 合并连续角色
                merged = []
                for c in cleaned:
                    if not merged:
                        merged.append(c)
                    elif merged[-1].get("role") == c.get("role"):
                        merged[-1]["parts"].extend(c.get("parts", []))
                    else:
                        merged.append(c)
                
                if merged and merged[0].get("role") == "model":
                    merged.insert(0, {"role": "user", "parts": [{"text": " "}]})
                    
                return merged

            safe_prompt_contents = clean_contents(prompt_contents)
            safe_response_parts = response_parts
            
            return await self._calculate_usage_with_session(safe_prompt_contents, safe_response_parts, model, clean_contents)
                    
        except Exception as e:
            logger.error(f"计算 usage metadata 失败: {e}")
            return {"promptTokenCount": 0, "candidatesTokenCount": 0, "totalTokenCount": 0}

    async def _calculate_usage_with_session(
        self,
        safe_prompt_contents: list[dict[str, Any]],
        response_parts: list[dict[str, Any]],
        model: str,
        clean_contents_fn: Any
    ) -> dict[str, Any]:
        prompt_token_count = await self._count_tokens_with_session(safe_prompt_contents, model)
        
        full_contents = list(safe_prompt_contents)
        if response_parts:
            # Deep copy to avoid modifying original safe_prompt_contents if extended
            import copy
            full_contents = copy.deepcopy(safe_prompt_contents)
            
            # response_parts should be correctly formatted. Make sure role is set to "model"
            # It might just be raw text dicts right now.
            model_reply = clean_contents_fn([{"parts": response_parts, "role": "model"}])
            
            if model_reply:
                # model_reply might have been transformed to have user as well? No, clean_contents ensures it's clean.
                # Actually clean_contents ensures if first role is model it inserts user. We need to prevent that here for appending.
                pass
            
            # Let's write a simple append logic
            if full_contents and full_contents[-1].get("role") == "model":
                full_contents[-1]["parts"].extend(response_parts)
            else:
                full_contents.append({"role": "model", "parts": response_parts})
            
            # Clean again to ensure all rules apply (like no empty parts, function calls formatted correctly)
            full_contents = clean_contents_fn(full_contents)
        
        total_token_count = await self._count_tokens_with_session(full_contents, model)
        
        if total_token_count < prompt_token_count:
            total_token_count = prompt_token_count
            
        candidates_token_count = total_token_count - prompt_token_count
        
        usage_metadata: dict[str, Any] = {
            "promptTokenCount": prompt_token_count,
            "candidatesTokenCount": candidates_token_count,
            "totalTokenCount": total_token_count
        }
        
        logger.debug(f"Token 计算结果: {usage_metadata}")
        return usage_metadata

    async def _count_tokens_with_session(self, contents: list[dict[str, Any]], model: str) -> int:
        try:
            target_model = self.model_builder.parse_model_name(model)

            url = f"{self.vcore_ai_anonymous_base_api}/v3/entityServices/AiplatformEntityService/schemas/AIPLATFORM_GRAPHQL:batchGraphql?key={self._api_key}&prettyPrint=false"
            
            async with self.network.create_session() as session:
                recaptcha_token = await self.network.fetch_recaptcha_token(session)
                if not recaptcha_token:
                    return 0
                
                # 移除 models/ 前缀以匹配示例
                if target_model.startswith("models/"):
                    target_model = target_model[7:]

                payload = {
                    "requestContext": {
                        "clientVersion": "boq_cloud-boq-clientweb-vcoreaistudio_20260402.09_p0",
                        "pagePath": "/vcore-ai/studio/multimodal",
                        "jurisdiction": "global",
                        "localizationData": {
                            "locale": "zh_CN",
                            "timezone": "Asia/Shanghai"
                        }
                    },
                    "querySignature": "2/mENOSldfC+HZM+tGhVuJLrl8M6gEyK3HRjUKuA5AM58=",
                    "operationName": "CountTokens",
                    "variables": {
                        "contents": contents,
                        "endpoint": "",
                        "model": target_model,
                        "region": "global",
                        "recaptchaToken": recaptcha_token
                    }
                }
                
                headers = {
                    "accept": "*/*",
                    "accept-language": "zh-CN,zh;q=0.9,en;q=0.8",
                    "content-type": "application/json",
                    "origin": "https://console.cloud.google.com",
                    "referer": "https://console.cloud.google.com/vcore-ai/studio/multimodal",
                    "x-goog-authuser": "0",
                }
                
                logger.debug_json("CountTokens 请求体", payload)
                
                response = await self.network.post_request(session, url, headers, payload)
                
                if response.status_code == 200:
                    data = response.json()
                    logger.debug_json("CountTokens 响应体", data)
                    try:
                        items = data if isinstance(data, list) else [data]
                        for entry in items:
                            if not isinstance(entry, dict): continue
                            if "errors" in entry:
                                logger.error(f"CountTokens 报错: {entry['errors']}")
                                continue
                            
                            results = entry.get("results", [])
                            for result in results:
                                if "errors" in result:
                                    logger.error(f"CountTokens 报错: {result['errors']}")
                                    continue
                                data_obj = result.get("data", {})
                                ui_data = data_obj.get("ui", {})
                                count_data = ui_data.get("countTokensV2") or data_obj.get("countTokensV2") or data_obj.get("countTokens")
                                if count_data and "totalTokens" in count_data:
                                    return int(count_data["totalTokens"])
                    except Exception as e:
                        logger.error(f"解析 CountTokens 响应失败: {e}")
                else:
                    logger.error(f"CountTokens API 请求失败: {response.status_code}")
                
            return 0
        except Exception as e:
            logger.error(f"远程 Token 计数失败: {e}")
            return 0

    async def count_tokens_remote(self, contents: list[dict[str, Any]], model: str = "gemini-2.5-flash") -> int:
        return await self._count_tokens_with_session(contents, model)

# 全局实例
_token_counter = TokenCounter()

async def calculate_usage_metadata(
    prompt_contents: list[dict[str, Any]], 
    response_parts: list[dict[str, Any]],
    request_context: dict[str, Any] | None = None
) -> dict[str, Any]:
    """便捷函数：计算完整的 usage metadata"""
    model = "gemini-2.5-flash"
    if request_context and isinstance(request_context, dict):
        downstream = request_context.get("downstream_payload", {})
        if isinstance(downstream, dict):
            model = downstream.get("model", model)
        
    return await _token_counter.calculate_usage_metadata_async(
        prompt_contents, 
        response_parts,
        model
    )