Jack1808/Claude_Code

1"""Request utility functions for API route handlers.2 3Contains token counting for API requests.4"""5 6import json7 8import tiktoken9from loguru import logger10 11from providers.common import get_block_attr12 13ENCODER = tiktoken.get_encoding("cl100k_base")14 15__all__ = ["get_token_count"]16 17 18def get_token_count(19    messages: list,20    system: str | list | None = None,21    tools: list | None = None,22) -> int:23    """Estimate token count for a request.24 25    Uses tiktoken cl100k_base encoding to estimate token usage.26    Includes system prompt, messages, tools, and per-message overhead.27    """28    total_tokens = 029 30    if system:31        if isinstance(system, str):32            total_tokens += len(ENCODER.encode(system))33        elif isinstance(system, list):34            for block in system:35                text = get_block_attr(block, "text", "")36                if text:37                    total_tokens += len(ENCODER.encode(str(text)))38        total_tokens += 4  # System block formatting overhead39 40    for msg in messages:41        if isinstance(msg.content, str):42            total_tokens += len(ENCODER.encode(msg.content))43        elif isinstance(msg.content, list):44            for block in msg.content:45                b_type = get_block_attr(block, "type") or None46 47                if b_type == "text":48                    text = get_block_attr(block, "text", "")49                    total_tokens += len(ENCODER.encode(str(text)))50                elif b_type == "thinking":51                    thinking = get_block_attr(block, "thinking", "")52                    total_tokens += len(ENCODER.encode(str(thinking)))53                elif b_type == "tool_use":54                    name = get_block_attr(block, "name", "")55                    inp = get_block_attr(block, "input", {})56                    block_id = get_block_attr(block, "id", "")57                    total_tokens += len(ENCODER.encode(str(name)))58                    total_tokens += len(ENCODER.encode(json.dumps(inp)))59                    total_tokens += len(ENCODER.encode(str(block_id)))60                    total_tokens += 1561                elif b_type == "image":62                    source = get_block_attr(block, "source")63                    if isinstance(source, dict):64                        data = source.get("data") or source.get("base64") or ""65                        if data:66                            total_tokens += max(85, len(data) // 3000)67                        else:68                            total_tokens += 76569                    else:70                        total_tokens += 76571                elif b_type == "tool_result":72                    content = get_block_attr(block, "content", "")73                    tool_use_id = get_block_attr(block, "tool_use_id", "")74                    if isinstance(content, str):75                        total_tokens += len(ENCODER.encode(content))76                    else:77                        total_tokens += len(ENCODER.encode(json.dumps(content)))78                    total_tokens += len(ENCODER.encode(str(tool_use_id)))79                    total_tokens += 880                else:81                    logger.debug(82                        "Unexpected block type %r, falling back to json/str encoding",83                        b_type,84                    )85                    try:86                        total_tokens += len(ENCODER.encode(json.dumps(block)))87                    except (TypeError, ValueError):88                        total_tokens += len(ENCODER.encode(str(block)))89 90    if tools:91        for tool in tools:92            tool_str = (93                tool.name + (tool.description or "") + json.dumps(tool.input_schema)94            )95            total_tokens += len(ENCODER.encode(tool_str))96 97    total_tokens += len(messages) * 498    if tools:99        total_tokens += len(tools) * 5100 101    return max(1, total_tokens)102