cf3f4ca58f
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Mert Unsal <mertunsal1905@gmail.com> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: maximevoisin-pm <maxime.voisin@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from vibe.core.utils.tokens import approx_token_count, truncate_middle_to_tokens
|
|
|
|
_MARKER = "\n\n[... truncated ...]\n\n"
|
|
|
|
|
|
def test_approx_token_count_empty() -> None:
|
|
assert approx_token_count("") == 0
|
|
|
|
|
|
def test_approx_token_count_rounds_up() -> None:
|
|
# 1..4 chars all round to 1 token; 5 rolls over to 2.
|
|
assert approx_token_count("a") == 1
|
|
assert approx_token_count("a" * 4) == 1
|
|
assert approx_token_count("a" * 5) == 2
|
|
|
|
|
|
def test_approx_token_count_budget_boundary() -> None:
|
|
# 80_000 chars == 20_000 tokens exactly; +1 char crosses the boundary.
|
|
assert approx_token_count("a" * 80_000) == 20_000
|
|
assert approx_token_count("a" * 80_001) == 20_001
|
|
|
|
|
|
def test_truncate_middle_returns_text_when_under_budget() -> None:
|
|
assert truncate_middle_to_tokens("hello", 10) == "hello"
|
|
|
|
|
|
def test_truncate_middle_empty_inputs() -> None:
|
|
assert truncate_middle_to_tokens("", 10) == ""
|
|
assert truncate_middle_to_tokens("anything", 0) == ""
|
|
assert truncate_middle_to_tokens("anything", -1) == ""
|
|
|
|
|
|
def test_truncate_middle_keeps_head_and_tail() -> None:
|
|
text = "HEAD" + ("x" * 1000) + "TAIL"
|
|
out = truncate_middle_to_tokens(text, max_tokens=20) # 80 char budget
|
|
assert _MARKER in out
|
|
assert out.startswith("HEAD")
|
|
assert out.endswith("TAIL")
|
|
assert len(out) <= 80
|
|
|
|
|
|
def test_truncate_middle_tight_budget_falls_back_to_head_cut() -> None:
|
|
# Budget too tight to fit the marker (24 chars). Function head-truncates.
|
|
text = "abcdefghijklmnopqrstuvwxyz"
|
|
out = truncate_middle_to_tokens(text, max_tokens=2) # 8 char budget
|
|
assert _MARKER not in out
|
|
assert out == "abcdefgh"
|
|
assert len(out) == 8
|