tokenizer.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/tokenizer.py
 
"""
Tokenizer-based token counting.
 
Purpose
-------
Provides accurate token counting using tiktoken when available,
with graceful fallback to character-based heuristics.
 
Responsibilities
----------------
- Count tokens in message lists and strings
- Lazy-load tiktoken encoding (cached per process)
- Provide consistent token counting across the codebase
 
Diagnostics
-----------
Domain: OC_CHAT.CTX
Levels:
    L2 — tokenizer fallback notice (tiktoken missing)
    L3 — tokenizer initialization, encoding selection
    L4 — per-call token counts
 
Contracts
---------
- Returns integer token count
- Works without tiktoken installed (falls back to heuristic)
- Uses cl100k_base encoding (GPT-4/DeepSeek compatible)
 
Dependencies
------------
- tiktoken (lazy; the cl100k_base encoding, optional — graceful fallback if not installed)
- oc_diagnostics (structured debug logging via _dbg)
"""
from __future__ import annotations
 
from typing import TYPE_CHECKING, Any
 
from oc_diagnostics import _dbg
 
if TYPE_CHECKING:
    pass
 
# Lazy-loaded tokenizer encoding
_encoding = None
 
 
def _get_encoding():
    """Lazy-load tiktoken encoding. Returns None if tiktoken not available.
 
    Uses cl100k_base encoding which is compatible with GPT-4 and DeepSeek
    tokenization schemes.
    """
    func = "_get_encoding"
    global _encoding
    if _encoding is not None:
        return _encoding
 
    try:
        import tiktoken
        _encoding = tiktoken.get_encoding("cl100k_base")
        _dbg("OC_CHAT.CTX", func, "tiktoken cl100k_base loaded", 3)
        return _encoding
    except ImportError:
        _dbg("OC_CHAT.CTX", func,
             "tiktoken not available, using fallback heuristic", 2)
        return None
 
 
def count_tokens(text: str) -> int:
    """Count tokens in a string using tiktoken or fallback heuristic.
    
    Uses cl100k_base encoding (GPT-4/DeepSeek compatible).
    Falls back to ~4 characters per token if tiktoken not installed.
    
    Args:
        text: The string to count tokens for.
        
    Returns:
        Estimated token count.
    """
    if not text:
        return 0
    
    enc = _get_encoding()
    if enc is not None:
        return len(enc.encode(text))
    
    # Fallback heuristic: ~4 characters per token
    return len(text) // 4
 
 
def count_message_tokens(messages: list[dict[str, Any]]) -> int:
    """Count tokens in a message list.
    
    This is the replacement for the old _estimate_tokens() heuristic.
    Counts content, reasoning_content, and tool_calls for each message.
    
    Args:
        messages: List of message dicts with 'role', 'content', etc.
        
    Returns:
        Total estimated token count for the message list.
    """
    if not messages:
        return 0
    
    total = 0
    
    for msg in messages:
        # Message structure overhead (~4 tokens per message for formatting)
        total += 4
        
        # Content
        content = msg.get("content", "")
        if isinstance(content, str):
            total += count_tokens(content)
        
        # Role token
        role = msg.get("role", "")
        if role:
            total += 1
        
        # Reasoning content (DeepSeek thinking mode)
        reasoning = msg.get("reasoning_content", "")
        if isinstance(reasoning, str):
            total += count_tokens(reasoning)
        
        # Tool calls (if present in assistant messages)
        tool_calls = msg.get("tool_calls", [])
        if tool_calls and isinstance(tool_calls, list):
            for tc in tool_calls:
                fn = tc.get("function", {})
                if isinstance(fn, dict):
                    total += count_tokens(fn.get("name", ""))
                    total += count_tokens(fn.get("arguments", ""))
                # Tool call ID and type overhead
                total += 4
    
    return total