pipelines/shared/token_utils.py

from typing import Dict, List

_ENCODER = None


def _get_encoder():
    global _ENCODER
    if _ENCODER is not None:
        return _ENCODER
    try:
        import tiktoken

        _ENCODER = tiktoken.get_encoding("cl100k_base")
    except Exception:
        _ENCODER = False
    return _ENCODER


def count_tokens(text: str) -> int:
    if not text:
        return 0
    encoder = _get_encoder()
    if encoder:
        return len(encoder.encode(text))
    return max(1, len(text.split()))


def select_under_budget(chunks: List[Dict], token_budget: int) -> List[Dict]:
    selected: List[Dict] = []
    used = 0

    for chunk in chunks:
        text = chunk.get("text", "")
        tokens = count_tokens(text)
        if selected and used + tokens > token_budget:
            continue
        if not selected and tokens > token_budget:
            truncated = " ".join(text.split()[: max(1, token_budget)])
            item = dict(chunk)
            item["text"] = truncated
            item["token_count"] = count_tokens(truncated)
            selected.append(item)
            break
        item = dict(chunk)
        item["token_count"] = tokens
        selected.append(item)
        used += tokens

    return selected