feat: implement download and configure tiktoken cache (#15)

Author: bsbds

deepdoc/common/tiktoken_cache.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import uuid
+from pathlib import Path
+from urllib.request import urlopen
+
+
+DEEPDOC_TIKTOKEN_CACHE_DIR_ENV = "DEEPDOC_TIKTOKEN_CACHE_DIR"
+CL100K_BASE_BLOB_URL = "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken"
+CL100K_BASE_EXPECTED_HASH = "223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7"
+
+
+def resolve_tiktoken_cache_dir(cache_dir: str | None = None, model_home: str | None = None) -> Path:
+    if cache_dir:
+        return Path(cache_dir).expanduser().resolve()
+
+    explicit_cache_dir = os.getenv(DEEPDOC_TIKTOKEN_CACHE_DIR_ENV) or os.getenv("TIKTOKEN_CACHE_DIR")
+    if explicit_cache_dir:
+        return Path(explicit_cache_dir).expanduser().resolve()
+
+    configured_model_home = model_home or os.getenv("DEEPDOC_MODEL_HOME")
+    if configured_model_home:
+        return Path(configured_model_home).expanduser().resolve().joinpath("tiktoken_cache")
+
+    return Path.home().joinpath(".cache", "deepdoc", "tiktoken_cache").resolve()
+
+
+def configure_tiktoken_cache_env(cache_dir: str | None = None, model_home: str | None = None) -> str:
+    resolved_cache_dir = resolve_tiktoken_cache_dir(cache_dir=cache_dir, model_home=model_home)
+    os.environ["TIKTOKEN_CACHE_DIR"] = str(resolved_cache_dir)
+    return str(resolved_cache_dir)
+
+
+def cl100k_base_cache_key(blob_url: str = CL100K_BASE_BLOB_URL) -> str:
+    return hashlib.sha1(blob_url.encode()).hexdigest()
+
+
+def _matches_expected_hash(data: bytes, expected_hash: str | None) -> bool:
+    if not expected_hash:
+        return True
+    return hashlib.sha256(data).hexdigest() == expected_hash
+
+
+def download_cl100k_base(
+    *,
+    cache_dir: str | None = None,
+    model_home: str | None = None,
+    offline: bool = False,
+    blob_url: str = CL100K_BASE_BLOB_URL,
+    expected_hash: str = CL100K_BASE_EXPECTED_HASH,
+    timeout: int = 60,
+) -> Path:
+    resolved_cache_dir = resolve_tiktoken_cache_dir(cache_dir=cache_dir, model_home=model_home)
+    cache_key = cl100k_base_cache_key(blob_url)
+    target_path = resolved_cache_dir.joinpath(cache_key)
+
+    if target_path.exists():
+        data = target_path.read_bytes()
+        if _matches_expected_hash(data, expected_hash):
+            return target_path
+        target_path.unlink()
+
+    if offline:
+        raise FileNotFoundError(
+            "Missing cached tiktoken encoder '{}'. Expected file at {}. Run the download command without --offline first."
+            .format(cache_key, target_path)
+        )
+
+    with urlopen(blob_url, timeout=timeout) as response:
+        data = response.read()
+
+    if not _matches_expected_hash(data, expected_hash):
+        raise ValueError(
+            "Hash mismatch for tiktoken encoder downloaded from {}. Expected SHA256 {}."
+            .format(blob_url, expected_hash)
+        )
+
+    resolved_cache_dir.mkdir(parents=True, exist_ok=True)
+    tmp_path = target_path.with_name("{}.{}.tmp".format(target_path.name, uuid.uuid4().hex))
+    tmp_path.write_bytes(data)
+    tmp_path.replace(target_path)
+    return target_path

deepdoc/common/token_utils.py

@@ -14,21 +14,26 @@
 #  limitations under the License.
 #
 
-import os
 import tiktoken
 
-from .file_utils import get_project_base_directory
+from .tiktoken_cache import configure_tiktoken_cache_env
 
-tiktoken_cache_dir = get_project_base_directory()
-os.environ["TIKTOKEN_CACHE_DIR"] = tiktoken_cache_dir
-# encoder = tiktoken.encoding_for_model("gpt-3.5-turbo")
-encoder = tiktoken.get_encoding("cl100k_base")
+
+_encoder = None
+
+
+def _get_encoder():
+    global _encoder
+    if _encoder is None:
+        configure_tiktoken_cache_env()
+        _encoder = tiktoken.get_encoding("cl100k_base")
+    return _encoder
 
 
 def num_tokens_from_string(string: str) -> int:
     """Returns the number of tokens in a text string."""
     try:
-        code_list = encoder.encode(string)
+        code_list = _get_encoder().encode(string)
         return len(code_list)
     except Exception:
         return 0
@@ -84,4 +89,8 @@ def total_token_count_from_response(resp):
 
 def truncate(string: str, max_len: int) -> str:
     """Returns truncated text if the length of text exceed max_len."""
-    return encoder.decode(encoder.encode(string)[:max_len])
+    try:
+        encoder = _get_encoder()
+        return encoder.decode(encoder.encode(string)[:max_len])
+    except Exception:
+        return string[:max_len]

deepdoc/depend/num_tokens_from_string.py

@@ -0,0 +1,80 @@
+import hashlib
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+from deepdoc.common import tiktoken_cache as tc
+
+
+class _FakeResponse:
+    def __init__(self, data: bytes) -> None:
+        self._data = data
+
+    def read(self) -> bytes:
+        return self._data
+
+    def __enter__(self) -> "_FakeResponse":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        return None
+
+
+class TestTiktokenCache(unittest.TestCase):
+    def setUp(self) -> None:
+        self._old_env = os.environ.copy()
+
+    def tearDown(self) -> None:
+        os.environ.clear()
+        os.environ.update(self._old_env)
+
+    def test_resolve_tiktoken_cache_dir_uses_model_home(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            os.environ.pop("TIKTOKEN_CACHE_DIR", None)
+            os.environ.pop(tc.DEEPDOC_TIKTOKEN_CACHE_DIR_ENV, None)
+            os.environ["DEEPDOC_MODEL_HOME"] = tmp
+
+            resolved = tc.resolve_tiktoken_cache_dir()
+
+            self.assertEqual(resolved, Path(tmp).resolve().joinpath("tiktoken_cache"))
+
+    def test_download_cl100k_base_writes_cache_key_named_file(self) -> None:
+        payload = b"test-tiktoken-data"
+        blob_url = "https://example.com/cl100k_base.tiktoken"
+        expected_hash = hashlib.sha256(payload).hexdigest()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch.object(tc, "urlopen", return_value=_FakeResponse(payload)) as mocked_urlopen:
+                target = tc.download_cl100k_base(
+                    cache_dir=tmp,
+                    blob_url=blob_url,
+                    expected_hash=expected_hash,
+                )
+
+            self.assertEqual(target, Path(tmp).resolve().joinpath(hashlib.sha1(blob_url.encode()).hexdigest()))
+            self.assertEqual(target.read_bytes(), payload)
+            mocked_urlopen.assert_called_once_with(blob_url, timeout=60)
+
+    def test_download_cl100k_base_offline_uses_existing_cache(self) -> None:
+        payload = b"cached-tiktoken-data"
+        blob_url = "https://example.com/cl100k_base.tiktoken"
+        expected_hash = hashlib.sha256(payload).hexdigest()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            target = Path(tmp).resolve().joinpath(hashlib.sha1(blob_url.encode()).hexdigest())
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_bytes(payload)
+
+            with patch.object(tc, "urlopen") as mocked_urlopen:
+                resolved = tc.download_cl100k_base(
+                    cache_dir=tmp,
+                    blob_url=blob_url,
+                    expected_hash=expected_hash,
+                    offline=True,
+                )
+
+            self.assertEqual(resolved, target)
+            mocked_urlopen.assert_not_called()