| 1 | import re |
| 2 | from typing import Literal |
| 3 | import tiktoken |
| 4 | |
| 5 | APPROX_BUFFER = 1.1 |
| 6 | TRIM_BUFFER = 0.8 |
| 7 | EMBEDDED_IMAGE_DATA_PLACEHOLDER = "[embedded image data omitted from token estimate]" |
| 8 | _EMBEDDED_IMAGE_DATA_URL_PATTERN = re.compile( |
| 9 | r"data:(image/[A-Za-z0-9.+-]+(?:;[A-Za-z0-9.+-]+=[A-Za-z0-9.+/=_-]+)*);base64,[A-Za-z0-9+/=_-]+" |
| 10 | ) |
| 11 | |
| 12 | |
| 13 | def count_tokens(text: str, encoding_name="cl100k_base") -> int: |
| 14 | if not text: |
| 15 | return 0 |
| 16 | |
| 17 | # Get the encoding |
| 18 | encoding = tiktoken.get_encoding(encoding_name) |
| 19 | |
| 20 | # Encode the text and count the tokens |
| 21 | tokens = encoding.encode(text, disallowed_special=()) |
| 22 | token_count = len(tokens) |
| 23 | |
| 24 | return token_count |
| 25 | |
| 26 | |
| 27 | def approximate_tokens( |
| 28 | text: str, |
| 29 | ) -> int: |
| 30 | return int(count_tokens(text) * APPROX_BUFFER) |
| 31 | |
| 32 | |
| 33 | def sanitize_embedded_image_data_urls(text: str) -> str: |
| 34 | if not text: |
| 35 | return text |
| 36 | |
| 37 | return _EMBEDDED_IMAGE_DATA_URL_PATTERN.sub( |
| 38 | f"data:\\1;base64,{EMBEDDED_IMAGE_DATA_PLACEHOLDER}", |
| 39 | text, |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | def approximate_prompt_tokens(text: str) -> int: |
| 44 | return approximate_tokens(sanitize_embedded_image_data_urls(text)) |
| 45 | |
| 46 | |
| 47 | def trim_to_tokens( |
| 48 | text: str, |
| 49 | max_tokens: int, |
| 50 | direction: Literal["start", "end"], |
| 51 | ellipsis: str = "...", |
| 52 | ) -> str: |
| 53 | chars = len(text) |
| 54 | tokens = count_tokens(text) |
| 55 | |
| 56 | if tokens <= max_tokens: |
| 57 | return text |
| 58 | |
| 59 | approx_chars = int(chars * (max_tokens / tokens) * TRIM_BUFFER) |
| 60 | |
| 61 | if direction == "start": |
| 62 | return text[:approx_chars] + ellipsis |
| 63 | return ellipsis + text[chars - approx_chars : chars] |