mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(main): use local tiktoken cache in lazy loading (#19774)
The lazy loading implementation for encoding in __getattr__ was calling tiktoken.get_encoding() directly without first setting TIKTOKEN_CACHE_DIR. This caused tiktoken to attempt downloading the encoding file from the internet instead of using the local copy bundled with litellm. This fix uses _get_default_encoding() from _lazy_imports which properly sets TIKTOKEN_CACHE_DIR before loading tiktoken, ensuring the local cache is used.
This commit is contained in:
parent
920ef665a3
commit
807ba011eb
2 changed files with 34 additions and 2 deletions
|
|
@ -7300,8 +7300,11 @@ def _get_encoding():
|
|||
def __getattr__(name: str) -> Any:
|
||||
"""Lazy import handler for main module"""
|
||||
if name == "encoding":
|
||||
# Lazy load encoding to avoid heavy tiktoken import at module load time
|
||||
_encoding = tiktoken.get_encoding("cl100k_base")
|
||||
# Use _get_default_encoding which properly sets TIKTOKEN_CACHE_DIR
|
||||
# before loading tiktoken, ensuring the local cache is used
|
||||
# instead of downloading from the internet
|
||||
from litellm._lazy_imports import _get_default_encoding
|
||||
_encoding = _get_default_encoding()
|
||||
# Cache it in the module's __dict__ for subsequent accesses
|
||||
import sys
|
||||
|
||||
|
|
|
|||
|
|
@ -78,6 +78,35 @@ def test_lazy_loading_default():
|
|||
assert len(tokens) > 0, "Encoding should work"
|
||||
|
||||
|
||||
def test_tiktoken_cache_dir_set_on_lazy_load():
|
||||
"""Test that TIKTOKEN_CACHE_DIR is set when encoding is lazy loaded.
|
||||
|
||||
This ensures the local tiktoken cache is used instead of downloading
|
||||
from the internet. Regression test for issue #19768.
|
||||
"""
|
||||
# Remove environment variables to ensure clean state
|
||||
if "LITELLM_DISABLE_LAZY_LOADING" in os.environ:
|
||||
del os.environ["LITELLM_DISABLE_LAZY_LOADING"]
|
||||
if "TIKTOKEN_CACHE_DIR" in os.environ:
|
||||
del os.environ["TIKTOKEN_CACHE_DIR"]
|
||||
|
||||
# Clear any cached modules
|
||||
modules_to_clear = [k for k in sys.modules.keys() if k.startswith("litellm")]
|
||||
for module in modules_to_clear:
|
||||
del sys.modules[module]
|
||||
|
||||
# Import litellm fresh
|
||||
import litellm
|
||||
|
||||
# Access encoding (triggers lazy load)
|
||||
_ = litellm.encoding
|
||||
|
||||
# Verify TIKTOKEN_CACHE_DIR is now set and points to local tokenizers
|
||||
assert "TIKTOKEN_CACHE_DIR" in os.environ, "TIKTOKEN_CACHE_DIR should be set after lazy loading encoding"
|
||||
cache_dir = os.environ["TIKTOKEN_CACHE_DIR"]
|
||||
assert "tokenizers" in cache_dir, f"TIKTOKEN_CACHE_DIR should point to tokenizers directory, got: {cache_dir}"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_env():
|
||||
"""Clean up environment variable after each test"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue