mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test(rust_bridge): drop route dispatch assertions, test the bridge directly (#42536)
Co-authored-by: ryan <ryan@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
569ccaece9
commit
97a6c27bee
4 changed files with 68 additions and 677 deletions
|
|
@ -1,10 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -24,16 +22,12 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
)
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
_get_team_member_budget_counter,
|
||||
count_request_input_tokens,
|
||||
estimate_request_max_cost,
|
||||
release_unbound_budget_reservation,
|
||||
reserve_budget_for_request,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm.rust_bridge import bindings, configuration
|
||||
from litellm.rust_bridge import token_counter as rust_token_counter
|
||||
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
TOKEN_COUNTING_ROUTES: Final = (
|
||||
|
|
@ -222,249 +216,6 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None:
|
|||
assert math.isclose(after, before * 1000)
|
||||
|
||||
|
||||
ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
CL100K_MODEL: Final = "gpt-4"
|
||||
O200K_MODEL: Final = "gpt-4o"
|
||||
RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
|
||||
RUST_INPUT_TOKENS: Final = 4_321
|
||||
RUST_INPUT_TOKENS_BY_TOKENIZER: Final = MappingProxyType(
|
||||
{"anthropic": RUST_INPUT_TOKENS, "cl100k_base": 1_234, "o200k_base": 2_345}
|
||||
)
|
||||
|
||||
|
||||
class _FakeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeUpstream(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
"""Stands in for one shared native `Tokenizer`; only its name identifies it."""
|
||||
|
||||
def __init__(self, name: str, json: str | None = None) -> None:
|
||||
self.name = name
|
||||
self.json = json
|
||||
|
||||
|
||||
def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None:
|
||||
"""Point the counter's tokenizer lookups at fakes; the codec path keeps falling back to Python."""
|
||||
fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")}
|
||||
anthropic: Final = _FakeTokenizer("anthropic", anthropic_json)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_encoding", fakes.__getitem__)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic)
|
||||
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _FakeDeclined
|
||||
RustUpstreamError = _FakeUpstream
|
||||
|
||||
|
||||
class _RecordingCounter:
|
||||
"""Stands in for one native counter; records `(tokenizer, body)` on the shared factory."""
|
||||
|
||||
def __init__(self, factory: _RecordingFactory, tokenizer: rust_token_counter.RustTokenizer) -> None:
|
||||
self.factory = factory
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
self.factory.calls.append((self.tokenizer, body))
|
||||
return {"model": "", "input_tokens": RUST_INPUT_TOKENS_BY_TOKENIZER[self.tokenizer]}
|
||||
|
||||
|
||||
class _RecordingFactory:
|
||||
"""Stands in for the native `TokenCounter` class, built over a loaded `Tokenizer`."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
|
||||
|
||||
def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, cast(rust_token_counter.RustTokenizer, tokenizer.name))
|
||||
|
||||
|
||||
class _DecliningCounter:
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise _FakeDeclined("unsupported content block")
|
||||
|
||||
|
||||
class _DecliningFactory:
|
||||
def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _DecliningCounter:
|
||||
return _DecliningCounter()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rust_counter(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
|
||||
_fake_native_tokenizers(monkeypatch)
|
||||
rust_token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
yield
|
||||
rust_token_counter.TOKEN_COUNTER.reset()
|
||||
rust_token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("route", "request_body"),
|
||||
(
|
||||
("/v1/messages", RUST_COUNTED_BODY),
|
||||
("/v1/chat/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "messages": ANTHROPIC_MESSAGES}),
|
||||
("/v1/completions", {"model": ANTHROPIC_TOKENIZER_MODEL, "prompt": "hi"}),
|
||||
("/v1/responses", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": "hi"}),
|
||||
("/v1/embeddings", {"model": ANTHROPIC_TOKENIZER_MODEL, "input": ["hi"]}),
|
||||
("/v1/rerank", {"model": ANTHROPIC_TOKENIZER_MODEL, "query": "hi", "documents": ["a"]}),
|
||||
),
|
||||
)
|
||||
async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
|
||||
rust_counter: None, route: str, request_body: dict
|
||||
) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
raw_body: Final = json.dumps(request_body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=request_body, route=route, llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS}
|
||||
assert factory.calls == [("anthropic", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", (CL100K_MODEL, "azure/gpt-35-turbo", "gemini/gemini-2.5-pro", "my-router-alias"))
|
||||
async def test_tiktoken_cl100k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"]}
|
||||
assert factory.calls == [("cl100k_base", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", (O200K_MODEL, "gpt-5", "o3", "gpt-4.1", "chatgpt-4o-latest"))
|
||||
async def test_tiktoken_o200k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"]}
|
||||
assert factory.calls == [("o200k_base", raw_body)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_model_request_counts_once_per_tokenizer_and_python_for_the_rest(rust_counter: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
models: Final = (
|
||||
CL100K_MODEL,
|
||||
ANTHROPIC_TOKENIZER_MODEL,
|
||||
"gemini/gemini-2.5-pro",
|
||||
O200K_MODEL,
|
||||
"gpt-5",
|
||||
"replicate/meta/llama-2-70b-chat",
|
||||
)
|
||||
body: Final = {"model": list(models), "messages": ANTHROPIC_MESSAGES}
|
||||
raw_body: Final = json.dumps(body).encode()
|
||||
python_counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None
|
||||
)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
|
||||
)
|
||||
|
||||
assert factory.calls == [("cl100k_base", raw_body), ("anthropic", raw_body), ("o200k_base", raw_body)]
|
||||
assert dict(counts) == {
|
||||
CL100K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
|
||||
"gemini/gemini-2.5-pro": RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
|
||||
ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS,
|
||||
O200K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
|
||||
"gpt-5": RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
|
||||
"replicate/meta/llama-2-70b-chat": python_counts["replicate/meta/llama-2-70b-chat"],
|
||||
}
|
||||
assert counts["replicate/meta/llama-2-70b-chat"] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", (ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL))
|
||||
async def test_rust_decline_falls_back_to_python_count(rust_counter: None, model: str) -> None:
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(_DecliningFactory())
|
||||
body: Final = {**RUST_COUNTED_BODY, "model": model}
|
||||
python_counts: Final = await count_request_input_tokens(request_body=body, route="/v1/messages", llm_router=None)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body,
|
||||
route="/v1/messages",
|
||||
llm_router=None,
|
||||
raw_body=json.dumps(body).encode(),
|
||||
)
|
||||
|
||||
assert dict(counts) == dict(python_counts)
|
||||
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL], "messages": ANTHROPIC_MESSAGES}
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body,
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
raw_body=json.dumps(body).encode(),
|
||||
)
|
||||
|
||||
assert factory.calls == []
|
||||
assert set(counts) == {ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL}
|
||||
assert not set(counts.values()) & set(RUST_INPUT_TOKENS_BY_TOKENIZER.values())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model", ("replicate/meta/llama-2-70b-chat", "meta-llama/Llama-3-8b", "text-davinci-003"))
|
||||
async def test_models_without_a_rust_tokenizer_stay_in_python(
|
||||
rust_counter: None, monkeypatch: pytest.MonkeyPatch, model: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"text-davinci-003"}
|
||||
)
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
rust_token_counter.TOKEN_COUNTER.override(factory)
|
||||
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
|
||||
python_counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None
|
||||
)
|
||||
|
||||
counts: Final = await count_request_input_tokens(
|
||||
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=json.dumps(body).encode()
|
||||
)
|
||||
|
||||
assert factory.calls == []
|
||||
assert dict(counts) == dict(python_counts)
|
||||
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"expiry_offset, expected_max_budget",
|
||||
|
|
|
|||
|
|
@ -2,180 +2,18 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.spend_tracking.input_tokens import (
|
||||
TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS,
|
||||
count_input_tokens,
|
||||
count_input_tokens_for_model,
|
||||
)
|
||||
from litellm.rust_bridge import bindings, configuration, token_counter
|
||||
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
|
||||
from litellm.rust_bridge.token_counter import RustTokenizer
|
||||
|
||||
ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
CL100K_MODEL: Final = "gpt-4"
|
||||
O200K_MODEL: Final = "gpt-4o"
|
||||
PYTHON_ONLY_MODEL: Final = "replicate/meta/llama-2-70b-chat"
|
||||
MESSAGES: Final = [{"role": "user", "content": "hello"}]
|
||||
RUST_TOKENS: Final = 777
|
||||
|
||||
|
||||
class _FakeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeUpstream(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
"""Stands in for one shared native `Tokenizer`; only its name identifies it."""
|
||||
|
||||
def __init__(self, name: str, json: str | None = None) -> None:
|
||||
self.name = name
|
||||
self.json = json
|
||||
|
||||
|
||||
def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None:
|
||||
"""Point the counter's tokenizer lookups at fakes; the codec path keeps falling back to Python."""
|
||||
fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")}
|
||||
anthropic: Final = _FakeTokenizer("anthropic", anthropic_json)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_encoding", fakes.__getitem__)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic)
|
||||
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _FakeDeclined
|
||||
RustUpstreamError = _FakeUpstream
|
||||
|
||||
|
||||
class _RecordingCounter:
|
||||
def __init__(self, factory: _RecordingFactory, tokenizer: RustTokenizer) -> None:
|
||||
self.factory = factory
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
self.factory.calls.append((self.tokenizer, body))
|
||||
return {"model": "", "input_tokens": RUST_TOKENS}
|
||||
|
||||
|
||||
class _RecordingFactory:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[RustTokenizer, bytes]] = []
|
||||
|
||||
def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RecordingCounter:
|
||||
return _RecordingCounter(self, cast(RustTokenizer, tokenizer.name))
|
||||
|
||||
|
||||
class _DecliningCounter:
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise _FakeDeclined("unsupported request shape")
|
||||
|
||||
|
||||
class _DecliningFactory:
|
||||
def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _DecliningCounter:
|
||||
return _DecliningCounter()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
|
||||
_fake_native_tokenizers(monkeypatch)
|
||||
token_counter.TOKEN_COUNTER.reset()
|
||||
token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
yield
|
||||
token_counter.TOKEN_COUNTER.reset()
|
||||
token_counter._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
def _body(model: object) -> tuple[dict[str, object], bytes]:
|
||||
body: Final = {"model": model, "messages": MESSAGES}
|
||||
return body, json.dumps(body).encode()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_models_sharing_a_tokenizer_are_counted_once_and_merged() -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
token_counter.TOKEN_COUNTER.override(factory)
|
||||
request_body, raw_body = _body([ANTHROPIC_MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", PYTHON_ONLY_MODEL])
|
||||
|
||||
counts: Final = await count_input_tokens(
|
||||
request_body=request_body,
|
||||
raw_body=raw_body,
|
||||
models=(ANTHROPIC_MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", PYTHON_ONLY_MODEL),
|
||||
)
|
||||
|
||||
assert factory.calls == [("anthropic", raw_body), ("cl100k_base", raw_body), ("o200k_base", raw_body)]
|
||||
assert dict(counts) == {
|
||||
ANTHROPIC_MODEL: RUST_TOKENS,
|
||||
CL100K_MODEL: RUST_TOKENS,
|
||||
O200K_MODEL: RUST_TOKENS,
|
||||
"gpt-5": RUST_TOKENS,
|
||||
PYTHON_ONLY_MODEL: count_input_tokens_for_model(request_body=request_body, model=PYTHON_ONLY_MODEL),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rust_disabled_counts_everything_in_python() -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
token_counter.TOKEN_COUNTER.override(factory)
|
||||
request_body, raw_body = _body([ANTHROPIC_MODEL, CL100K_MODEL])
|
||||
|
||||
counts: Final = await count_input_tokens(
|
||||
request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL, CL100K_MODEL)
|
||||
)
|
||||
|
||||
assert factory.calls == []
|
||||
assert dict(counts) == {
|
||||
model: count_input_tokens_for_model(request_body=request_body, model=model)
|
||||
for model in (ANTHROPIC_MODEL, CL100K_MODEL)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_raw_body_counts_in_python() -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
token_counter.TOKEN_COUNTER.override(factory)
|
||||
request_body, _ = _body(ANTHROPIC_MODEL)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=None, models=(ANTHROPIC_MODEL,))
|
||||
|
||||
assert factory.calls == []
|
||||
assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_binding_counts_in_python() -> None:
|
||||
litellm.rust(True)
|
||||
token_counter.TOKEN_COUNTER.override(None)
|
||||
request_body, raw_body = _body(ANTHROPIC_MODEL)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL,))
|
||||
|
||||
assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_declined_request_counts_in_python() -> None:
|
||||
litellm.rust(True)
|
||||
token_counter.TOKEN_COUNTER.override(_DecliningFactory())
|
||||
request_body, raw_body = _body(ANTHROPIC_MODEL)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL,))
|
||||
|
||||
assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL)
|
||||
assert counts[ANTHROPIC_MODEL] != RUST_TOKENS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
"""Tests for the Rust input token counter bridge.
|
||||
"""Tests for the Rust input token counter bridge, called directly rather than through the route catalog.
|
||||
|
||||
The native factory is dependency-injected through ``TOKEN_COUNTER.override``
|
||||
so the fallback cases run without the compiled extension present. The parity
|
||||
cases need the extension and are skipped when it is not built.
|
||||
The factory is passed into ``native_count`` so the caching cases run without the compiled extension
|
||||
present. The parity cases need the extension and are skipped when it is not built.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -16,8 +15,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
|
||||
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model
|
||||
from litellm.rust_bridge import bindings, configuration
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model
|
||||
from litellm.rust_bridge import token_counter as bridge
|
||||
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
|
||||
from litellm.rust_bridge._native import Tokenizer
|
||||
|
|
@ -38,14 +36,6 @@ def _counted(body: dict[str, object], model: str) -> tuple[bytes, dict[str, obje
|
|||
return raw, json.loads(raw)
|
||||
|
||||
|
||||
class _FakeDeclined(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeUpstream(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
"""Stands in for one shared native `Tokenizer`; only its name identifies it."""
|
||||
|
||||
|
|
@ -54,29 +44,6 @@ class _FakeTokenizer:
|
|||
self.json = json
|
||||
|
||||
|
||||
def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None:
|
||||
"""Point the counter's tokenizer lookups at fakes while the bridge is faked; the codec path
|
||||
keeps falling back to Python. Parity tests that restore the real extension get the real
|
||||
lookups back."""
|
||||
fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")}
|
||||
anthropic: Final = _FakeTokenizer("anthropic", anthropic_json)
|
||||
real_encoding: Final = tokenizer_dispatch.native_encoding
|
||||
real_anthropic: Final = tokenizer_dispatch.native_anthropic
|
||||
|
||||
def faked() -> bool:
|
||||
return isinstance(bindings.get_native_bridge(), _FakeNative)
|
||||
|
||||
monkeypatch.setattr(
|
||||
tokenizer_dispatch, "native_encoding", lambda name: fakes[name] if faked() else real_encoding(name)
|
||||
)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic if faked() else real_anthropic())
|
||||
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _FakeDeclined
|
||||
RustUpstreamError = _FakeUpstream
|
||||
|
||||
|
||||
class _RecordingCounter:
|
||||
def __init__(self, tokenizer: _FakeTokenizer, fast: bool) -> None:
|
||||
self.tokenizer = tokenizer
|
||||
|
|
@ -100,57 +67,25 @@ class _RecordingFactory:
|
|||
return counter
|
||||
|
||||
|
||||
class _RaisingCounter:
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.error = error
|
||||
|
||||
async def acount_request(self, body: bytes) -> object:
|
||||
raise self.error
|
||||
|
||||
|
||||
class _RaisingFactory:
|
||||
"""Every counter it builds, for either tokenizer, raises `error` on count."""
|
||||
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.error = error
|
||||
|
||||
def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RaisingCounter:
|
||||
return _RaisingCounter(self.error)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
|
||||
bridge.TOKEN_COUNTER.reset()
|
||||
def _reset_counters():
|
||||
bridge._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
|
||||
_fake_native_tokenizers(monkeypatch, anthropic_json=claude_json_str)
|
||||
yield
|
||||
bridge.TOKEN_COUNTER.reset()
|
||||
bridge._counter.cache_clear()
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_tokenizers(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Point the counter's tokenizer lookups at fakes so a recording factory sees which one it was built over."""
|
||||
fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")}
|
||||
anthropic: Final = _FakeTokenizer("anthropic", claude_json_str)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_encoding", fakes.__getitem__)
|
||||
monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.RustTokenizer) -> None:
|
||||
async def test_native_count_returns_typed_count_and_reuses_one_counter(fake_tokenizers: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(False)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
model: Final = MODEL_BY_TOKENIZER[tokenizer]
|
||||
raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,))
|
||||
|
||||
assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model)
|
||||
assert factory.counters == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
first: Final = await bridge.native_count(factory, "anthropic", BODY)
|
||||
second: Final = await bridge.native_count(factory, "anthropic", BODY)
|
||||
|
|
@ -166,10 +101,10 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base"))
|
||||
async def test_tiktoken_counter_is_built_over_the_shared_encoding_once(tokenizer: bridge.RustTokenizer) -> None:
|
||||
async def test_tiktoken_counter_is_built_over_the_shared_encoding_once(
|
||||
fake_tokenizers: None, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
first: Final = await bridge.native_count(factory, tokenizer, BODY)
|
||||
second: Final = await bridge.native_count(factory, tokenizer, BODY)
|
||||
|
|
@ -183,10 +118,8 @@ async def test_tiktoken_counter_is_built_over_the_shared_encoding_once(tokenizer
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
|
||||
async def test_each_tokenizer_gets_its_own_cached_counter(fake_tokenizers: None) -> None:
|
||||
factory: Final = _RecordingFactory()
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(factory)
|
||||
|
||||
await bridge.native_count(factory, "anthropic", BODY)
|
||||
await bridge.native_count(factory, "cl100k_base", BODY)
|
||||
|
|
@ -198,43 +131,6 @@ async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
|
|||
assert [len(counter.bodies) for counter in factory.counters] == [2, 1, 2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
litellm.rust(True)
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
|
||||
raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, MODEL)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(MODEL,))
|
||||
|
||||
assert counts[MODEL] == count_input_tokens_for_model(request_body=request_body, model=MODEL)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages")))
|
||||
model: Final = MODEL_BY_TOKENIZER[tokenizer]
|
||||
raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,))
|
||||
|
||||
assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None:
|
||||
litellm.rust(True)
|
||||
bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed")))
|
||||
model: Final = MODEL_BY_TOKENIZER[tokenizer]
|
||||
raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,))
|
||||
|
||||
assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected"),
|
||||
(
|
||||
|
|
@ -416,39 +312,34 @@ PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
|
|||
@pytest.mark.parametrize(("model", "tokenizer"), PARITY_MODELS)
|
||||
@pytest.mark.parametrize("request_body", PARITY_REQUESTS)
|
||||
async def test_native_count_matches_python_budget_counter(
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
|
||||
request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
body: Final = json.dumps(request_body).replace(MODEL, model)
|
||||
parsed: Final = json.loads(body)
|
||||
|
||||
request_body_parsed: Final = json.loads(body)
|
||||
counts: Final = await count_input_tokens(request_body=request_body_parsed, raw_body=body.encode(), models=(model,))
|
||||
python_count: Final = count_input_tokens_for_model(request_body=request_body_parsed, model=model)
|
||||
counted: Final = await bridge.native_count(native.TokenCounter, tokenizer, body.encode())
|
||||
|
||||
assert counts[model] == python_count
|
||||
assert counted.input_tokens == count_input_tokens_for_model(request_body=parsed, model=model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("model", "tokenizer"), ((CL100K_MODEL, "cl100k_base"), (O200K_MODEL, "o200k_base")))
|
||||
async def test_tiktoken_counts_long_text_exactly_where_python_chunks(
|
||||
monkeypatch: pytest.MonkeyPatch, model: str, tokenizer: bridge.RustTokenizer
|
||||
model: str, tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
text: Final = "x " * 20_000
|
||||
body: Final = {"model": model, "messages": [{"role": "user", "content": text}]}
|
||||
encoding: Final = Tokenizer.from_tiktoken(tokenizer)
|
||||
exact: Final = 3 + encoding.count("user") + encoding.count(text) + 3
|
||||
chunks: Final = -(-len(text) // TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS)
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=body, raw_body=json.dumps(body).encode(), models=(model,))
|
||||
counted: Final = await bridge.native_count(native.TokenCounter, tokenizer, json.dumps(body).encode())
|
||||
python_count: Final = count_input_tokens_for_model(request_body=body, model=model)
|
||||
|
||||
assert counts[model] == exact
|
||||
assert counted.input_tokens == exact
|
||||
assert python_count is not None
|
||||
assert exact < python_count <= exact + chunks
|
||||
|
||||
|
|
@ -470,14 +361,10 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
|
|||
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
|
||||
@pytest.mark.parametrize("request_body", DECLINED_REQUESTS)
|
||||
async def test_native_declines_shapes_python_prices_differently(
|
||||
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], tokenizer: bridge.RustTokenizer
|
||||
request_body: dict[str, object], tokenizer: bridge.RustTokenizer
|
||||
) -> None:
|
||||
native: Final = pytest.importorskip("litellm.rust_bridge._native")
|
||||
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
|
||||
litellm.rust(True)
|
||||
model: Final = MODEL_BY_TOKENIZER[tokenizer]
|
||||
raw, parsed = _counted(request_body, model)
|
||||
raw, _ = _counted(request_body, MODEL_BY_TOKENIZER[tokenizer])
|
||||
|
||||
counts: Final = await count_input_tokens(request_body=parsed, raw_body=raw, models=(model,))
|
||||
|
||||
assert counts.get(model) == count_input_tokens_for_model(request_body=parsed, model=model)
|
||||
with pytest.raises(native.RustBridgeDeclined):
|
||||
await bridge.native_count(native.TokenCounter, tokenizer, raw)
|
||||
|
|
|
|||
|
|
@ -1,134 +1,49 @@
|
|||
from collections.abc import Generator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import tiktoken
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
|
||||
from litellm.rust_bridge import configuration, tokenizer
|
||||
from litellm.utils import _select_tokenizer
|
||||
from litellm.rust_bridge import tokenizer
|
||||
from litellm.utils import claude_json_str
|
||||
from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
|
||||
monkeypatch.delenv("LITELLM_RUST", raising=False)
|
||||
configuration.reset_rust_configuration()
|
||||
yield
|
||||
tokenizer.TOKENIZER.reset()
|
||||
configuration.reset_rust_configuration()
|
||||
TEXTS: Final = ("hello <|endoftext|> world", "café 漢字 🙂", " def f():\n return 1\n", "<SOS>hello<EOT> again")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("environment", (None, "0", "1"))
|
||||
@pytest.mark.parametrize("process", (None, False, True))
|
||||
def test_tokenizer_factories_follow_rollout(
|
||||
monkeypatch: pytest.MonkeyPatch, environment: str | None, process: bool | None
|
||||
) -> None:
|
||||
configuration.rust(process)
|
||||
if environment is not None:
|
||||
monkeypatch.setenv("LITELLM_RUST", environment)
|
||||
enabled: Final = environment == "1" if environment is not None else process is True
|
||||
encoding: Final = tokenizer.get_encoding("cl100k_base")
|
||||
custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)
|
||||
@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base"))
|
||||
@pytest.mark.parametrize("text", TEXTS)
|
||||
def test_native_encoding_matches_tiktoken(name: str, text: str) -> None:
|
||||
native: Final = tokenizer.native_encoding(name)
|
||||
if native is None:
|
||||
pytest.skip("native extension is not built")
|
||||
encoding: Final = OpenAIEncoding.wrap(native)
|
||||
reference: Final = tiktoken.get_encoding(name)
|
||||
|
||||
ids: Final = encoding.encode(text, disallowed_special=())
|
||||
assert ids == reference.encode(text, disallowed_special=())
|
||||
assert encoding.decode(ids) == reference.decode(ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", TEXTS)
|
||||
def test_native_anthropic_tokenizer_matches_python(text: str) -> None:
|
||||
native: Final = tokenizer.native_anthropic()
|
||||
if native is None:
|
||||
pytest.skip("native extension is not built")
|
||||
reference: Final = Tokenizer.from_str(claude_json_str)
|
||||
|
||||
ids: Final = HuggingFaceTokenizer(native).encode(text).ids
|
||||
assert ids == reference.encode(text).ids
|
||||
assert HuggingFaceTokenizer(native).decode(ids) == reference.decode(ids)
|
||||
|
||||
|
||||
def test_native_custom_tokenizer_matches_python() -> None:
|
||||
factory: Final = tokenizer.TOKENIZER.load()
|
||||
if factory is None:
|
||||
pytest.skip("native extension is not built")
|
||||
native: Final = HuggingFaceTokenizer(factory.from_json(TOKENIZER_JSON))
|
||||
reference: Final = Tokenizer.from_str(TOKENIZER_JSON)
|
||||
|
||||
assert isinstance(encoding, OpenAIEncoding if enabled else tiktoken.Encoding)
|
||||
assert isinstance(custom["tokenizer"], HuggingFaceTokenizer if enabled else Tokenizer)
|
||||
assert encoding.encode("café 漢字 🙂") == tiktoken.get_encoding(encoding.name).encode("café 漢字 🙂")
|
||||
assert litellm.encode(text="Hello World", custom_tokenizer=custom) == reference.encode("Hello World").ids
|
||||
assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(reference.encode("Hello World"))
|
||||
|
||||
|
||||
def test_missing_native_binding_keeps_python_tokenizer_api() -> None:
|
||||
configuration.rust(True)
|
||||
tokenizer.TOKENIZER.override(None)
|
||||
encoding: Final = tokenizer.get_encoding("cl100k_base")
|
||||
custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)["tokenizer"]
|
||||
|
||||
assert isinstance(encoding, tiktoken.Encoding)
|
||||
assert isinstance(custom, Tokenizer)
|
||||
custom.enable_padding(pad_id=0, pad_token="[UNK]")
|
||||
assert [item.ids for item in custom.encode_batch(["Hello", "Hello World"])] == [[3, 1, 0], [3, 1, 2]]
|
||||
|
||||
|
||||
def test_cached_selection_follows_backend_changes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
|
||||
configuration.rust(True)
|
||||
native: Final = _select_tokenizer("dispatch-fixture")["tokenizer"]
|
||||
configuration.rust(False)
|
||||
python: Final = _select_tokenizer("dispatch-fixture")["tokenizer"]
|
||||
|
||||
assert isinstance(native, OpenAIEncoding)
|
||||
assert isinstance(python, tiktoken.Encoding)
|
||||
assert native.encode("hello") == python.encode("hello")
|
||||
|
||||
|
||||
def test_declined_native_factory_falls_back_before_tokenizing() -> None:
|
||||
from litellm.rust_bridge._native import RustBridgeDeclined
|
||||
|
||||
class UnavailableTokenizer:
|
||||
@staticmethod
|
||||
def from_json(json: str) -> None:
|
||||
raise RustBridgeDeclined("huggingface feature is disabled")
|
||||
|
||||
configuration.rust(True)
|
||||
binding: Final = tokenizer._as_factory(UnavailableTokenizer)
|
||||
tokenizer.TOKENIZER.override(binding)
|
||||
custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)
|
||||
|
||||
assert isinstance(custom["tokenizer"], Tokenizer)
|
||||
assert (
|
||||
litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom)
|
||||
== "Hello World"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "text"),
|
||||
(
|
||||
("gpt-4o", "hello <|endoftext|> world"),
|
||||
("gpt-3.5-turbo", "café 漢字 🙂"),
|
||||
("text-davinci-003", " def f():\n return 1\n"),
|
||||
("tokenizer-parity-fixture", "<SOS>hello<EOT> again"),
|
||||
),
|
||||
)
|
||||
def test_public_token_api_is_identical_across_backends(monkeypatch: pytest.MonkeyPatch, model: str, text: str) -> None:
|
||||
"""`litellm.token_counter`, `encode` and `decode` return the same values whichever backend
|
||||
the catalog picks; only the object types differ."""
|
||||
monkeypatch.setattr(litellm, "anthropic_models", {*litellm.anthropic_models, "tokenizer-parity-fixture"})
|
||||
messages: Final = [{"role": "user", "content": text}, {"role": "assistant", "content": "ok"}]
|
||||
|
||||
def observe() -> tuple[int, int, list[int], str]:
|
||||
ids: Final = litellm.encode(model=model, text=text)
|
||||
return (
|
||||
litellm.token_counter(model=model, text=text),
|
||||
litellm.token_counter(model=model, messages=messages),
|
||||
ids,
|
||||
litellm.decode(model=model, tokens=ids),
|
||||
)
|
||||
|
||||
configuration.rust(False)
|
||||
python: Final = observe()
|
||||
configuration.rust(True)
|
||||
rust: Final = observe()
|
||||
|
||||
assert rust == python
|
||||
|
||||
|
||||
def test_cached_huggingface_tokenizers_follow_backend_changes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer as RustHuggingFaceTokenizer
|
||||
from litellm.utils import _load_huggingface_tokenizer
|
||||
|
||||
monkeypatch.setattr(litellm, "anthropic_models", {*litellm.anthropic_models, "tokenizer-cache-fixture"})
|
||||
_load_huggingface_tokenizer.cache_clear()
|
||||
configuration.rust(True)
|
||||
native: Final = _select_tokenizer("tokenizer-cache-fixture")["tokenizer"]
|
||||
configuration.rust(False)
|
||||
python: Final = _select_tokenizer("tokenizer-cache-fixture")["tokenizer"]
|
||||
configuration.rust(True)
|
||||
|
||||
assert isinstance(native, RustHuggingFaceTokenizer)
|
||||
assert isinstance(python, Tokenizer)
|
||||
assert _select_tokenizer("tokenizer-cache-fixture")["tokenizer"] is native
|
||||
assert native.encode("Hello World").ids == reference.encode("Hello World").ids
|
||||
assert native.decode(reference.encode("Hello World").ids) == reference.decode(reference.encode("Hello World").ids)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue