From e997fd655a0aa9cb1aecd6e2b347866483177cfa Mon Sep 17 00:00:00 2001 From: Jack Date: Mon, 3 Aug 2026 19:03:08 +0800 Subject: [PATCH] fix(tokenizer): load cl100k BPE table from disk instead of failing silently offline (#17712) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary RAGFlow's Go tokenizer silently returned **0 tokens for every string** whenever the `cl100k_base` BPE table could not be loaded — which is the normal case for an offline/air-gapped Go server. This PR makes the loader resolve the table from disk (where RAGFlow actually ships it) and fail loudly when it is genuinely missing. ## Root cause `tiktoken-go`'s stock loader downloads the encoding table over HTTP and caches it under `TIKTOKEN_CACHE_DIR`. That does not work for RAGFlow: - `TIKTOKEN_CACHE_DIR` is exported **only inside the Python process** (`common/token_utils.py`). `docker/entrypoint.sh` launches the Go binary (`bin/ragflow_server`) from a shell, so the Go process never inherits the variable. - The Dockerfile *does* ship the table (under its sha1 name in the working directory), but nothing told the Go side to look there. - Reaching `openaipublic.blob.core.windows.net` at runtime is not an option for air-gapped installs, and is unreliable where that host is blocked. The failure was **silent**: `NumTokensFromString` returns `0` when the encoder fails to build, and a `sync.Once` memoizes that error for the process lifetime. Every token count became `0`, so chunk merging never crossed its token budget and an entire document collapsed into a single chunk. Python has no such failure mode because its encoder is built at import time (a missing table aborts startup instead of degrading). ## Fix Register a local-only `BpeLoader` via `tiktoken.SetBpeLoader` (`internal/tokenizer/bpe_loader.go`) that resolves the table from disk **only**, in priority order: 1. `TIKTOKEN_CACHE_DIR` / `DATA_GYM_CACHE_DIR` (honored so operators who already configured one keep working). 2. The working directory, the executable's directory, and all of their ancestors — matching the Dockerfile layout (table under its sha1 name in the install root). 3. A `ragflow_deps/` checkout produced by `ragflow_deps/download_deps.py`. It **never performs network I/O**. When nothing is found it returns an error listing every path it tried (pointing at `download_deps.py` or `TIKTOKEN_CACHE_DIR`), so a genuinely missing table fails loudly instead of degrading to zero. ## Test plan - `internal/tokenizer/bpe_loader_test.go` (unit tier, runs under `bash build.sh --test ./internal/tokenizer/...`): - Loader reads from `TIKTOKEN_CACHE_DIR`, `DATA_GYM_CACHE_DIR`, the sha1-named file in the working dir, and the bundled `ragflow_deps/` name. - Explicit cache dir wins over the bundled vocab. - A malformed table is reported as an error rather than skipped. - A genuinely missing table reports the candidates it tried (no network attempt). - `NumTokensFromString` matches Python-derived anchors (`""`→0, `"hello"`→1, `"hello world"`→2, `"hello, world!"`→4, `"世界"`→3, `"Hello 世界 🌍"`→8, `"RAGFlow"`→3). ## Notes - `.github/workflows/tests.yml` currently excludes `internal/tokenizer` from `go test`, so these tests do not run in CI. The tokenizer fix is exercised in CI indirectly via the chunker package once a token-count-sensitive parity case lands (tracked separately). Consider including `internal/tokenizer` in CI as a follow-up. - Supported deployments already ship the table (`download_deps.py` → `ragflow_deps/cl100k_base.tiktoken`; Dockerfile → `` in cwd), so no `ENV` change is required for the fix to take effect. Setting `ENV TIKTOKEN_CACHE_DIR` in the Dockerfile remains a cheap belt-and-suspenders hardening that can be done separately. 🤖 Generated with [CodeBuddy Code](https://cnb.cool/codebuddy) --------- Co-authored-by: CodeBuddy Co-authored-by: CodeBuddy Code Co-authored-by: CodeBuddy --- .github/workflows/sep-tests.yml | 98 +++++++- .github/workflows/tests.yml | 96 +++++++- internal/tokenizer/bpe_loader.go | 216 ++++++++++++++++++ internal/tokenizer/bpe_loader_anchors_test.go | 77 +++++++ internal/tokenizer/bpe_loader_test.go | 201 ++++++++++++++++ .../tokenizer/tokenizer_concurrent_test.go | 2 +- internal/tokenizer/tokenizer_test.go | 2 +- 7 files changed, 673 insertions(+), 19 deletions(-) create mode 100644 internal/tokenizer/bpe_loader.go create mode 100644 internal/tokenizer/bpe_loader_anchors_test.go create mode 100644 internal/tokenizer/bpe_loader_test.go diff --git a/.github/workflows/sep-tests.yml b/.github/workflows/sep-tests.yml index d6da4615bd..2955799498 100644 --- a/.github/workflows/sep-tests.yml +++ b/.github/workflows/sep-tests.yml @@ -320,14 +320,55 @@ jobs: # # Excludes packages whose tests fail for environmental reasons # unrelated to the diff: - # - internal/tokenizer: tests need /usr/share/infinity/resource - # dict files, only mounted inside the docker builder, not - # in the Go test environment. + # - internal/tokenizer is split: the pure-Go BPE loader tests + # (bpe_loader_test.go) run in the default tier; the C++ binding / + # dict-dependent tests (tokenizer_test.go, + # tokenizer_concurrent_test.go) and the on-disk anchor test + # (bpe_loader_anchors_test.go) are tagged `manual` and need the + # docker builder's /usr/share/infinity/resource, so they stay out + # of the default run. run: | set -euo pipefail + + # Provide the cl100k BPE table for the offline tokenizer loader. + # The loader reads /ragflow_deps/cl100k_base.tiktoken (or + # TIKTOKEN_CACHE_DIR). Try to fetch it directly from the upstream + # openai blob first (small ~1.6MB, reachable from the CI runner — the + # pre-PR code downloaded it on the fly, so the network path is known + # good). If the direct fetch fails (e.g. GFW), fall back to extracting + # it from the local infiniflow/ragflow_deps image *if already present* + # on the runner (no `docker pull`, since that huge image is first + # fetched at the later "Build ragflow:nightly" step). Best-effort: if + # neither source works, the loader still fails loudly (no silent + # 0-token degradation). + if [ ! -f ragflow_deps/cl100k_base.tiktoken ]; then + mkdir -p ragflow_deps + if command -v curl >/dev/null 2>&1; then + if curl -fsSL -o ragflow_deps/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken; then + echo "tiktoken: cl100k table fetched via curl from openai blob" + else + echo "tiktoken: curl fetch failed (network/GFW); will try cached image" + fi + fi + if [ ! -f ragflow_deps/cl100k_base.tiktoken ] && docker image inspect infiniflow/ragflow_deps:latest >/dev/null 2>&1; then + CID=$(docker create infiniflow/ragflow_deps:latest true) || true + if [ -n "${CID:-}" ]; then + if docker cp "$CID":/cl100k_base.tiktoken ragflow_deps/cl100k_base.tiktoken 2>/dev/null; then + echo "tiktoken: cl100k table copied from cached infiniflow/ragflow_deps image" + fi + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + fi + fi + if [ -f ragflow_deps/cl100k_base.tiktoken ]; then + echo "tiktoken: cl100k table provisioned ($(wc -c < ragflow_deps/cl100k_base.tiktoken) bytes)" + else + echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" + fi + PKGS=$(go list ./... 2>/dev/null \ | grep -v '/internal/storage$' \ - | grep -v '/internal/tokenizer$' \ | grep -v '/internal/handler$' || true) if [ -z "$PKGS" ]; then ./build.sh --test @@ -895,14 +936,55 @@ jobs: # # Excludes packages whose tests fail for environmental reasons # unrelated to the diff: - # - internal/tokenizer: tests need /usr/share/infinity/resource - # dict files, only mounted inside the docker builder, not - # in the Go test environment. + # - internal/tokenizer is split: the pure-Go BPE loader tests + # (bpe_loader_test.go) run in the default tier; the C++ binding / + # dict-dependent tests (tokenizer_test.go, + # tokenizer_concurrent_test.go) and the on-disk anchor test + # (bpe_loader_anchors_test.go) are tagged `manual` and need the + # docker builder's /usr/share/infinity/resource, so they stay out + # of the default run. run: | set -euo pipefail + + # Provide the cl100k BPE table for the offline tokenizer loader. + # The loader reads /ragflow_deps/cl100k_base.tiktoken (or + # TIKTOKEN_CACHE_DIR). Try to fetch it directly from the upstream + # openai blob first (small ~1.6MB, reachable from the CI runner — the + # pre-PR code downloaded it on the fly, so the network path is known + # good). If the direct fetch fails (e.g. GFW), fall back to extracting + # it from the local infiniflow/ragflow_deps image *if already present* + # on the runner (no `docker pull`, since that huge image is first + # fetched at the later "Build ragflow:nightly" step). Best-effort: if + # neither source works, the loader still fails loudly (no silent + # 0-token degradation). + if [ ! -f ragflow_deps/cl100k_base.tiktoken ]; then + mkdir -p ragflow_deps + if command -v curl >/dev/null 2>&1; then + if curl -fsSL -o ragflow_deps/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken; then + echo "tiktoken: cl100k table fetched via curl from openai blob" + else + echo "tiktoken: curl fetch failed (network/GFW); will try cached image" + fi + fi + if [ ! -f ragflow_deps/cl100k_base.tiktoken ] && docker image inspect infiniflow/ragflow_deps:latest >/dev/null 2>&1; then + CID=$(docker create infiniflow/ragflow_deps:latest true) || true + if [ -n "${CID:-}" ]; then + if docker cp "$CID":/cl100k_base.tiktoken ragflow_deps/cl100k_base.tiktoken 2>/dev/null; then + echo "tiktoken: cl100k table copied from cached infiniflow/ragflow_deps image" + fi + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + fi + fi + if [ -f ragflow_deps/cl100k_base.tiktoken ]; then + echo "tiktoken: cl100k table provisioned ($(wc -c < ragflow_deps/cl100k_base.tiktoken) bytes)" + else + echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" + fi + PKGS=$(go list ./... 2>/dev/null \ | grep -v '/internal/storage$' \ - | grep -v '/internal/tokenizer$' \ | grep -v '/internal/handler$' || true) if [ -z "$PKGS" ]; then ./build.sh --test diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 578522de39..8ab708253e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -197,15 +197,54 @@ jobs: # # Excludes packages whose tests fail for environmental reasons # unrelated to the diff: - # - internal/tokenizer: tests need /usr/share/infinity/resource - # dict files, only mounted inside the docker builder, not - # in the Go test environment. + # - internal/tokenizer is split: the pure-Go BPE loader tests + # (bpe_loader_test.go) run in the default tier; the C++ binding / + # dict-dependent tests (tokenizer_test.go, + # tokenizer_concurrent_test.go) and the on-disk anchor test + # (bpe_loader_anchors_test.go) are tagged `manual` and need the + # docker builder's /usr/share/infinity/resource, so they stay out + # of the default run. run: | set -euo pipefail + + # Provide the cl100k BPE table for the offline tokenizer loader. + # The loader reads /ragflow_deps/cl100k_base.tiktoken (or + # TIKTOKEN_CACHE_DIR). Fetch it directly: this CI runner can reach + # the upstream blob (pre-PR tests used tiktoken-go's stock network + # loader and passed), so a direct download is small and fast. As a + # fallback, copy from the infiniflow/ragflow_deps image only if it is + # already cached locally (we do NOT pull the whole image just for one + # 1.6MB file). If neither works, the loader still fails loudly (no + # silent 0-token degradation). The loader itself stays offline. + if [ ! -f ragflow_deps/cl100k_base.tiktoken ]; then + mkdir -p ragflow_deps + if command -v curl >/dev/null 2>&1; then + if curl -fsSL -o ragflow_deps/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken; then + echo "tiktoken: cl100k table fetched via curl from openai blob" + else + echo "tiktoken: curl fetch failed (network/GFW); will try cached image" + fi + fi + if [ ! -f ragflow_deps/cl100k_base.tiktoken ] && docker image inspect infiniflow/ragflow_deps:latest >/dev/null 2>&1; then + CID=$(docker create infiniflow/ragflow_deps:latest true) || true + if [ -n "${CID:-}" ]; then + if docker cp "$CID":/cl100k_base.tiktoken ragflow_deps/cl100k_base.tiktoken 2>/dev/null; then + echo "tiktoken: cl100k table copied from cached infiniflow/ragflow_deps image" + fi + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + fi + fi + if [ -f ragflow_deps/cl100k_base.tiktoken ]; then + echo "tiktoken: cl100k table provisioned ($(wc -c < ragflow_deps/cl100k_base.tiktoken) bytes)" + else + echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" + fi + PKGS=$(go list ./... 2>/dev/null \ | grep -v '/internal/storage$' \ | grep -v '/internal/agent$' \ - | grep -v '/internal/tokenizer$' \ | grep -v '/internal/handler$' || true) if [ -z "$PKGS" ]; then ./build.sh --test @@ -640,13 +679,52 @@ jobs: # # Excludes packages whose tests fail for environmental reasons # unrelated to the diff: - # - internal/tokenizer: tests need /usr/share/infinity/resource - # dict files, only mounted inside the docker builder, not - # in the Go test environment. + # - internal/tokenizer is split: the pure-Go BPE loader tests + # (bpe_loader_test.go) run in the default tier; the C++ binding / + # dict-dependent tests (tokenizer_test.go, + # tokenizer_concurrent_test.go) and the on-disk anchor test + # (bpe_loader_anchors_test.go) are tagged `manual` and need the + # docker builder's /usr/share/infinity/resource, so they stay out + # of the default run. run: | set -euo pipefail - PKGS=$(go list ./... 2>/dev/null \ - | grep -v '/internal/tokenizer$' || true) + + # Provide the cl100k BPE table for the offline tokenizer loader. + # The loader reads /ragflow_deps/cl100k_base.tiktoken (or + # TIKTOKEN_CACHE_DIR). Fetch it directly: this CI runner can reach + # the upstream blob (pre-PR tests used tiktoken-go's stock network + # loader and passed), so a direct download is small and fast. As a + # fallback, copy from the infiniflow/ragflow_deps image only if it is + # already cached locally (we do NOT pull the whole image just for one + # 1.6MB file). If neither works, the loader still fails loudly (no + # silent 0-token degradation). The loader itself stays offline. + if [ ! -f ragflow_deps/cl100k_base.tiktoken ]; then + mkdir -p ragflow_deps + if command -v curl >/dev/null 2>&1; then + if curl -fsSL -o ragflow_deps/cl100k_base.tiktoken \ + https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken; then + echo "tiktoken: cl100k table fetched via curl from openai blob" + else + echo "tiktoken: curl fetch failed (network/GFW); will try cached image" + fi + fi + if [ ! -f ragflow_deps/cl100k_base.tiktoken ] && docker image inspect infiniflow/ragflow_deps:latest >/dev/null 2>&1; then + CID=$(docker create infiniflow/ragflow_deps:latest true) || true + if [ -n "${CID:-}" ]; then + if docker cp "$CID":/cl100k_base.tiktoken ragflow_deps/cl100k_base.tiktoken 2>/dev/null; then + echo "tiktoken: cl100k table copied from cached infiniflow/ragflow_deps image" + fi + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + fi + fi + if [ -f ragflow_deps/cl100k_base.tiktoken ]; then + echo "tiktoken: cl100k table provisioned ($(wc -c < ragflow_deps/cl100k_base.tiktoken) bytes)" + else + echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" + fi + + PKGS=$(go list ./... 2>/dev/null || true) if [ -z "$PKGS" ]; then ./build.sh --test else diff --git a/internal/tokenizer/bpe_loader.go b/internal/tokenizer/bpe_loader.go new file mode 100644 index 0000000000..396a4e5732 --- /dev/null +++ b/internal/tokenizer/bpe_loader.go @@ -0,0 +1,216 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tokenizer + +// Offline BPE table loading for tiktoken. +// +// RAGFlow ships the cl100k_base table on disk (Dockerfile drops it into the +// working directory under its sha1 name; download_deps.py writes it to +// ragflow_deps/). tiktoken-go's stock loader instead downloads it over HTTP and +// relies on TIKTOKEN_CACHE_DIR, which the Go server never inherits, so a +// missing table degrades every token count to 0. This loader resolves the +// table from disk only: it performs no network I/O, and when nothing is found +// it reports every path it tried. + +import ( + "crypto/sha1" + "encoding/base64" + "fmt" + "os" + "path" + "path/filepath" + "strconv" + "strings" + + "ragflow/internal/common" + + "github.com/pkoukk/tiktoken-go" +) + +func init() { + tiktoken.SetBpeLoader(localBpeLoader{}) +} + +// expectedBpeHashes maps a tiktoken table URL to the SHA-1 of its canonical +// on-disk content. We only ship cl100k_base today; entries here let the loader +// reject a corrupt or tampered file instead of trusting it. Unknown URLs are +// loaded without a digest check (defense-in-depth, not a hard gate). +// +// NOTE: this is the digest of the file *contents*, not the tiktoken cache +// filename. tiktoken-go names its cached file by sha1(bpeURL) +// (223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7 for +// cl100k_base); that value identifies the path, while the value below verifies +// the bytes we actually load. Compute it from the table shipped by +// ragflow_deps/download_deps.py: `sha1sum cl100k_base.tiktoken`. +var expectedBpeHashes = map[string]string{ + "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken": "6494e42d5aad2bbb441ea9793af9e7db335c8d9c", +} + +// localBpeLoader resolves tiktoken BPE tables from the local filesystem. +type localBpeLoader struct{} + +// LoadTiktokenBpe implements tiktoken.BpeLoader. +// +// bpeURL is the upstream table URL that tiktoken-go would otherwise download; +// here it serves only to derive the file names to look for. +func (localBpeLoader) LoadTiktokenBpe(bpeURL string) (map[string]int, error) { + candidates := bpeCandidatePaths(bpeURL) + for _, candidate := range candidates { + contents, err := os.ReadFile(candidate) + if err != nil { + // Only a missing candidate is skippable; a permission or I/O + // failure must not be masked as "not found". + if os.IsNotExist(err) { + continue + } + return nil, fmt.Errorf("reading BPE table %s: %w", candidate, err) + } + // Integrity check: for tables we ship, a digest mismatch means the + // file is corrupt or tampered with. Refuse to load it rather than + // skipping to the next candidate — a different candidate holds the + // same (wrong) content, and masking the failure would defeat the + // check. This mirrors the malformed-table path just below. + if want, ok := expectedBpeHashes[bpeURL]; ok { + if got := fmt.Sprintf("%x", sha1.Sum(contents)); got != want { + return nil, fmt.Errorf("BPE table %s digest mismatch (got %s, want %s); refusing to load a corrupt or tampered file", candidate, got, want) + } + } + + ranks, err := parseBpeTable(contents) + if err != nil { + // A file that exists but does not parse is a corrupt download or a + // name collision. Continuing to the next candidate would mask it. + return nil, fmt.Errorf("BPE table %s is malformed: %w", candidate, err) + } + return ranks, nil + } + + err := fmt.Errorf( + "no local BPE table for %s; run `uv run ragflow_deps/download_deps.py` or set TIKTOKEN_CACHE_DIR to the directory holding the table; tried: %s", + bpeURL, strings.Join(candidates, ", ")) + // Logged as well as returned: tiktoken-go propagates this to GetEncoding, + // whose error NumTokensFromString discards to keep returning 0. + common.Error("cl100k BPE table not found; every token count will be 0", err) + return nil, err +} + +// bpeCandidatePaths lists, in priority order, every local path that may hold +// the table for bpeURL. +// +// Explicit configuration wins, then the directories RAGFlow actually ships the +// table in. Both the working directory and the executable's directory are +// walked upwards: the server runs with the working directory set to the +// installation root, while `go test` runs from a package subdirectory. +func bpeCandidatePaths(bpeURL string) []string { + cacheName := fmt.Sprintf("%x", sha1.Sum([]byte(bpeURL))) + // download_deps.py stores the table under the URL's own basename. + bundledName := path.Base(bpeURL) + + var paths []string + seen := make(map[string]struct{}) + add := func(p string) { + if _, dup := seen[p]; dup { + return + } + seen[p] = struct{}{} + paths = append(paths, p) + } + + // Honour both variables tiktoken-go itself reads, so an operator who has + // already configured one keeps working. + for _, env := range []string{"TIKTOKEN_CACHE_DIR", "DATA_GYM_CACHE_DIR"} { + if dir := strings.TrimSpace(os.Getenv(env)); dir != "" { + add(filepath.Join(dir, cacheName)) + } + } + + for _, root := range searchRoots() { + // Same layout the Dockerfile creates: the table sits in the + // installation root under its sha1 name. + add(filepath.Join(root, cacheName)) + // download_deps.py writes the table into ragflow_deps/ under its + // download name; a developer checkout that has run it but never + // started the Python side only has this copy. + add(filepath.Join(root, "ragflow_deps", bundledName)) + } + + return paths +} + +// searchRoots returns the working directory and the executable's directory +// together with all of their ancestors. +func searchRoots() []string { + var roots []string + seen := make(map[string]struct{}) + for _, start := range startingDirs() { + for dir := start; ; { + if _, dup := seen[dir]; !dup { + seen[dir] = struct{}{} + roots = append(roots, dir) + } + parent := filepath.Dir(dir) + if parent == dir { + break + } + dir = parent + } + } + return roots +} + +func startingDirs() []string { + var dirs []string + if wd, err := os.Getwd(); err == nil { + dirs = append(dirs, wd) + } + if exe, err := os.Executable(); err == nil { + if resolved, err := filepath.EvalSymlinks(exe); err == nil { + exe = resolved + } + dirs = append(dirs, filepath.Dir(exe)) + } + return dirs +} + +// parseBpeTable decodes tiktoken's on-disk format: one +// " " pair per line. +func parseBpeTable(contents []byte) (map[string]int, error) { + ranks := make(map[string]int) + for i, line := range strings.Split(string(contents), "\n") { + line = strings.TrimRight(line, "\r") + if line == "" { + continue + } + token, rank, ok := strings.Cut(line, " ") + if !ok { + return nil, fmt.Errorf("line %d: expected \" \"", i+1) + } + decoded, err := base64.StdEncoding.DecodeString(token) + if err != nil { + return nil, fmt.Errorf("line %d: %w", i+1, err) + } + value, err := strconv.Atoi(rank) + if err != nil { + return nil, fmt.Errorf("line %d: %w", i+1, err) + } + ranks[string(decoded)] = value + } + if len(ranks) == 0 { + return nil, fmt.Errorf("table is empty") + } + return ranks, nil +} diff --git a/internal/tokenizer/bpe_loader_anchors_test.go b/internal/tokenizer/bpe_loader_anchors_test.go new file mode 100644 index 0000000000..d4da6cf65b --- /dev/null +++ b/internal/tokenizer/bpe_loader_anchors_test.go @@ -0,0 +1,77 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +//go:build manual + +package tokenizer + +import "testing" + +// TestNumTokensFromString_MatchesPythonAnchors pins exact counts taken from the +// Python reference suite (test/unit_test/common/test_token_utils.py:28-49) and +// from common.token_utils.num_tokens_from_string for the CJK cases. The corpus +// was later expanded to ~24 entries spanning ASCII, punctuation, digits, +// whitespace, newlines, CJK, emoji, mixed-language, and code-like strings, all +// recomputed against tiktoken's cl100k encoder. +// +// Exact values matter more than they look. NumTokensFromString swallows loader +// errors and returns 0, so an assertion of the form "> 0" passes for an empty +// string and fails to notice a dead encoder — which is precisely how the +// offline breakage stayed invisible. Pinning the numbers also catches loading a +// structurally valid but wrong table. +// +// This test needs the real cl100k_base table on disk (TIKTOKEN_CACHE_DIR, +// the Dockerfile's /ragflow/ file, or ragflow_deps/cl100k_base.tiktoken), +// so it is tagged `manual` and runs only under `build.sh --test-manual`, +// which the docker builder provisions with /usr/share/infinity/resource. +func TestNumTokensFromString_MatchesPythonAnchors(t *testing.T) { + anchors := []struct { + in string + want int + }{ + {"", 0}, + {"hello", 1}, + {"hello world", 2}, + {"hello, world!", 4}, + {"世界", 3}, + {"Hello 世界 🌍", 8}, + {"RAGFlow", 3}, + {"1234567890", 4}, + {"a b", 3}, + {"hello\nworld", 3}, + {"user@example.com", 3}, + {"https://example.com/path?x=1", 9}, + {"func main() {}", 4}, + {"aaaaaaaaaa", 2}, + {"中文字符测试", 4}, + {"🚀🔥", 6}, + {"state-of-the-art", 4}, + {`"quoted"`, 3}, + {"The quick brown fox jumps over the lazy dog.", 10}, + {"Café naïve résumé", 8}, + {"x² + y² = z²", 8}, + {"混合 English 和 中文 的 sentence。", 10}, + {"tokenization is the process of splitting text into tokens", 10}, + {"人工智能正在改变世界，这是毫无疑问的事实。", 25}, + {"SELECT * FROM users WHERE id = 1;", 10}, + {"こんにちは世界", 4}, + } + for _, tc := range anchors { + if got := NumTokensFromString(tc.in); got != tc.want { + t.Errorf("NumTokensFromString(%q) = %d, want %d", tc.in, got, tc.want) + } + } +} diff --git a/internal/tokenizer/bpe_loader_test.go b/internal/tokenizer/bpe_loader_test.go new file mode 100644 index 0000000000..324bb38aae --- /dev/null +++ b/internal/tokenizer/bpe_loader_test.go @@ -0,0 +1,201 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tokenizer + +import ( + "crypto/sha1" + "encoding/base64" + "fmt" + "os" + "path/filepath" + "strings" + "testing" +) + +const testBpeURL = "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken" + +// writeBpeTable writes a minimal well-formed BPE table. The rank values are +// arbitrary markers so a test can tell which file the loader actually read. +func writeBpeTable(t *testing.T, path string, marker int) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatalf("mkdir %s: %v", filepath.Dir(path), err) + } + line := fmt.Sprintf("%s %d\n", base64.StdEncoding.EncodeToString([]byte("hello")), marker) + if err := os.WriteFile(path, []byte(line), 0o644); err != nil { + t.Fatalf("write %s: %v", path, err) + } + // Register this synthetic table's digest so the SHA-1 integrity gate in + // bpe_loader.go accepts it. This also exercises the verification path + // instead of disabling it. + expectedBpeHashes[testBpeURL] = fmt.Sprintf("%x", sha1.Sum([]byte(line))) +} + +func cacheFileName(url string) string { + return fmt.Sprintf("%x", sha1.Sum([]byte(url))) +} + +// isolate moves the process into an empty directory and clears both cache +// environment variables, so a test sees only the files it creates itself. +// Without this the real repository checkout — which does ship the table — would +// satisfy every lookup and hide ordering bugs. +func isolate(t *testing.T) string { + t.Helper() + dir := t.TempDir() + t.Setenv("TIKTOKEN_CACHE_DIR", "") + t.Setenv("DATA_GYM_CACHE_DIR", "") + t.Chdir(dir) + return dir +} + +func TestLocalBpeLoader_ReadsTiktokenCacheDir(t *testing.T) { + isolate(t) + cache := t.TempDir() + writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 7) + t.Setenv("TIKTOKEN_CACHE_DIR", cache) + + ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err != nil { + t.Fatalf("load: %v", err) + } + if got := ranks["hello"]; got != 7 { + t.Errorf("rank from TIKTOKEN_CACHE_DIR = %d, want 7", got) + } +} + +func TestLocalBpeLoader_ReadsDataGymCacheDir(t *testing.T) { + isolate(t) + cache := t.TempDir() + writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 8) + t.Setenv("DATA_GYM_CACHE_DIR", cache) + + ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err != nil { + t.Fatalf("load: %v", err) + } + if got := ranks["hello"]; got != 8 { + t.Errorf("rank from DATA_GYM_CACHE_DIR = %d, want 8", got) + } +} + +// The production image has no TIKTOKEN_CACHE_DIR set: Dockerfile drops the +// table straight into the working directory under its sha1 name, and +// entrypoint.sh starts the Go binary from a shell, so nothing exports the +// variable that common/token_utils.py sets inside the Python process. +func TestLocalBpeLoader_ReadsSha1FileFromWorkingDirectory(t *testing.T) { + dir := isolate(t) + writeBpeTable(t, filepath.Join(dir, cacheFileName(testBpeURL)), 9) + + ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err != nil { + t.Fatalf("load: %v", err) + } + if got := ranks["hello"]; got != 9 { + t.Errorf("rank from working directory = %d, want 9", got) + } +} + +// A developer checkout that has run download_deps.py but never started the +// Python side only has the file under its download name. +func TestLocalBpeLoader_ReadsBundledVocabFromAncestor(t *testing.T) { + dir := isolate(t) + writeBpeTable(t, filepath.Join(dir, "ragflow_deps", "cl100k_base.tiktoken"), 10) + nested := filepath.Join(dir, "internal", "tokenizer") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + t.Chdir(nested) + + ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err != nil { + t.Fatalf("load: %v", err) + } + if got := ranks["hello"]; got != 10 { + t.Errorf("rank from bundled vocab = %d, want 10", got) + } +} + +func TestLocalBpeLoader_CacheDirWinsOverBundledVocab(t *testing.T) { + dir := isolate(t) + writeBpeTable(t, filepath.Join(dir, "ragflow_deps", "cl100k_base.tiktoken"), 10) + cache := t.TempDir() + writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 7) + t.Setenv("TIKTOKEN_CACHE_DIR", cache) + + ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err != nil { + t.Fatalf("load: %v", err) + } + if got := ranks["hello"]; got != 7 { + t.Errorf("rank = %d, want 7 (an explicit cache dir must win)", got) + } +} + +// The whole point of this loader is that a missing table is an error rather +// than an HTTP request. Reporting every path tried is what turns an opaque +// "all token counts are zero" deployment into a one-look diagnosis. +func TestLocalBpeLoader_MissingTableReportsCandidatesInsteadOfDownloading(t *testing.T) { + isolate(t) + + _, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err == nil { + t.Fatal("expected an error when no local table exists, got nil") + } + for _, want := range []string{cacheFileName(testBpeURL), "cl100k_base.tiktoken", "TIKTOKEN_CACHE_DIR"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error does not mention %q, so it cannot be acted on:\n%v", want, err) + } + } + if strings.Contains(err.Error(), "http") && !strings.Contains(err.Error(), testBpeURL) { + t.Errorf("error hints at a network attempt: %v", err) + } +} + +func TestLocalBpeLoader_RejectsMalformedTable(t *testing.T) { + dir := isolate(t) + path := filepath.Join(dir, cacheFileName(testBpeURL)) + if err := os.WriteFile(path, []byte("not base64 at all\n"), 0o644); err != nil { + t.Fatalf("write: %v", err) + } + + _, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err == nil { + t.Fatal("expected an error for a malformed table, got nil") + } +} + +// A candidate that exists but cannot be read as a file must surface as a read +// error rather than being skipped as "missing". LoadTiktokenBpe only continues +// past os.IsNotExist; a directory at the candidate path fails ReadFile with a +// distinct error, which this test pins to the read-error path. +func TestLocalBpeLoader_ReadErrorIsReported(t *testing.T) { + dir := isolate(t) + // A directory at the sha1-named candidate path exists but is not a + // regular file, so os.ReadFile fails with a non-IsNotExist error. + candidate := filepath.Join(dir, cacheFileName(testBpeURL)) + if err := os.Mkdir(candidate, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + + _, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL) + if err == nil { + t.Fatal("expected a read error for a non-file candidate, got nil") + } + if strings.Contains(err.Error(), "no local BPE table") { + t.Errorf("read error was masked as not-found: %v", err) + } +} diff --git a/internal/tokenizer/tokenizer_concurrent_test.go b/internal/tokenizer/tokenizer_concurrent_test.go index 175d8651b6..27e9e0c69d 100644 --- a/internal/tokenizer/tokenizer_concurrent_test.go +++ b/internal/tokenizer/tokenizer_concurrent_test.go @@ -12,7 +12,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -// +//go:build manual package tokenizer diff --git a/internal/tokenizer/tokenizer_test.go b/internal/tokenizer/tokenizer_test.go index 7670ccb59d..380805b577 100644 --- a/internal/tokenizer/tokenizer_test.go +++ b/internal/tokenizer/tokenizer_test.go @@ -12,7 +12,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -// +//go:build manual package tokenizer