diff --git a/.github/workflows/sep-tests.yml b/.github/workflows/sep-tests.yml index 283c6ed82d..3ca4a9bb45 100644 --- a/.github/workflows/sep-tests.yml +++ b/.github/workflows/sep-tests.yml @@ -484,6 +484,40 @@ jobs: echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" fi + # Provide the tokenizer assets the four non-cl100k counters load + # (ragflow_deps/huggingface.co//). Their only no-network source + # is the infiniflow/ragflow_deps image — the runtime images read them from + # exactly the same path — so this is the same "if the image is already + # cached" fallback the table above uses (again, no `docker pull`). Without + # them counter_properties_test / asset_dir_test silently drop their + # non-cl100k coverage, which is how a counter regression goes unnoticed. + TOKENIZER_ASSET_PATHS=( + BAAI/bge-m3/sentencepiece.bpe.model + BAAI/bge-large-en-v1.5/vocab.txt + Qwen/Qwen3-Embedding-0.6B/tokenizer.json + intfloat/e5-mistral-7b-instruct/tokenizer.json + ) + if 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 + for asset in "${TOKENIZER_ASSET_PATHS[@]}"; do + target="ragflow_deps/huggingface.co/$asset" + if [ -f "$target" ]; then + continue + fi + mkdir -p "$(dirname "$target")" + if docker cp "$CID:/huggingface.co/$asset" "$target" 2>/dev/null; then + echo "tokenizer: $asset copied from cached infiniflow/ragflow_deps image" + else + echo "tokenizer: $asset absent from the cached image — its counter will skip" + fi + done + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + else + echo "tokenizer: infiniflow/ragflow_deps not cached on this runner — non-cl100k counters will skip" + fi + PKGS=$(go list ./... 2>/dev/null \ | grep -v '/internal/storage$' \ | grep -v '/internal/handler$' || true) @@ -493,17 +527,20 @@ jobs: ./build.sh --test -- $PKGS fi - - name: Run ORT deps-download regression guard + - name: Run ragflow_deps regression guards if: env.API_PROXY_SCHEME == 'go' - # Pure-Python, local, needs no API keys or model downloads; guards the + # Pure-Python, local, needs no API keys or model downloads; guards (a) the # ONNX Runtime extraction logic the in-process (Go) DeepDoc backend build - # depends on (notably the version-bump path C2 hardened). Run it in - # isolation so the test/testcases conftest (which requires + # depends on (notably the version-bump path C2 hardened) and (b) the + # tokenizer-asset chain: the downloader's asset list, the three runtime + # Dockerfiles' copy loops and the loaders' SHA-1 pins have to agree, or the + # Go counters silently degrade to the calibrated cl100k estimate. Run them + # in isolation so the test/testcases conftest (which requires # SILICONFLOW_API_KEY) cannot fail the collection. run: | set -euo pipefail uv sync --python 3.13 --group test --frozen - uv run pytest ragflow_deps/test_download_go_deps.py -q --durations=100 + uv run pytest ragflow_deps/test_download_go_deps.py ragflow_deps/test_tokenizer_assets.py -q --durations=100 - name: Build ragflow:nightly run: | @@ -1181,6 +1218,40 @@ jobs: echo "tiktoken: cl100k table NOT provisioned — offline loader will fail loudly" fi + # Provide the tokenizer assets the four non-cl100k counters load + # (ragflow_deps/huggingface.co//). Their only no-network source + # is the infiniflow/ragflow_deps image — the runtime images read them from + # exactly the same path — so this is the same "if the image is already + # cached" fallback the table above uses (again, no `docker pull`). Without + # them counter_properties_test / asset_dir_test silently drop their + # non-cl100k coverage, which is how a counter regression goes unnoticed. + TOKENIZER_ASSET_PATHS=( + BAAI/bge-m3/sentencepiece.bpe.model + BAAI/bge-large-en-v1.5/vocab.txt + Qwen/Qwen3-Embedding-0.6B/tokenizer.json + intfloat/e5-mistral-7b-instruct/tokenizer.json + ) + if 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 + for asset in "${TOKENIZER_ASSET_PATHS[@]}"; do + target="ragflow_deps/huggingface.co/$asset" + if [ -f "$target" ]; then + continue + fi + mkdir -p "$(dirname "$target")" + if docker cp "$CID:/huggingface.co/$asset" "$target" 2>/dev/null; then + echo "tokenizer: $asset copied from cached infiniflow/ragflow_deps image" + else + echo "tokenizer: $asset absent from the cached image — its counter will skip" + fi + done + docker rm -f "$CID" >/dev/null 2>&1 || true + fi + else + echo "tokenizer: infiniflow/ragflow_deps not cached on this runner — non-cl100k counters will skip" + fi + PKGS=$(go list ./... 2>/dev/null \ | grep -v '/internal/storage$' \ | grep -v '/internal/handler$' || true) @@ -1190,17 +1261,20 @@ jobs: ./build.sh --test -- $PKGS fi - - name: Run ORT deps-download regression guard + - name: Run ragflow_deps regression guards if: env.API_PROXY_SCHEME == 'go' - # Pure-Python, local, needs no API keys or model downloads; guards the + # Pure-Python, local, needs no API keys or model downloads; guards (a) the # ONNX Runtime extraction logic the in-process (Go) DeepDoc backend build - # depends on (notably the version-bump path C2 hardened). Run it in - # isolation so the test/testcases conftest (which requires + # depends on (notably the version-bump path C2 hardened) and (b) the + # tokenizer-asset chain: the downloader's asset list, the three runtime + # Dockerfiles' copy loops and the loaders' SHA-1 pins have to agree, or the + # Go counters silently degrade to the calibrated cl100k estimate. Run them + # in isolation so the test/testcases conftest (which requires # SILICONFLOW_API_KEY) cannot fail the collection. run: | set -euo pipefail uv sync --python 3.13 --group test --frozen - uv run pytest ragflow_deps/test_download_go_deps.py -q --durations=100 + uv run pytest ragflow_deps/test_download_go_deps.py ragflow_deps/test_tokenizer_assets.py -q --durations=100 - name: Build ragflow:nightly run: | diff --git a/Dockerfile b/Dockerfile index b6bee7c474..9d0fca24f9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -27,6 +27,29 @@ RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/,target=/deps cp /deps/tika-server-standard-3.3.0.jar /deps/tika-server-standard-3.3.0.jar.md5 /ragflow/ && \ cp /deps/cl100k_base.tiktoken /ragflow/9b5ad71b2ce5302211f9c61530b329a4922fc6a4 +# Embedding tokenizer assets (internal/tokenizer/embedding_token_limits.md). The Go +# counters load them from ragflow_deps/huggingface.co//; an image without them +# counts every tagged model with the calibrated cl100k estimate, which is the less precise +# path these counters exist to replace (cl100k under-counts XLM-R on some content, and an +# under-count is what makes a provider answer 400) - so a missing asset FAILS THE BUILD +# instead of shipping a degraded counter nobody notices. +# The tokenizer.json files that download_deps.py fetches as cross-check oracles are +# test-only and deliberately not shipped here. +RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/huggingface.co,target=/huggingface.co \ + for asset in \ + BAAI/bge-m3/sentencepiece.bpe.model \ + BAAI/bge-large-en-v1.5/vocab.txt \ + Qwen/Qwen3-Embedding-0.6B/tokenizer.json \ + intfloat/e5-mistral-7b-instruct/tokenizer.json ; do \ + if [ -f "/huggingface.co/$asset" ]; then \ + mkdir -p "/ragflow/ragflow_deps/huggingface.co/$(dirname "$asset")" && \ + cp "/huggingface.co/$asset" "/ragflow/ragflow_deps/huggingface.co/$asset" ; \ + else \ + echo "ERROR: tokenizer asset $asset is missing from the infiniflow/ragflow_deps image; this image would count with the calibrated estimate instead of the model's own tokenizer" >&2 ; \ + exit 1 ; \ + fi ; \ + done + ENV TIKA_SERVER_JAR="file:///ragflow/tika-server-standard-3.3.0.jar" ENV DEBIAN_FRONTEND=noninteractive diff --git a/Dockerfile_base b/Dockerfile_base index 6c3e5ce8cb..e67bfd6bb0 100644 --- a/Dockerfile_base +++ b/Dockerfile_base @@ -24,6 +24,27 @@ RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/,target=/deps cp /deps/tika-server-standard-3.3.0.jar /deps/tika-server-standard-3.3.0.jar.md5 /ragflow/ && \ cp /deps/cl100k_base.tiktoken /ragflow/9b5ad71b2ce5302211f9c61530b329a4922fc6a4 +# Embedding tokenizer assets (internal/tokenizer/embedding_token_limits.md). The Go +# counters load them from ragflow_deps/huggingface.co//; an image without them +# counts every tagged model with the calibrated cl100k estimate, the less precise path +# these counters exist to replace - so a missing asset FAILS THE BUILD instead of shipping +# a degraded counter nobody notices. The tokenizer.json files that download_deps.py +# fetches as cross-check oracles are test-only and not shipped here. +RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/huggingface.co,target=/huggingface.co \ + for asset in \ + BAAI/bge-m3/sentencepiece.bpe.model \ + BAAI/bge-large-en-v1.5/vocab.txt \ + Qwen/Qwen3-Embedding-0.6B/tokenizer.json \ + intfloat/e5-mistral-7b-instruct/tokenizer.json ; do \ + if [ -f "/huggingface.co/$asset" ]; then \ + mkdir -p "/ragflow/ragflow_deps/huggingface.co/$(dirname "$asset")" && \ + cp "/huggingface.co/$asset" "/ragflow/ragflow_deps/huggingface.co/$asset" ; \ + else \ + echo "ERROR: tokenizer asset $asset is missing from the infiniflow/ragflow_deps image; this image would count with the calibrated estimate instead of the model's own tokenizer" >&2 ; \ + exit 1 ; \ + fi ; \ + done + ENV TIKA_SERVER_JAR="file:///ragflow/tika-server-standard-3.3.0.jar" ENV DEBIAN_FRONTEND=noninteractive diff --git a/Dockerfile_go b/Dockerfile_go index 6632a1907e..1b28061439 100644 --- a/Dockerfile_go +++ b/Dockerfile_go @@ -50,6 +50,27 @@ RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/cl100k_base.ti mkdir -p /ragflow/ragflow_deps && \ cp /tmp/cl100k_base.tiktoken /ragflow/ragflow_deps/cl100k_base.tiktoken +# Embedding tokenizer assets for the Go counters (internal/tokenizer). Same story as +# the cl100k table above: the deps image carries them and the Go image has to put them +# where the loaders search (/ragflow_deps/huggingface.co//). A +# missing asset FAILS THE BUILD here rather than shipping an image where every tagged +# model silently counts with the calibrated cl100k estimate. The tokenizer.json files +# that exist only as cross-check oracles are deliberately not shipped. +RUN --mount=type=bind,from=infiniflow/ragflow_deps:latest,source=/huggingface.co,target=/huggingface.co \ + for asset in \ + BAAI/bge-m3/sentencepiece.bpe.model \ + BAAI/bge-large-en-v1.5/vocab.txt \ + Qwen/Qwen3-Embedding-0.6B/tokenizer.json \ + intfloat/e5-mistral-7b-instruct/tokenizer.json ; do \ + if [ -f "/huggingface.co/$asset" ]; then \ + mkdir -p "/ragflow/ragflow_deps/huggingface.co/$(dirname "$asset")" && \ + cp "/huggingface.co/$asset" "/ragflow/ragflow_deps/huggingface.co/$asset" ; \ + else \ + echo "ERROR: tokenizer asset $asset is missing from the infiniflow/ragflow_deps image; this image would count with the calibrated estimate instead of the model's own tokenizer" >&2 ; \ + exit 1 ; \ + fi ; \ + done + ENV DEBIAN_FRONTEND=noninteractive # Setup apt diff --git a/build.sh b/build.sh index 0d8fb48666..610af90740 100755 --- a/build.sh +++ b/build.sh @@ -69,6 +69,68 @@ _seed_from_system() { return 1 } +# ── cl100k_base BPE table ───────────────────────────────────────────── +# internal/tokenizer loads this table from disk only (no network) and PANICS in +# NumTokensFromString / TrimContentToTokenLimit when it is absent: a deployment +# that forgot the file must not silently zero every token count (see +# internal/tokenizer/bpe_loader.go and failfast_test.go). CI provisions it in a +# workflow step before running the Go tests; build.sh does the same so the +# documented local test commands are self-contained. Order mirrors CI: system +# pre-seed first, then the network. A failure here is deliberately NOT fatal — +# the loader stays the final gate, and its panic message names the file. +CL100K_TABLE_URL="https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken" +CL100K_TABLE_RELPATH="ragflow_deps/cl100k_base.tiktoken" +SYSTEM_DEPS_TOKENIZER="/opt/ragflow_deps" + +# Provision the cl100k_base BPE table into ragflow_deps/ (no-op if present). +# Returns 0 when the table is in place, 1 when it could not be provisioned. +_ensure_cl100k_table() { + local dest="${PROJECT_ROOT}/${CL100K_TABLE_RELPATH}" + local sys_copy="${SYSTEM_DEPS_TOKENIZER}/cl100k_base.tiktoken" + + echo "check if the cl100k_base BPE table exists at ${dest}" + + if [ -s "$dest" ]; then + echo " cl100k_base.tiktoken → ${dest} (already present)" + return 0 + fi + + if [ -s "$sys_copy" ]; then + mkdir -p "$(dirname "$dest")" + cp "$sys_copy" "$dest" + echo " cl100k_base.tiktoken → ${dest} (system)" + return 0 + fi + + # Download to a temp file and move it into place only after a successful + # transfer: an interrupted download would otherwise leave a truncated table + # behind, which the loader rejects on digest. + local tmp="${dest}.download" + mkdir -p "$(dirname "$dest")" + if command -v curl >/dev/null 2>&1; then + if curl -fsSL -o "$tmp" "$CL100K_TABLE_URL"; then + mv "$tmp" "$dest" + echo " cl100k_base.tiktoken → ${dest} (downloaded)" + return 0 + fi + elif command -v wget >/dev/null 2>&1; then + if wget -q -O "$tmp" "$CL100K_TABLE_URL"; then + mv "$tmp" "$dest" + echo " cl100k_base.tiktoken → ${dest} (downloaded)" + return 0 + fi + fi + rm -f "$tmp" + + echo -e " ${YELLOW}cl100k_base BPE table NOT provisioned${NC} (offline, or the network is blocked)" + echo " internal/tokenizer panics on a missing table by design, so every test that" + echo " counts tokens will fail loudly. Recover with any of:" + echo " - re-run this command (a transient network error usually clears);" + echo " - curl -fsSL -o ${CL100K_TABLE_RELPATH} ${CL100K_TABLE_URL}" + echo " - or point TIKTOKEN_CACHE_DIR at a directory holding the table." + return 1 +} + echo -e "${GREEN}=== RAGFlow Go Server Build Script ===${NC}" # Function to print section headers @@ -727,6 +789,11 @@ setup_cgo_env() { run_go_tests() { print_section "Running Go tests" + # The tokenizer panics instead of counting 0 when the cl100k table is + # missing, so provision it before the tests, as CI does. Non-fatal: without + # it only the token-counting tests fail, and they say why. + _ensure_cl100k_table || true + cd "$PROJECT_ROOT" setup_cgo_env @@ -816,6 +883,10 @@ run_go_tests_tagged() { local tags="$1"; shift print_section "Running Go tests (tags: ${tags})" + # Same cl100k table requirement as the unit tier (integration/e2e runs count + # tokens too); non-fatal, see _ensure_cl100k_table. + _ensure_cl100k_table || true + cd "$PROJECT_ROOT" setup_cgo_env diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 94e4439b6a..797571e660 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -499,7 +499,11 @@ func main() { logLevel = "debug" } - if err = common.InitLogger(logLevel, common.FileOutput{Filename: logFileName, Path: "logs"}, serverName); err != nil { + // Temporary pre-config logger: STDOUT ONLY (empty FileOutput). The port + // is not known yet, so a file here would be an orphaned log (e.g. + // logs/api_server.log next to the real logs/api_server_9384.log); the + // real file sink is attached by the post-config re-initialization below. + if err = common.InitLogger(logLevel, common.FileOutput{}, serverName); err != nil { panic("failed to initialize logger: " + err.Error()) } @@ -517,9 +521,15 @@ func main() { globalConfig := server.GetConfig() // override default port if provided + // NOTE: this switch must stay side-effect-free on the LOG (no + // registerNativeDeepDoc here): it runs while only the temporary + // stdout-only logger exists, so anything it logs is lost from the file. + // Side effects that log (DeepDoc registration) move below, after the + // real file-backed logger is up. + needNativeDeepDoc := false switch *arguments.mode { case "api": - registerNativeDeepDoc() + needNativeDeepDoc = true apiServerConfig := globalConfig.GetAPIServerConfig() port := apiServerConfig.HTTPPort if arguments.port != nil { @@ -540,7 +550,7 @@ func main() { serverName = fmt.Sprintf("admin_server_%d", port) } case "ingestor": - registerNativeDeepDoc() + needNativeDeepDoc = true if serverName == "" { uuid := utility.GenerateUUID() serverName = fmt.Sprintf("ingestor_server_%s", uuid) @@ -590,6 +600,14 @@ func main() { common.Error("Failed to reinitialize logger with configured level", err) } + // Wire the in-process DeepDoc backend only after the REAL file-backed + // logger exists: its registration lines (and the Fatal abort on a missing + // backend) must land in the run's log file, not in the pre-config + // stdout-only window. + if needNativeDeepDoc { + registerNativeDeepDoc() + } + // Print all configuration settings common.Info(fmt.Sprintf("Starting %s server: %s, mode: %s", *arguments.mode, serverName, globalConfig.GetMode())) server.PrintAll() @@ -908,6 +926,7 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg if err := tokenizer.InitCL100KEncoder(); err != nil { common.Fatal("Failed to initialize cl100k_base tokenizer", zap.Error(err)) } + logTokenizerCounters() // The dataset-level post-processing consumer cluster (§11) is owned and run by // the Ingestor: it is started inside ingestor.Start() and joined inside @@ -1118,6 +1137,11 @@ func startServer(ctx context.Context, args *serverArgs) error { docEngine := engine.Get() documentDAO := dao.NewDocumentDAO() retrievalEnhancer := retrievalbridge.NewEnhancer(docEngine, metadataService) + // Keep the concrete adapter: it is both the canvas/agent-tool backend and the + // target of the agentic-RAG bridge wired below. The bridge must hold it + // directly (not look it up in the registry) because the registry is shared — + // agenttool.SetRetrievalService and runtime.SetRetrievalService write the same + // singleton, so a bridge that resolved its target per call would find itself. retrievalAdapter := agenttool.NewNLPRetrievalAdapterFromDeps( docEngine, documentDAO, @@ -1142,11 +1166,11 @@ func startServer(ctx context.Context, args *serverArgs) error { common.Info("agent: retrieval service adapter installed") // Wire the agentic-RAG runtime as the Go chat pipeline's evidence engine - // (internal/service/chat_pipeline.retrieveViaHarness): it searches through the - // runtime retrieval singleton below and runs each request on a model resolved - // from the caller's ModelID. Activating the agentic loop here enables the full - // planner/SCA path for reasoning chats. - runtime.SetRetrievalService(retrievalbridge.NewRuntimeAdapter()) + // (internal/service/chat_pipeline.retrieveViaHarness): it runs each request on a + // model resolved from the caller's ModelID and searches through the adapter + // above. Activating the agentic loop here enables the full planner/SCA path for + // reasoning chats. + runtime.SetRetrievalService(retrievalbridge.NewRuntimeAdapter(retrievalAdapter)) agentic_rag.SetAgenticLoop(agentic_rag.NewAgenticLoop()) service.SetHarnessRetriever(func(ctx context.Context, req service.HarnessRequest) (service.HarnessResult, error) { // Resolve the tenant's actual chat model, mirroring Python where RAGTools @@ -1793,6 +1817,28 @@ func registerNativeDeepDoc() { zap.Int("gomaxprocs", goruntime.GOMAXPROCS(0))) } +// logTokenizerCounters reports, once at startup, which embedding tokenizers this process +// can count with. An unavailable counter is not fatal here: the process still starts, and +// untagged models keep counting with the calibrated estimate by design. But the models that +// declare it cannot be ingested - the embedder refuses to substitute the calibrated count - +// so this report is what tells an operator which asset to restore before documents fail. +func logTokenizerCounters() { + var available, unavailable []string + for _, status := range tokenizer.CounterStatuses() { + if status.Available { + available = append(available, status.ID) + continue + } + unavailable = append(unavailable, status.ID) + } + common.Info("embedding tokenizer counters", zap.Strings("available", available)) + if len(unavailable) > 0 { + common.Warn("embedding tokenizers unavailable; models that declare them cannot be ingested until the asset is restored", + zap.Strings("unavailable", unavailable), + zap.String("hint", "run `uv run ragflow_deps/download_go_deps.py`, or set "+common.EnvModelAssetsDir+" to a directory holding them")) + } +} + // resolveDeepDocModelDir picks the model directory: the explicit DEEPDOC_MODEL_DIR // env, else the RAGFlow default (rag/res/deepdoc, mirroring deepdoc_server.py), // else the snapshot fetched by ragflow_deps/download_deps.py. The first @@ -1802,10 +1848,13 @@ func resolveDeepDocModelDir() string { return v } wd, _ := os.Getwd() - candidates := []string{ + // MODEL_ASSETS_DIR first: the shared model-asset root keeps this layout too, so one + // variable can point at the DeepDoc weights as well as the embedding tokenizers. + candidates := append([]string(nil), common.ModelAssetCandidates("huggingface.co/InfiniFlow/deepdoc")...) + candidates = append(candidates, filepath.Join(wd, "rag", "res", "deepdoc"), filepath.Join(wd, "huggingface.co", "InfiniFlow", "deepdoc"), - } + ) for _, c := range candidates { if dirHasModels(c) { return c diff --git a/conf/all_models.json b/conf/all_models.json index e9c5414b99..297c9df807 100644 --- a/conf/all_models.json +++ b/conf/all_models.json @@ -1075,6 +1075,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 512, "max_output": 512, "max_batch_size": 2048 @@ -1084,6 +1085,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 1024, "context_length": 512, "max_output": 512, @@ -1105,6 +1107,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 8192, "max_output": 8192, "max_batch_size": 2048 @@ -1117,6 +1120,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "max_dimension": 1024, "dimensions": [], "context_length": 8192, @@ -1132,6 +1136,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "context_length": 8192, "max_output": 8192, "max_batch_size": 2048 @@ -2098,6 +2103,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_batch_size": 2048 }, { @@ -11940,6 +11946,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 512, "max_output": 512, "max_batch_size": 2048 @@ -11953,6 +11960,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 512, "max_output": 512, "max_batch_size": 2048 @@ -11966,6 +11974,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "context_length": 512, "max_output": 512, "max_batch_size": 2048 @@ -11980,6 +11989,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "context_length": 512, "max_output": 512, "max_batch_size": 2048 @@ -13468,6 +13478,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "max_dimension": 1024, "dimensions": [ 32, @@ -13490,6 +13501,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "max_dimension": 1024, "dimensions": [ 32, @@ -13512,6 +13524,7 @@ "model_types": [ "embedding" ], + "tokenizer": "xlmr-spm", "max_dimension": 1024, "dimensions": [ 32, @@ -13639,6 +13652,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 768, "dimensions": [ 768 @@ -13655,6 +13669,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 768, "dimensions": [ 768 @@ -13671,6 +13686,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 768, "dimensions": [ 768 @@ -13687,6 +13703,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 768, "dimensions": [ 768 @@ -13703,6 +13720,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 768, "dimensions": [ 768 @@ -13719,6 +13737,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "max_dimension": 512, "dimensions": [ 512 @@ -24130,6 +24149,8 @@ "model_types": [ "embedding" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "max_dimension": 3072, "dimensions": [], "max_batch_size": 2048 @@ -24140,6 +24161,8 @@ "model_types": [ "embedding" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "max_dimension": 1536, "dimensions": [], "max_batch_size": 2048 @@ -24150,6 +24173,8 @@ "model_types": [ "embedding" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "max_dimension": 1536, "dimensions": [], "max_batch_size": 2048 @@ -25486,6 +25511,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "context_length": 32768, "max_output": 32768, "max_batch_size": 2048 @@ -25499,6 +25525,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "context_length": 32768, "max_output": 32768, "max_batch_size": 2048 @@ -25512,6 +25539,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "context_length": 32768, "max_output": 32768, "max_batch_size": 2048 @@ -27953,6 +27981,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 1024, "dimensions": [], "context_length": 32768, @@ -27967,6 +27996,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 1024, "dimensions": [], "context_length": 32768, @@ -27981,6 +28011,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 2560, "dimensions": [], "context_length": 32768, @@ -27995,6 +28026,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 2560, "dimensions": [], "context_length": 32768, @@ -28009,6 +28041,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 4096, "dimensions": [], "context_length": 32768, @@ -28023,6 +28056,7 @@ "model_types": [ "embedding" ], + "tokenizer": "qwen-bpe", "max_dimension": 4096, "dimensions": [], "context_length": 32768, @@ -36864,6 +36898,8 @@ "embedding", "chat" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "alias": [], "max_batch_size": 2048 }, @@ -36873,6 +36909,8 @@ "embedding", "chat" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "alias": [], "max_batch_size": 2048 }, @@ -36882,6 +36920,8 @@ "embedding", "chat" ], + "tokenizer": "cl100k_base", + "context_length": 8191, "alias": [], "max_batch_size": 2048 }, @@ -36936,6 +36976,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 512, "max_output": 512, "max_batch_size": 1 @@ -36949,6 +36990,7 @@ "model_types": [ "embedding" ], + "tokenizer": "bert-wordpiece", "context_length": 512, "max_output": 512, "max_batch_size": 1 diff --git a/conf/mapping.json b/conf/mapping.json index 495f7c7763..1d8a8e65a5 100644 --- a/conf/mapping.json +++ b/conf/mapping.json @@ -133,9 +133,11 @@ "match_pattern": "regex", "match": "^.*_(with_weight|list)$", "mapping": { - "type": "text", - "index": "false", - "store": true + "type": "keyword", + "index": true, + "store": true, + "ignore_above": 10000, + "similarity": "boolean" } } }, @@ -209,4 +211,4 @@ } ] } -} \ No newline at end of file +} diff --git a/docker/docker-compose-base.yml b/docker/docker-compose-base.yml index 5dce3477c1..20961d8b86 100644 --- a/docker/docker-compose-base.yml +++ b/docker/docker-compose-base.yml @@ -223,14 +223,43 @@ services: env_file: .env environment: - MYSQL_ROOT_PASSWORD=${MYSQL_PASSWORD} + # Ingestion throughput tuning, OPT-IN. The ingestion sink commits per progress + # event, so MySQL's own defaults (innodb_flush_log_at_trx_commit=1 + + # sync_binlog=1 on a 128MB buffer pool) spent 68ms of CPU and two fsyncs on + # every commit: 211k commits had burned 14.4k seconds of MySQL CPU and capped + # ingestion at ~23 documents/minute, with mysqld the busiest process on the + # box. Relaxing the two settings brings the same commit down to 1.6ms and + # ingestion to ~56 documents/minute (measured, same load). + # + # They are NOT the default because they apply to every write in this instance, + # not only to the ingestion bookkeeping they were measured for: after a + # host-level crash (not a mysqld restart) a client can lose transactions it + # already observed as committed. A deployment that treats the whole database as + # rebuildable opts in through docker/.env: + # MYSQL_FLUSH_LOG_AT_TRX_COMMIT=2 + # MYSQL_SYNC_BINLOG=0 command: - --max_connections=1000 - --character-set-server=utf8mb4 - --collation-server=utf8mb4_unicode_ci - --default-authentication-plugin=mysql_native_password - --tls_version="TLSv1.2,TLSv1.3" - --init-file /data/application/init.sql - --binlog_expire_logs_seconds=604800 + - --max_connections=1000 + - --character-set-server=utf8mb4 + - --collation-server=utf8mb4_unicode_ci + - --default-authentication-plugin=mysql_native_password + - --tls_version=TLSv1.2,TLSv1.3 + - --init-file=/data/application/init.sql + - --binlog_expire_logs_seconds=604800 + - --innodb-flush-log-at-trx-commit=${MYSQL_FLUSH_LOG_AT_TRX_COMMIT:-1} + - --sync-binlog=${MYSQL_SYNC_BINLOG:-1} + # 1G, not 4G. The database is 3.6G but its hot working set is under 1G + # (document/file/file2document/ingestion_task_log sum to ~880M), and 2.66G + # of the 3.6G is pipeline_operation_log, a log table that is written but + # not re-read. Ingestion is embedding-bound, not database-bound: the + # per-stage database cost is 0.03-0.10s, unchanged when the pool was + # shrunk from 4G to 1G online. The 4G pool simply wasted 3G of a box with + # no swap. Restart mysqld after changing this (an online shrink returns + # only part of the memory to the OS). + - --innodb-buffer-pool-size=1G + - --innodb-io-capacity=2000 + - --innodb-io-capacity-max=4000 + - --innodb-redo-log-capacity=1G ports: - ${EXPOSE_MYSQL_PORT}:3306 volumes: @@ -271,7 +300,7 @@ services: # services use a separate `kvrocks` service (defined below) so the Python # path is left unchanged by the Valkey -> apache/kvrocks migration. image: valkey/valkey:8 - command: ["redis-server", "--requirepass", "${REDIS_PASSWORD}", "--maxmemory", "128mb", "--maxmemory-policy", "volatile-lru"] + command: ["redis-server", "--requirepass", "${REDIS_PASSWORD}", "--maxmemory", "1gb", "--maxmemory-policy", "volatile-lru"] env_file: .env ports: - ${REDIS_PORT}:6379 diff --git a/go.mod b/go.mod index a405d10d50..b777148565 100644 --- a/go.mod +++ b/go.mod @@ -1,8 +1,9 @@ module ragflow -go 1.26.4 +go 1.27 require ( + cloud.google.com/go v0.123.0 cloud.google.com/go/bigquery v1.82.0 cloud.google.com/go/storage v1.63.0 github.com/AkmalOt/gomsg v0.0.0-20260407083308-985c3a1a76b7 @@ -11,6 +12,7 @@ require ( github.com/DATA-DOG/go-sqlmock v1.5.2 github.com/JohannesKaufmann/html-to-markdown v1.6.0 github.com/LuxorLabs/tenki-sdk-go/sandbox v0.7.0 + github.com/PuerkitoBio/goquery v1.9.2 github.com/alibabacloud-go/agentrun-20250910/v5 v5.8.4 github.com/alibabacloud-go/darabonba-openapi/v2 v2.2.1 github.com/alicebob/miniredis/v2 v2.38.0 @@ -23,10 +25,13 @@ require ( github.com/aws/aws-sdk-go-v2/service/sts v1.41.8 github.com/aws/smithy-go v1.24.2 github.com/browserbase/stagehand-go/v3 v3.21.0 + github.com/bytedance/sonic v1.15.1 github.com/cenkalti/backoff/v5 v5.0.3 github.com/cespare/xxhash/v2 v2.3.0 - github.com/cloudwego/eino v0.9.14 + github.com/cloudwego/eino v0.10.0-alpha.33 github.com/denisenkom/go-mssqldb v0.12.3 + github.com/dop251/goja v0.0.0-20260820211235-95a30dcd3fa5 + github.com/eino-contrib/jsonschema v1.0.3 github.com/elastic/go-elasticsearch/v8 v8.19.1 github.com/emersion/go-imap/v2 v2.0.0-beta.8 github.com/emersion/go-message v0.18.2 @@ -40,11 +45,14 @@ require ( github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 github.com/infiniflow/infinity-go-sdk v0.0.0-00010101000000-000000000000 + github.com/infiniflow/onnxruntime_go v1.23.1 + github.com/jackc/pgx/v5 v5.10.0 github.com/kaptinlin/jsonrepair v0.4.8 github.com/larksuite/oapi-sdk-go/v3 v3.9.9 github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728 github.com/lib/pq v1.10.9 github.com/minio/minio-go/v7 v7.0.99 + github.com/modelcontextprotocol/go-sdk v1.8.0 github.com/nats-io/nats-server/v2 v2.14.3 github.com/nats-io/nats.go v1.52.0 github.com/nikolalohinski/gonja v1.5.3 @@ -70,6 +78,7 @@ require ( go.uber.org/zap v1.27.1 golang.org/x/crypto v0.54.0 golang.org/x/image v0.45.0 + golang.org/x/mod v0.38.0 golang.org/x/net v0.57.0 golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.22.0 @@ -84,29 +93,9 @@ require ( gorm.io/gorm v1.25.7 ) -require ( - github.com/PuerkitoBio/goquery v1.9.2 // indirect - github.com/andybalholm/cascadia v1.3.2 // indirect - github.com/apache/arrow/go/v15 v15.0.2 // indirect - github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect - github.com/goccy/go-json v0.10.5 // indirect - github.com/google/flatbuffers v23.5.26+incompatible // indirect - github.com/google/jsonschema-go v0.4.3 // indirect - github.com/json-iterator/go v1.1.12 // indirect - github.com/pierrec/lz4/v4 v4.1.18 // indirect - github.com/segmentio/asm v1.1.3 // indirect - github.com/segmentio/encoding v0.5.4 // indirect - github.com/yosida95/uritemplate/v3 v3.0.2 // indirect - golang.org/x/mod v0.38.0 // indirect - golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 // indirect - golang.org/x/tools v0.48.0 // indirect - golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect -) - require ( buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.11-20260415201107-50325440f8f2.1 // indirect cel.dev/expr v0.25.2 // indirect - cloud.google.com/go v0.123.0 cloud.google.com/go/auth v0.21.0 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect cloud.google.com/go/compute/metadata v0.9.0 // indirect @@ -122,7 +111,9 @@ require ( github.com/alibabacloud-go/tea v1.5.0 // indirect github.com/alibabacloud-go/tea-utils/v2 v2.0.9 // indirect github.com/aliyun/credentials-go v1.4.5 // indirect + github.com/andybalholm/cascadia v1.3.2 // indirect github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect + github.com/apache/arrow/go/v15 v15.0.2 // indirect github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.19 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.19 // indirect @@ -139,17 +130,17 @@ require ( github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.2 // indirect github.com/bytedance/gopkg v0.1.4 // indirect - github.com/bytedance/sonic v1.15.1 github.com/bytedance/sonic/loader v0.5.1 // indirect github.com/clbanning/mxj/v2 v2.7.0 // indirect github.com/cloudwego/base64x v0.1.7 // indirect github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dlclark/regexp2 v1.10.0 // indirect + github.com/dlclark/regexp2/v2 v2.5.2 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/ebitengine/purego v0.10.1 // indirect - github.com/eino-contrib/jsonschema v1.0.3 github.com/elastic/elastic-transport-go/v8 v8.8.0 // indirect + github.com/emersion/go-sasl v0.0.0-20241020182733-b788ff22d5a6 // indirect github.com/envoyproxy/go-control-plane/envoy v1.37.0 // indirect github.com/envoyproxy/protoc-gen-validate v1.3.3 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect @@ -165,23 +156,27 @@ require ( github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-playground/validator/v10 v10.30.2 // indirect + github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect + github.com/goccy/go-json v0.10.5 // indirect github.com/goccy/go-yaml v1.19.2 // indirect github.com/gogo/protobuf v1.3.2 // indirect github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe // indirect github.com/golang-sql/sqlexp v0.1.0 // indirect + github.com/google/flatbuffers v23.5.26+incompatible // indirect github.com/google/go-tpm v0.9.8 // indirect + github.com/google/jsonschema-go v0.4.3 // indirect + github.com/google/pprof v0.0.0-20230207041349-798e818bf904 // indirect github.com/google/s2a-go v0.1.9 // indirect github.com/googleapis/enterprise-certificate-proxy v0.3.17 // indirect github.com/googleapis/gax-go/v2 v2.23.0 // indirect github.com/goph/emperror v0.17.2 // indirect github.com/hashicorp/hcl v1.0.0 // indirect - github.com/infiniflow/onnxruntime_go v1.23.1 github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect - github.com/jackc/pgx/v5 v5.10.0 github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect + github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/compress v1.18.6 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/klauspost/crc32 v1.3.0 // indirect @@ -194,7 +189,6 @@ require ( github.com/minio/highwayhash v1.0.4 // indirect github.com/minio/md5-simd v1.1.2 // indirect github.com/mitchellh/mapstructure v1.5.0 // indirect - github.com/modelcontextprotocol/go-sdk v1.8.0 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/nats-io/jwt/v2 v2.8.2 // indirect @@ -204,6 +198,7 @@ require ( github.com/pelletier/go-toml/v2 v2.3.1 // indirect github.com/philhofer/fwd v1.2.0 // indirect github.com/phpdave11/gofpdi v1.0.14-0.20211212211723-1f10f9844311 // indirect + github.com/pierrec/lz4/v4 v4.1.18 // indirect github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.59.1 // indirect @@ -213,6 +208,8 @@ require ( github.com/rs/xid v1.6.0 // indirect github.com/sagikazarmark/locafero v0.4.0 // indirect github.com/sagikazarmark/slog-shim v0.1.0 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/sirupsen/logrus v1.9.3 // indirect github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect github.com/sourcegraph/conc v0.3.0 // indirect @@ -234,6 +231,7 @@ require ( github.com/xuri/efp v0.0.1 // indirect github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 // indirect github.com/yargevad/filepathx v1.0.0 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/yuin/gopher-lua v1.1.1 // indirect go.mongodb.org/mongo-driver/v2 v2.6.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect @@ -247,7 +245,10 @@ require ( golang.org/x/arch v0.27.0 // indirect golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect golang.org/x/sys v0.47.0 // indirect + golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 // indirect golang.org/x/time v0.15.0 // indirect + golang.org/x/tools v0.48.0 // indirect + golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect google.golang.org/genproto v0.0.0-20260519071638-aa98bba5eb94 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260630182238-925bb5da69e7 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260630182238-925bb5da69e7 // indirect diff --git a/go.sum b/go.sum index 010aaa20df..735fc5c47f 100644 --- a/go.sum +++ b/go.sum @@ -59,6 +59,8 @@ github.com/JohannesKaufmann/html-to-markdown v1.6.0 h1:04VXMiE50YYfCfLboJCLcgqF5 github.com/JohannesKaufmann/html-to-markdown v1.6.0/go.mod h1:NUI78lGg/a7vpEJTz/0uOcYMaibytE4BUOQS8k78yPQ= github.com/LuxorLabs/tenki-sdk-go/sandbox v0.7.0 h1:PBQ+ad8fZiLH/1ij1hoynR9aAPVmI3EzT7DiLNmUdBM= github.com/LuxorLabs/tenki-sdk-go/sandbox v0.7.0/go.mod h1:rBypesed6hSrj+OgpJwyMVBF76PaQ1iOqc2RGUL4RmI= +github.com/Masterminds/semver/v3 v3.5.0 h1:kQceYJfbupGfZOKZQg0kou0DgAKhzDg2NZPAwZ/2OOE= +github.com/Masterminds/semver/v3 v3.5.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM= github.com/PuerkitoBio/goquery v1.9.2 h1:4/wZksC3KgkQw7SQgkKotmKljk0M6V8TUvA8Wb4yPeE= github.com/PuerkitoBio/goquery v1.9.2/go.mod h1:GHPCaP0ODyyxqcNoFGYlAprUFH81NuRPd0GX3Zu2Mvk= github.com/RaveNoX/go-jsoncommentstrip v1.0.0/go.mod h1:78ihd09MekBnJnxpICcwzCMzGrKSKYe4AqU6PDYYpjk= @@ -162,7 +164,10 @@ github.com/aws/smithy-go v1.24.2/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqx github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmatcuk/doublestar v1.1.1 h1:YroD6BJCZBYx06yYFEWvUuKVWQn3vLLQAVmDmvTSaiQ= github.com/bmatcuk/doublestar v1.1.1/go.mod h1:UD6OnuiIn0yFxxA2le/rnRU1G4RaI4UvFv1sNto9p6w= +github.com/bmatcuk/doublestar/v4 v4.10.0 h1:zU9WiOla1YA122oLM6i4EXvGW62DvKZVxIe6TYWexEs= +github.com/bmatcuk/doublestar/v4 v4.10.0/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc= github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= github.com/browserbase/stagehand-go/v3 v3.21.0 h1:4DMtqcibQSobwLvuODla8SdkF0sk3BzsbEqG3eG+KI4= github.com/browserbase/stagehand-go/v3 v3.21.0/go.mod h1:2w/HwHce4GmhHWsKamELt/TLLtDJFXmCUHbYqPCH62k= @@ -191,8 +196,8 @@ github.com/clbanning/mxj/v2 v2.7.0/go.mod h1:hNiWqW14h+kc+MdF9C6/YoRfjEJoR3ou6tn github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= github.com/cloudwego/base64x v0.1.7 h1:NppS+Fgzg5ovhn4NkUXaDT3x9jldgH5ToMCqzBSi2zI= github.com/cloudwego/base64x v0.1.7/go.mod h1:Cu1PV9zfrSf7ET2tIbWbbEy7jO7HHJ13q4X2SQ8aWYg= -github.com/cloudwego/eino v0.9.14 h1:suNVibjtkPMiW7csFBdBqN3FRG0nRlCqoGKE7t0UDwY= -github.com/cloudwego/eino v0.9.14/go.mod h1:OBD1mrkfkt/pJa4rkg1P0VnaMeOVl7l8IAdEqY//3IQ= +github.com/cloudwego/eino v0.10.0-alpha.33 h1:2cpoeyjgzpweISH79XwN1buF0ruLSFyetVJ1t0MvdbQ= +github.com/cloudwego/eino v0.10.0-alpha.33/go.mod h1:OBD1mrkfkt/pJa4rkg1P0VnaMeOVl7l8IAdEqY//3IQ= github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= @@ -206,7 +211,11 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dlclark/regexp2/v2 v2.5.2 h1:HAsucWRhsqcDzl6Ua9aR8JwYOTzrZyPrF0/FNxJVAI0= +github.com/dlclark/regexp2/v2 v2.5.2/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU= github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ= +github.com/dop251/goja v0.0.0-20260820211235-95a30dcd3fa5 h1:vtyXXUeZQM4iNxud6zzTtdxT95nBJrelu77+DS0Vyi0= +github.com/dop251/goja v0.0.0-20260820211235-95a30dcd3fa5/go.mod h1:LiIEzozrcvNXorsG/3+ypGqdTUAqZryhzSsqi0oU/Qg= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/ebitengine/purego v0.10.1 h1:dewVBCBT2GaMu1SrNTYxQhgQBethzfhiwvZiLGP/qyY= @@ -275,6 +284,8 @@ github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJn github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= github.com/go-playground/validator/v10 v10.30.2 h1:JiFIMtSSHb2/XBUbWM4i/MpeQm9ZK2xqPNk8vgvu5JQ= github.com/go-playground/validator/v10 v10.30.2/go.mod h1:mAf2pIOVXjTEBrwUMGKkCWKKPs9NheYGabeB04txQSc= +github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU= +github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg= github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc= github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI= github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= @@ -321,8 +332,8 @@ github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+ github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc= github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0= -github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ= -github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo= +github.com/google/pprof v0.0.0-20230207041349-798e818bf904 h1:4/hN5RUoecvl+RmJRE2YxKWtnnQls6rQjjW5oV7qg2U= +github.com/google/pprof v0.0.0-20230207041349-798e818bf904/go.mod h1:uglQLonpP8qtYCYyzA+8c/9qtqgA3qsXGYqCPKARAFg= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index adf9eecc1e..e6305f6b2d 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -31,7 +31,6 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity/models" - "ragflow/internal/tokenizer" "go.uber.org/zap" ) @@ -245,7 +244,7 @@ func (e *einoChatInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvok Thinking: out.ReasoningContent, Model: modelName, Stopped: true, - Usage: usageFromCM(cm), + Usage: usageFromMessage(out), } if calls := nativeToolCalls(out.ToolCalls); len(calls) > 0 { resp.ToolCalls = calls @@ -253,17 +252,28 @@ func (e *einoChatInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvok return resp, nil } -// usageFromCM copies a ChatModel's LastUsage split into the chat.Response Usage, -// and mirrors its total into Tokens for the callers that still read the legacy -// single counter. -func usageFromCM(cm *models.ChatModel) *chat.Usage { - if cm == nil || cm.LastUsage == nil { +// usageFromMessage copies the call's own token split - which the model attaches to the +// message it returns (ResponseMeta.Usage) - into the chat.Response Usage. The split +// travels with the call instead of being read back off the shared ChatModel, where a +// concurrent call could have replaced it before the read. +func usageFromMessage(msg *schema.Message) *chat.Usage { + if msg == nil || msg.ResponseMeta == nil { + return nil + } + return usageFromEinoUsage(msg.ResponseMeta.Usage) +} + +// usageFromEinoUsage converts eino's own usage metadata into the chat.Response Usage. +// A nil usage yields a nil Response.Usage, which is the honest answer for a call whose +// provider reported none - better than a value left behind by another call. +func usageFromEinoUsage(u *schema.TokenUsage) *chat.Usage { + if u == nil { return nil } return &chat.Usage{ - PromptTokens: cm.LastUsage.PromptTokens, - CompletionTokens: cm.LastUsage.CompletionTokens, - TotalTokens: cm.LastUsage.TotalTokens, + PromptTokens: u.PromptTokens, + CompletionTokens: u.CompletionTokens, + TotalTokens: u.TotalTokens, } } @@ -410,6 +420,11 @@ func streamEinoChat(ctx context.Context, wrapper *models.EinoChatModel, msgs []* var content, reasoning strings.Builder pendingCalls := map[int]*schema.ToolCall{} + // A streamed call reports its usage on the last message that carries it + // (stream_options.include_usage), so keep the newest one and read it out below. + // It is this call's own split: the model no longer parks it on a field shared by + // every call on the ChatModel. + var streamUsage *schema.TokenUsage for { chunk, recvErr := sr.Recv() if recvErr == io.EOF { @@ -421,6 +436,9 @@ func streamEinoChat(ctx context.Context, wrapper *models.EinoChatModel, msgs []* if chunk == nil { continue } + if chunk.ResponseMeta != nil && chunk.ResponseMeta.Usage != nil { + streamUsage = chunk.ResponseMeta.Usage + } // Streaming tool-call deltas are merged before any content handling so a // chunk that carries both is processed exactly once. for _, tc := range chunk.ToolCalls { @@ -469,7 +487,7 @@ func streamEinoChat(ctx context.Context, wrapper *models.EinoChatModel, msgs []* Thinking: reasoning.String(), Model: modelName, Stopped: true, - Usage: usageFromCM(cm), + Usage: usageFromEinoUsage(streamUsage), } if len(pendingCalls) > 0 { // Preserve provider-given index order (Python's final_tool_calls.values() @@ -553,7 +571,7 @@ func (c *resolvedModelInvoker) modelFor(req ChatInvokeRequest) string { // Invoke satisfies ChatInvoker. func (c *resolvedModelInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRequest) (*ChatInvokeResponse, error) { modelName := c.modelFor(req) - wrapper, cm, err := c.resolvedChatWrapper(modelName, req) + wrapper, _, err := c.resolvedChatWrapper(modelName, req) if err != nil { return nil, err } @@ -566,7 +584,7 @@ func (c *resolvedModelInvoker) Invoke(ctx context.Context, db *gorm.DB, req Chat Thinking: out.ReasoningContent, Model: modelName, Stopped: true, - Usage: usageFromCM(cm), + Usage: usageFromMessage(out), } if calls := nativeToolCalls(out.ToolCalls); len(calls) > 0 { resp.ToolCalls = calls @@ -665,7 +683,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[strin return nil, fmt.Errorf("component: LLM.Invoke: resolve model: %w", err) } // Resolve the model's context window (content_length) for message - // fitting. 0 means the model is unknown → fitMessages falls back to + // fitting. 0 means the model is unknown → chat.FitMessages falls back to // 8192, matching Python's chat_mdl.max_length = model_config.get("max_tokens") or 8192. // tenantID scopes composite-reference resolution to the tenant's own rows // so a per-model "max_tokens" override in tenant_model.extra is honored. @@ -677,8 +695,8 @@ func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[strin } contentLength := dao.ResolveModelContentLength(ctx, db, tenantID, originalModelID, p.Driver, p.ModelID) if contentLength <= 0 { - // A 0 makes fitMessages fall back to the 8192 default budget, which can - // silently discard most of a large-context prompt, so surface the + // A 0 makes chat.FitMessages fall back to the 8192 default budget, which + // can silently discard most of a large-context prompt, so surface the // resolution failure for diagnosis. common.Warn("llm: content_length not resolved, falling back to 8192", zap.String("model_ref", originalModelID), @@ -805,7 +823,7 @@ func (c *LLMComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[strin // generation length only. { // The system prompt is already embedded as the first message - // in msgs by buildMessagesWithImages; pass "" so fitMessages + // in msgs by buildMessagesWithImages; pass "" so chat.FitMessages // does not duplicate it. fitted, fitErr := chat.FitMessages("", msgs, contentLength) if fitErr != "" { @@ -1367,157 +1385,6 @@ func mergeLLMParam(base LLMParam, inputs map[string]any) LLMParam { return p } -// effectiveContextLength returns maxLength if positive, otherwise 8192. -// Mirrors Python's LLM.effective_context_length in PR #16413 — prevents -// zero/negative context windows from silently trimming all prompt content. -func effectiveContextLength(maxLength int) int { - if maxLength > 0 { - return maxLength - } - return 8192 -} - -// contextFitBudget returns 97% of the effective context length as the -// token budget for message_fit_in. Mirrors Python's LLM.context_fit_budget -// in PR #16413. -func contextFitBudget(maxLength int) int { - return int(float64(effectiveContextLength(maxLength)) * 0.97) -} - -// validateFittedMessages checks that the fitted message list is non-empty -// and the last message is a non-empty user turn (content or multi-modal -// parts). Returns an error string on failure, empty string on success. -// Python requires len >= 2 because the system prompt is always injected -// upstream; Go allows len >= 1 because the system message may be embedded -// inside msgs (from buildMessagesWithImages) or absent entirely. -func validateFittedMessages(msgFit []schema.Message) string { - if len(msgFit) == 0 { - return "**ERROR**: message_fit_in produced insufficient messages for LLM" - } - last := msgFit[len(msgFit)-1] - if last.Role != schema.User { - return "**ERROR**: LLM last message is not a user turn after prompt fitting; check model content_length context setting" - } - if strings.TrimSpace(last.Content) == "" && len(last.UserInputMultiContent) == 0 { - return "**ERROR**: LLM user message is empty after prompt fitting; check model content_length context setting" - } - return "" -} - -// FitMessages exposes message_fit_in semantics (LLM.fit_messages, PR #16413) -// for packages that compose final-answer prompts outside the agent loop. -func FitMessages(systemPrompt string, msgs []schema.Message, maxLength int) ([]schema.Message, string) { - return fitMessages(systemPrompt, msgs, maxLength) -} - -// fitMessages calls message_fit_in semantics on the given messages and -// validates that the result ends with a non-empty user turn. Returns the -// fitted messages and an error string (empty on success). -// Mirrors Python's LLM.fit_messages in PR #16413. -func fitMessages(systemPrompt string, msgs []schema.Message, maxLength int) ([]schema.Message, string) { - // Deep-copy msgs (mirrors Python's deepcopy) to avoid mutating caller's slice. - copied := make([]schema.Message, len(msgs)) - for i, m := range msgs { - cloned := slices.Clone(m.UserInputMultiContent) - for j, p := range cloned { - if p.Image != nil { - imgCopy := *p.Image - if p.Image.URL != nil { - u := *p.Image.URL - imgCopy.URL = &u - } - cloned[j].Image = &imgCopy - } - } - copied[i] = schema.Message{ - Role: m.Role, - Content: m.Content, - UserInputMultiContent: cloned, - } - } - - // Convert to tokenizer.Message. Track where each entry's text lives - // (plain Content or a multi-modal text part) so the fitted text can be - // written back to the right field. Entries with no text at all - // (image-only turns) carry an empty Content for fitting and survive - // fitting when kept. - type fitSource struct { - copiedIdx int // index into copied; -1 for the synthetic system prompt - multiIdx int // -1 means the text lives in Content - textInContent bool // the original message carried text in Content - } - all := make([]tokenizer.Message, 0, 1+len(copied)) - sources := make([]fitSource, 0, 1+len(copied)) - - if systemPrompt != "" { - all = append(all, tokenizer.Message{Role: "system", Content: systemPrompt}) - sources = append(sources, fitSource{copiedIdx: -1, multiIdx: 0}) - } - - for i := range copied { - text := copied[i].Content - multiIdx := -1 - hadText := text != "" - if !hadText { - // Fold every non-empty text part into the token budget: only the - // first text part is written back, so leaving later parts out - // would let text exceed the fitted budget after reconstruction. - var textParts []string - for j, p := range copied[i].UserInputMultiContent { - if p.Type == schema.ChatMessagePartTypeText && p.Text != "" { - textParts = append(textParts, p.Text) - if multiIdx < 0 { - multiIdx = j - } - } - } - if len(textParts) > 0 { - text = strings.Join(textParts, "\n\n") - hadText = true - } - } - all = append(all, tokenizer.Message{Role: string(copied[i].Role), Content: text}) - sources = append(sources, fitSource{copiedIdx: i, multiIdx: multiIdx, textInContent: copied[i].Content != ""}) - } - - // Use 97% of effective context as the token budget. - budget := contextFitBudget(maxLength) - kept, keptIdx, _ := tokenizer.Fit(all, budget) - - // Convert back to []schema.Message. tokenizer.Fit reports exactly which - // entries are kept (keptIdx); dropped entries are simply absent, so no - // empty-content sentinel is needed and image-only turns are preserved. - result := make([]schema.Message, 0, len(kept)) - for j, i := range keptIdx { - src := sources[i] - if src.copiedIdx < 0 { - result = append(result, schema.Message{Role: schema.System, Content: kept[j].Content}) - continue - } - m := copied[src.copiedIdx] - if src.multiIdx >= 0 && src.multiIdx < len(m.UserInputMultiContent) { - m.UserInputMultiContent[src.multiIdx].Text = kept[j].Content - // Drop any additional text parts: their content was folded into - // the first part before fitting, so keeping them would re-introduce - // text outside the token budget. - keptParts := m.UserInputMultiContent[:0] - for k, part := range m.UserInputMultiContent { - if part.Type == schema.ChatMessagePartTypeText && k != src.multiIdx { - continue - } - keptParts = append(keptParts, part) - } - m.UserInputMultiContent = keptParts - } else if src.textInContent { - // Always write the fitted text back (even when trimmed to empty): - // leaving the original would send untrimmed content past the budget. - m.Content = kept[j].Content - } - result = append(result, m) - } - return result, validateFittedMessages(result) -} - // stringFrom extracts a string from inputs[name], accepting both string and // fmt.Stringer-able values. func stringFrom(inputs map[string]any, name string) (string, bool) { diff --git a/internal/agent/retrievalbridge/runtime_adapter.go b/internal/agent/retrievalbridge/runtime_adapter.go index 2daf46945d..565947ee1c 100644 --- a/internal/agent/retrievalbridge/runtime_adapter.go +++ b/internal/agent/retrievalbridge/runtime_adapter.go @@ -29,22 +29,32 @@ import ( // RuntimeAdapter exposes the agent tool's NLPRetrievalAdapter (the production // hybrid-retrieval backend installed by the server) through the runtime // RetrievalService singleton that the agentic-RAG harness reads via -// harness.RuntimeRetriever. The wrapped tool service is read on every call (not -// captured at construction) because the server installs it during boot. +// harness.RuntimeRetriever. +// +// The wrapped service is captured at construction, NOT resolved from the runtime +// singleton on every call: agenttool.SetRetrievalService and +// runtime.SetRetrievalService write the SAME registry, and the server installs +// this adapter into it (after installing the NLP adapter). Resolving the target +// at call time therefore returned this adapter itself, and Search called itself +// until the stack overflowed on every agentic retrieval. // // This keeps a single retrieval implementation: the harness and the canvas both // search through the same NLPRetrievalAdapter. -type RuntimeAdapter struct{} +type RuntimeAdapter struct { + svc agenttool.RetrievalService +} -// NewRuntimeAdapter creates the runtime RetrievalService bridge. -func NewRuntimeAdapter() *RuntimeAdapter { - return &RuntimeAdapter{} +// NewRuntimeAdapter creates the runtime RetrievalService bridge over svc — the +// concrete NLPRetrievalAdapter installed at boot. A nil svc makes Search report +// agenttool.ErrRetrievalServiceMissing, the sentinel the registry itself used. +func NewRuntimeAdapter(svc agenttool.RetrievalService) *RuntimeAdapter { + return &RuntimeAdapter{svc: svc} } // Search forwards a runtime retrieval request to the installed tool service and // translates the result back to the runtime chunk shape. func (a *RuntimeAdapter) Search(ctx context.Context, db *gorm.DB, req agentrunt.RetrievalRequest) ([]agentrunt.RetrievalChunk, error) { - svc := agenttool.GetRetrievalService() + svc := a.svc if svc == nil { return nil, agenttool.ErrRetrievalServiceMissing } @@ -56,20 +66,12 @@ func (a *RuntimeAdapter) Search(ctx context.Context, db *gorm.DB, req agentrunt. db = dao.DB } toolReq := agenttool.RetrievalRequest{ - Query: req.Query, - DatasetIDs: req.DatasetIDs, - MemoryIDs: req.MemoryIDs, - TopN: req.TopN, - RerankCandidatesCount: req.RerankCandidatesCount, - TopK: req.TopK, - // VectorSimilarityWeight is the VECTOR weight (Python - // vector_similarity_weight) and is forwarded to the nlp layer - // UN-inverted. KeywordsSimilarityWeight (canvas semantics, keyword - // weight) stays on its own inversion path in nlpRequestFromRetrieval — - // conflating the two turned the agentic hybrid leg vector-dominant and - // the BM25 legs into pure-vector searches. - VectorSimilarityWeight: req.VectorSimilarityWeight, - DisableVectorLeg: req.DisableVectorLeg, + Query: req.Query, + DatasetIDs: req.DatasetIDs, + MemoryIDs: req.MemoryIDs, + TopN: req.TopN, + RerankCandidatesCount: req.RerankCandidatesCount, + TopK: req.TopK, KeywordsSimilarityWeight: req.KeywordsSimilarityWeight, UseKG: req.UseKG, SimilarityThreshold: req.SimilarityThreshold, @@ -90,33 +92,24 @@ func (a *RuntimeAdapter) Search(ctx context.Context, db *gorm.DB, req agentrunt. if err != nil { return nil, err } - out := make([]agentrunt.RetrievalChunk, 0, len(chunks)) - for _, c := range chunks { - out = append(out, agentrunt.RetrievalChunk{ - ID: c.ID, - Content: c.Content, - DocumentID: c.DocumentID, - DocumentName: c.DocumentName, - DatasetID: c.DatasetID, - ImageID: c.ImageID, - DocType: c.DocType, - URL: c.URL, - Positions: c.Positions, - Score: c.Score, - TermSimilarity: c.TermSimilarity, - VectorSimilarity: c.VectorSimilarity, - MomID: c.MomID, - }) - } - return out, nil + // agenttool.RetrievalChunk is a type ALIAS of agentrunt.RetrievalChunk (single + // owner: internal/agent/runtime), so the result already is the runtime slice: + // returning it directly is exact. Hand-writing a field-by-field copy here is + // how ChunkIndex/PageNum were silently dropped — every field the shared type + // gains now flows through without this file having to know about it. + return chunks, nil } -// rankFeatureOrNil converts a map-valued rank feature into the pointer form the -// tool request expects, returning nil when empty so a missing feature is +// rankFeatureOrNil converts a rank feature into the pointer form the tool +// request expects, returning nil when empty so a missing feature is // distinguished from an empty one (Python passes None for "no rank feature"). -func rankFeatureOrNil(m map[string]float64) *map[string]float64 { - if len(m) == 0 { +// +// Upstream carried the feature as a value map in the runtime request and +// converted it at this boundary; the request now carries the pointer form +// directly, so the same invariant is enforced on the pointer instead. +func rankFeatureOrNil(m *map[string]float64) *map[string]float64 { + if m == nil || len(*m) == 0 { return nil } - return &m + return m } diff --git a/internal/agent/retrievalbridge/runtime_adapter_test.go b/internal/agent/retrievalbridge/runtime_adapter_test.go index e58c46b8af..65c60144a5 100644 --- a/internal/agent/retrievalbridge/runtime_adapter_test.go +++ b/internal/agent/retrievalbridge/runtime_adapter_test.go @@ -2,6 +2,8 @@ package retrievalbridge import ( "context" + "errors" + "reflect" "testing" agentrunt "ragflow/internal/agent/runtime" @@ -11,20 +13,19 @@ import ( ) type captureRetrievalService struct { - req agenttool.RetrievalRequest + req agenttool.RetrievalRequest + chunks []agenttool.RetrievalChunk + err error } func (s *captureRetrievalService) Search(_ context.Context, _ *gorm.DB, req agenttool.RetrievalRequest) ([]agenttool.RetrievalChunk, error) { s.req = req - return nil, nil + return s.chunks, s.err } func TestRuntimeAdapterDisablesDenseFallback(t *testing.T) { - previous := agenttool.GetRetrievalService() service := &captureRetrievalService{} - agenttool.SetRetrievalService(service) - t.Cleanup(func() { agenttool.SetRetrievalService(previous) }) - _, err := NewRuntimeAdapter().Search(t.Context(), nil, agentrunt.RetrievalRequest{Query: "test"}) + _, err := NewRuntimeAdapter(service).Search(t.Context(), nil, agentrunt.RetrievalRequest{Query: "test"}) if err != nil { t.Fatal(err) } @@ -32,3 +33,72 @@ func TestRuntimeAdapterDisablesDenseFallback(t *testing.T) { t.Fatal("Agentic runtime retrieval must disable dense fallback") } } + +// TestRuntimeAdapterDoesNotResolveItsTargetFromTheSharedRegistry reproduces the +// server's boot order (cmd/ragflow_server.go): the NLP adapter is registered, +// then the bridge REPLACES it in the same shared registry +// (agenttool.SetRetrievalService and runtime.SetRetrievalService write one +// singleton). The bridge must still reach the adapter it was built with; +// resolving the registry on every call made Search call itself, which is a stack +// overflow — so a regression here crashes this test rather than failing it. +func TestRuntimeAdapterDoesNotResolveItsTargetFromTheSharedRegistry(t *testing.T) { + previous := agenttool.GetRetrievalService() + t.Cleanup(func() { agenttool.SetRetrievalService(previous) }) + + service := &captureRetrievalService{chunks: []agenttool.RetrievalChunk{{ID: "from-adapter"}}} + agenttool.SetRetrievalService(service) + bridge := NewRuntimeAdapter(service) + // boot order, second half: the bridge itself becomes the registry's value + agenttool.SetRetrievalService(bridge) + + chunks, err := bridge.Search(t.Context(), nil, agentrunt.RetrievalRequest{Query: "test"}) + if err != nil { + t.Fatal(err) + } + if len(chunks) != 1 || chunks[0].ID != "from-adapter" { + t.Fatalf("bridge did not reach the adapter it was built with: %+v", chunks) + } +} + +// TestRuntimeAdapterPreservesEveryChunkField pins the field-preservation contract +// at the bridge. A hand-written field list is how ChunkIndex/PageNum were +// silently dropped, and those two are what let the deep-read tools sort a +// document into reading order. +func TestRuntimeAdapterPreservesEveryChunkField(t *testing.T) { + want := agenttool.RetrievalChunk{ + ID: "c1", + Content: "content", + DocumentID: "doc", + DocumentName: "doc.md", + DatasetID: "kb", + ImageID: "img", + DocType: "text", + URL: "http://example.invalid", + Positions: []any{1.0, 2.0}, + ChunkIndex: 7, + PageNum: 3, + MomID: "mom", + Score: 0.9, + TermSimilarity: 0.5, + VectorSimilarity: 0.4, + } + service := &captureRetrievalService{chunks: []agenttool.RetrievalChunk{want}} + + got, err := NewRuntimeAdapter(service).Search(t.Context(), nil, agentrunt.RetrievalRequest{Query: "q"}) + if err != nil { + t.Fatal(err) + } + if len(got) != 1 { + t.Fatalf("got %d chunks, want 1", len(got)) + } + if !reflect.DeepEqual(got[0], want) { + t.Fatalf("chunk lost fields across the bridge:\n got %+v\nwant %+v", got[0], want) + } +} + +func TestRuntimeAdapterWithoutServiceReportsTheRegistrySentinel(t *testing.T) { + _, err := NewRuntimeAdapter(nil).Search(t.Context(), nil, agentrunt.RetrievalRequest{Query: "q"}) + if !errors.Is(err, agenttool.ErrRetrievalServiceMissing) { + t.Fatalf("want ErrRetrievalServiceMissing, got %v", err) + } +} diff --git a/internal/agent/runtime/maps.go b/internal/agent/runtime/maps.go new file mode 100644 index 0000000000..a52ce5f9fa --- /dev/null +++ b/internal/agent/runtime/maps.go @@ -0,0 +1,73 @@ +// +// 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 runtime + +import "encoding/json" + +// IntFromMap returns raw[key] as an int, tolerating the numeric types JSON +// decoding and the doc engine produce (float64, float32, int, int64). Shared by +// the retrieval tools so chunk metadata is parsed uniformly. +func IntFromMap(raw map[string]any, key string) int { + if f, ok := NumberFromMap(raw, key); ok { + return int(f) + } + return 0 +} + +// StringFromMap returns raw[key].(string) or "" if missing / wrong type. +func StringFromMap(raw map[string]any, key string) string { + if v, ok := raw[key]; ok { + if s, ok := v.(string); ok { + return s + } + } + return "" +} + +// FirstStringFromMap returns the first non-empty string among the given keys. +func FirstStringFromMap(raw map[string]any, keys ...string) string { + for _, key := range keys { + if value := StringFromMap(raw, key); value != "" { + return value + } + } + return "" +} + +// NumberFromMap returns raw[key].(float64) with a tolerant path for ints, JSON +// numbers and json.Number (which arises when decoding with UseNumber). +func NumberFromMap(raw map[string]any, key string) (float64, bool) { + v, ok := raw[key] + if !ok { + return 0, false + } + switch x := v.(type) { + case float64: + return x, true + case float32: + return float64(x), true + case int: + return float64(x), true + case int64: + return float64(x), true + case json.Number: + if f, err := x.Float64(); err == nil { + return f, true + } + } + return 0, false +} diff --git a/internal/agent/runtime/retrieval.go b/internal/agent/runtime/retrieval.go index 325efb9bb7..92e55ee6e3 100644 --- a/internal/agent/runtime/retrieval.go +++ b/internal/agent/runtime/retrieval.go @@ -14,10 +14,9 @@ // limitations under the License. // -// Retrieval contracts shared by the canvas agent runtime (internal/agent/tool) -// and the smart-reasoning agent (internal/agentic_rag). Keeping these here — -// in the engine-agnostic runtime package — means neither agent layer depends on -// the other: both depend on this shared +// Retrieval contracts shared by the canvas agent runtime (internal/agent/tool). +// Keeping these in the engine-agnostic runtime package means the tool layer +// re-exports them instead of owning a second copy. package runtime import ( @@ -63,36 +62,39 @@ type RetrievalChunk struct { // RetrievalRequest is the input to RetrievalService.Search. type RetrievalRequest struct { - Query string - DatasetIDs []string - MemoryIDs []string - TopN int - RerankCandidatesCount int - TopK int - // VectorSimilarityWeight is the VECTOR leg's weight (Python - // vector_similarity_weight), forwarded verbatim by the agentic harness. The - // canvas path keeps using KeywordsSimilarityWeight (the keyword weight the - // adapter inverts); see retrievalbridge/runtime_adapter.go. - VectorSimilarityWeight *float64 - // DisableVectorLeg mirrors Python passing embd_mdl=None: the backend runs - // the keyword-only branch with NO dense leg (not even a weight-0 one). - DisableVectorLeg bool + Query string + DatasetIDs []string + MemoryIDs []string + TopN int + // RerankCandidatesCount caps the candidate set pulled for reranking. Zero + // means "use the backend default". + RerankCandidatesCount int + TopK int KeywordsSimilarityWeight *float64 UseKG bool SimilarityThreshold *float64 + AllowDenseFallback *bool RerankID string CrossLanguages []string TOCEnhance bool MetaDataFilter map[string]any RetrievalFrom string - // DocScope restricts retrieval to a set of document ids. Empty = no doc filter. + // DocScope restricts retrieval to a set of document ids (the doc_id list + // routed by the dataset_navigation_by_tree tool). Empty = no doc filter. DocScope []string // TenantID is the calling tenant (== user_id in RAGFlow's data model). TenantID string // RankFeature is the label_question term→weight map passed through to the // engine so retrieval is biased toward the query's predicted topic class. // Mirrors engine nlp.RetrievalRequest.RankFeature. - RankFeature map[string]float64 + RankFeature *map[string]float64 + // UserID optionally filters memory messages by the user_id they were + // recorded with (the Retrieval node's "User ID" field, e.g. resolved + // from sys.user_id). Empty = no user filter. Only meaningful for + // retrieval_from=memory. + UserID string + // ExcludeCompiled excludes compiled-product rows from plain retrieval. + ExcludeCompiled bool // OnlyOriginalText, when true, restricts retrieval to ordinary document // text chunks (available_int=1 and no compile_kwd), excluding // knowledge-compiled products. diff --git a/internal/agent/tool/agentic_search.go b/internal/agent/tool/agentic_search.go index df444d885c..2bcb3b9b86 100644 --- a/internal/agent/tool/agentic_search.go +++ b/internal/agent/tool/agentic_search.go @@ -112,10 +112,11 @@ func (a *AgenticSearchTool) InvokableRun(ctx context.Context, argumentsInJSON st svc := GetRetrievalService() tenantID := canvasTenantID(ctx) - datasetIDs := args.KbIDs - if len(datasetIDs) == 0 { - datasetIDs = canvasDatasetIDs(ctx, nil) - } + // Route through canvasDatasetIDs so a trusted WithScope injection (when + // present) intersects the model-supplied kb_ids, preventing escalation + // outside the conversation's KBs. In the canvas runtime (no injection) + // explicit kb_ids remain authoritative, preserving prior behavior. + datasetIDs := canvasDatasetIDs(ctx, args.KbIDs) if svc == nil || tenantID == "" || len(datasetIDs) == 0 { return jsonChunksEmpty(), nil } diff --git a/internal/agent/tool/dataset_navigation.go b/internal/agent/tool/dataset_navigation.go index ff368e4cd5..297bdd4e2e 100644 --- a/internal/agent/tool/dataset_navigation.go +++ b/internal/agent/tool/dataset_navigation.go @@ -222,16 +222,16 @@ func (d *DatasetNavigationByTree) InvokableRun(ctx context.Context, argumentsInJ // as doc_ids to the retrieval), and collect()'s inScope still applies. if len(docs) == 0 { threshold := datasetNavRecallMinScore - w := datasetNavRecallVectorWeight + keywordsWeight := 1 - datasetNavRecallVectorWeight chunks, err := GetRetrievalService().Search(ctx, dao.DB, RetrievalRequest{ - Query: query, - DatasetIDs: datasetIDs, - TopN: datasetNavRecallTopN, - SimilarityThreshold: &threshold, - VectorSimilarityWeight: &w, - TenantID: tenantID, - DocScope: docScope, - RetrievalFrom: "dataset", + Query: query, + DatasetIDs: datasetIDs, + TopN: datasetNavRecallTopN, + SimilarityThreshold: &threshold, + KeywordsSimilarityWeight: &keywordsWeight, + TenantID: tenantID, + DocScope: docScope, + RetrievalFrom: "dataset", }) if err != nil { log.Printf("[Dataset navigation] content-recall retrieval failed: %v", err) diff --git a/internal/agent/tool/dataset_navigation_test.go b/internal/agent/tool/dataset_navigation_test.go index 57743911a3..f4c0ec9224 100644 --- a/internal/agent/tool/dataset_navigation_test.go +++ b/internal/agent/tool/dataset_navigation_test.go @@ -232,6 +232,30 @@ func TestDatasetNavigation_BlankRequestScopeFallsBackToDefaults(t *testing.T) { } } +// TestCanvasTenantID_FromCanvasState asserts the tenant id resolves from canvas +// state, falling back to user_id. +func TestCanvasTenantID_FromCanvasState(t *testing.T) { + state := runtime.NewCanvasState("run-1", "task-1") + state.Sys["tenant_id"] = "canvas-tenant" + ctx := runtime.WithState(context.Background(), state) + if got := canvasTenantID(ctx); got != "canvas-tenant" { + t.Errorf("canvasTenantID = %q, want canvas tenant canvas-tenant", got) + } + + // Fall back to user_id when tenant_id is absent. + state2 := runtime.NewCanvasState("run-1", "task-1") + state2.Sys["user_id"] = "user-1" + ctx2 := runtime.WithState(context.Background(), state2) + if got := canvasTenantID(ctx2); got != "user-1" { + t.Errorf("canvasTenantID = %q, want fallback user-1", got) + } + + // No canvas state at all → empty. + if got := canvasTenantID(context.Background()); got != "" { + t.Errorf("canvasTenantID = %q, want empty without canvas state", got) + } +} + func containsStr(s, sub string) bool { return len(s) > 0 && len(sub) > 0 && (s == sub || containsSub(s, sub)) } diff --git a/internal/agent/tool/retrieval.go b/internal/agent/tool/retrieval.go index f4d3b2bb51..fc9e50531a 100644 --- a/internal/agent/tool/retrieval.go +++ b/internal/agent/tool/retrieval.go @@ -38,14 +38,6 @@ import ( // the Python Canvas. var ErrGraphRAGNotSupported = errors.New("GraphRAG 检索暂不支持，请使用 Python Canvas 或关闭 use_kg") -// ErrRetrievalServiceMissing is returned when the -// internal/service/nlp RetrievalService is not registered. Wire a -// real implementation via SetRetrievalService at boot to resolve. -var ErrRetrievalServiceMissing = errors.New( - "Retrieval service not yet implemented (service not registered) — " + - "use Python Canvas or implement internal/service/nlp/retrieval.go", -) - // retrievalToolName preserves the Python typo ("dateset") for backward // compatibility with existing Canvas DSLs that reference the tool by name. const retrievalToolName = "search_my_dateset" diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 334cac4e81..b80bd0393f 100644 --- a/internal/agent/tool/retrieval_nlp.go +++ b/internal/agent/tool/retrieval_nlp.go @@ -67,6 +67,7 @@ import ( "regexp" "strings" + "ragflow/internal/agent/runtime" "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/engine" @@ -258,6 +259,7 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri } } query = retrievalUserPrefixPattern.ReplaceAllString(query, "") + vectorOnly := req.KeywordsSimilarityWeight != nil && *req.KeywordsSimilarityWeight < 0.01 // rank_feature (Python retrieve: rank_feature=label_question(question, // self.kbs)). Prefer a feature supplied on the request (computed by RAGTools // from its own KB objects) so the agentic tool stays authoritative; fall @@ -265,7 +267,10 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri var rankFeature map[string]float64 if req.RankFeature != nil && len(*req.RankFeature) > 0 { rankFeature = *req.RankFeature - } else if a.enhancer != nil { + } else if a.enhancer != nil && !vectorOnly { + // Rank features are extra scoring clauses on the TEXT leg; a + // vector-only request has no text leg, so labelling the question would + // be a tag lookup whose result nothing consumes. rankFeature = a.enhancer.LabelQuestion(ctx, query, datasets.kbs) } rerankModel, err := a.resolveRerankModel(ctx, req, datasets.kbs[0].TenantID) @@ -336,6 +341,7 @@ func nlpRequestFromRetrieval( Aggs: boolPtr(false), Highlight: boolPtr(false), AllowDenseFallback: req.AllowDenseFallback, + VectorOnly: req.KeywordsSimilarityWeight != nil && *req.KeywordsSimilarityWeight < 0.01, } if excludeCompiled { // Python hybrid_search excludes compiled products from plain retrieval @@ -358,22 +364,27 @@ func nlpRequestFromRetrieval( if req.SimilarityThreshold != nil { nlpReq.SimilarityThreshold = req.SimilarityThreshold } - if req.DisableVectorLeg { - // Python embd_mdl=None (bm25_search / grep_search / retrieve with - // using_embedding=False): keyword-only search, NO dense leg at all — - // not even a weight-0 one, which would still constrain the candidate - // pool through the KNN similarity option. - nlpReq.EmbeddingModel = nil - } - if req.VectorSimilarityWeight != nil { - // Agentic harness path: the vector weight arrives already in vector - // semantics (Python vector_similarity_weight) — forward verbatim. - nlpReq.VectorSimilarityWeight = req.VectorSimilarityWeight - } else if req.KeywordsSimilarityWeight != nil { - // Canvas path: keywords_similarity_weight is the KEYWORD weight - // (user-facing config); the vector weight is its complement. + if req.KeywordsSimilarityWeight != nil { + // KeywordsSimilarityWeight is the user-facing keyword weight. The NPL + // backend takes the complementary vector weight. Near-total keyword + // weighting also disables the dense leg entirely, as its zero-weight + // KNN filter would otherwise still constrain the candidate set. vectorSimilarityWeight := 1 - *req.KeywordsSimilarityWeight nlpReq.VectorSimilarityWeight = &vectorSimilarityWeight + if *req.KeywordsSimilarityWeight > 0.99 { + nlpReq.EmbeddingModel = nil + } + } + // Restrict to ordinary document text chunks with the SAME filter grep_chunks + // applies: available_int=1 (must_not available_int<1 semantics so chunks + // whose available_int is absent still pass) and must_not exists compile_kwd + // to exclude knowledge-compiled products. Setting available_int explicitly + // here matches grep exactly (both go through buildBoolQueryFromCondition). + if req.OnlyOriginalText { + nlpReq.Filter = map[string]interface{}{ + "available_int": 1, + "must_not": map[string]interface{}{"exists": "compile_kwd"}, + } } return nlpReq } @@ -570,40 +581,33 @@ func (a *NLPRetrievalAdapter) resolveRerankModel( // can't break the whole result list. func translateChunk(raw map[string]any) RetrievalChunk { return RetrievalChunk{ - ID: stringFromMap(raw, "chunk_id"), + ID: StringFromMap(raw, "chunk_id"), Content: contentFromMap(raw), - DocumentID: stringFromMap(raw, "doc_id"), - DocumentName: stringFromMap(raw, "docnm_kwd"), - DatasetID: stringFromMap(raw, "kb_id"), - ImageID: firstStringFromMap(raw, "image_id", "img_id"), - DocType: firstStringFromMap(raw, "doc_type_kwd", "doc_type"), - URL: firstStringFromMap(raw, "url", "document_url", "doc_url"), + DocumentID: StringFromMap(raw, "doc_id"), + DocumentName: StringFromMap(raw, "docnm_kwd"), + DatasetID: StringFromMap(raw, "kb_id"), + ImageID: FirstStringFromMap(raw, "image_id", "img_id"), + DocType: FirstStringFromMap(raw, "doc_type_kwd", "doc_type"), + URL: FirstStringFromMap(raw, "url", "document_url", "doc_url"), Positions: firstValueFromMap(raw, "positions", "position_int"), - MomID: stringFromMap(raw, "mom_id"), + MomID: StringFromMap(raw, "mom_id"), + ChunkIndex: IntFromMap(raw, "chunk_order_int"), + PageNum: IntFromMap(raw, "page_num_int"), Score: scoreFromMap(raw), TermSimilarity: scoreValueFromMap(raw, "term_similarity"), VectorSimilarity: scoreValueFromMap(raw, "vector_similarity"), } } -// stringFromMap returns raw[key].(string) or "" if missing / wrong -// type. Keeps the translator compact. -func stringFromMap(raw map[string]any, key string) string { - if v, ok := raw[key]; ok { - if s, ok := v.(string); ok { - return s - } - } - return "" +// IntFromMap, StringFromMap, FirstStringFromMap and NumberFromMap are +// re-exported from internal/agent/runtime (single owner) so the canvas tool +// package keeps its helper names without owning a second copy. +func IntFromMap(raw map[string]any, key string) int { return runtime.IntFromMap(raw, key) } +func StringFromMap(raw map[string]any, key string) string { + return runtime.StringFromMap(raw, key) } - -func firstStringFromMap(raw map[string]any, keys ...string) string { - for _, key := range keys { - if value := stringFromMap(raw, key); value != "" { - return value - } - } - return "" +func FirstStringFromMap(raw map[string]any, keys ...string) string { + return runtime.FirstStringFromMap(raw, keys...) } func firstValueFromMap(raw map[string]any, keys ...string) any { @@ -621,10 +625,10 @@ func firstValueFromMap(raw map[string]any, keys ...string) any { // the model sees in Python; we use it here too. Empty / missing → // fall back to content_ltks; both empty → empty string. func contentFromMap(raw map[string]any) string { - if v := stringFromMap(raw, "content_with_weight"); v != "" { + if v := StringFromMap(raw, "content_with_weight"); v != "" { return v } - return stringFromMap(raw, "content_ltks") + return StringFromMap(raw, "content_ltks") } // scoreFromMap returns the chunk's similarity score. nlp populates @@ -633,11 +637,11 @@ func contentFromMap(raw map[string]any) string { // zero, average the two sub-scores. Wrong-type values → fall through // to sub-scores; missing sub-scores → 0. func scoreFromMap(raw map[string]any) float64 { - if f, ok := numberFromMap(raw, "similarity"); ok { + if f, ok := NumberFromMap(raw, "similarity"); ok { return f } - term, termOK := numberFromMap(raw, "term_similarity") - vec, vecOK := numberFromMap(raw, "vector_similarity") + term, termOK := NumberFromMap(raw, "term_similarity") + vec, vecOK := NumberFromMap(raw, "vector_similarity") if termOK && vecOK { return (term + vec) / 2 } @@ -651,28 +655,12 @@ func scoreFromMap(raw map[string]any) float64 { } func scoreValueFromMap(raw map[string]any, key string) float64 { - value, _ := numberFromMap(raw, key) + value, _ := NumberFromMap(raw, key) return value } -// numberFromMap returns raw[key].(float64) with a tolerant path -// for ints. JSON unmarshaling can produce either. -func numberFromMap(raw map[string]any, key string) (float64, bool) { - v, ok := raw[key] - if !ok { - return 0, false - } - switch x := v.(type) { - case float64: - return x, true - case float32: - return float64(x), true - case int: - return float64(x), true - case int64: - return float64(x), true - } - return 0, false +func NumberFromMap(raw map[string]any, key string) (float64, bool) { + return runtime.NumberFromMap(raw, key) } func boolPtr(b bool) *bool { return &b } diff --git a/internal/agent/tool/retrieval_nlp_test.go b/internal/agent/tool/retrieval_nlp_test.go index eb2e92eb34..80d53ab27d 100644 --- a/internal/agent/tool/retrieval_nlp_test.go +++ b/internal/agent/tool/retrieval_nlp_test.go @@ -276,6 +276,44 @@ func TestNLPRequestFromRetrieval_ThreadsSearchControls(t *testing.T) { if got.EmbeddingModel != embeddingModel { t.Fatal("EmbeddingModel was not passed to nlp retrieval request") } + if got.VectorOnly { + t.Fatal("mixed retrieval must not be vector-only") + } +} + +func TestNLPRequestFromRetrieval_DerivesPureModesFromKeywordWeight(t *testing.T) { + embeddingModel := &modelModule.EmbeddingModel{} + vectorOnlyWeight := 0.0 + vectorOnly := nlpRequestFromRetrieval(RetrievalRequest{ + Query: "hi", + DatasetIDs: []string{"kb-1"}, + KeywordsSimilarityWeight: &vectorOnlyWeight, + }, []string{"tenant-a"}, 3, embeddingModel, false) + if !vectorOnly.VectorOnly { + t.Fatal("keywords weight 0 must derive vector-only search") + } + if vectorOnly.EmbeddingModel != embeddingModel { + t.Fatal("vector-only search must retain the embedding model") + } + if vectorOnly.VectorSimilarityWeight == nil || *vectorOnly.VectorSimilarityWeight != 1 { + t.Fatalf("vector-only vector weight = %v, want 1", vectorOnly.VectorSimilarityWeight) + } + + keywordOnlyWeight := 1.0 + keywordOnly := nlpRequestFromRetrieval(RetrievalRequest{ + Query: "hi", + DatasetIDs: []string{"kb-1"}, + KeywordsSimilarityWeight: &keywordOnlyWeight, + }, []string{"tenant-a"}, 3, embeddingModel, false) + if keywordOnly.VectorOnly { + t.Fatal("keywords weight 1 must not derive vector-only search") + } + if keywordOnly.EmbeddingModel != nil { + t.Fatal("keywords weight 1 must disable the vector leg") + } + if keywordOnly.VectorSimilarityWeight == nil || *keywordOnly.VectorSimilarityWeight != 0 { + t.Fatalf("keyword-only vector weight = %v, want 0", keywordOnly.VectorSimilarityWeight) + } } func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) { diff --git a/internal/agent/tool/retrieval_service.go b/internal/agent/tool/retrieval_service.go index e615056df5..527d9652e7 100644 --- a/internal/agent/tool/retrieval_service.go +++ b/internal/agent/tool/retrieval_service.go @@ -14,261 +14,42 @@ // limitations under the License. // -// RetrievalService is the abstract interface for retrieval. -// The Go path exposes only the parameters currently threaded into -// internal/service/nlp. Extra Python-only options should be added -// here only when the adapter can pass them through. +// Retrieval contracts live in internal/agent/runtime (the engine-agnostic +// package the canvas agent depends on). This file re-exports them under the +// historical tool.XXX names so the canvas tool package keeps its public API +// stable without owning a second copy. package tool import ( - "context" - "errors" - "fmt" - "sync" - - "gorm.io/gorm" + "ragflow/internal/agent/runtime" ) -// RetrievalChunk is the minimal shape RetrievalService returns. The -// full Chunk type (with document_id, docnm_kwd, position, etc.) -// lives in internal/entity and is wired in by a follow-up phase. -type RetrievalChunk struct { - ID string - Content string - DocumentID string - DocumentName string - DatasetID string - ImageID string - URL string - Positions any - Score float64 - TermSimilarity float64 - VectorSimilarity float64 - // MomID is the parent chunk id when this chunk is a child fragment; empty - // for top-level chunks. Threaded through so the harness can run - // retrieval_by_children (child fragments are promoted to their parent chunk). - MomID string - // DocType is the engine's doc_type_kwd ("text" / "image" / "table"). It is - // what lets the answer reference card render an image chunk as an image - // instead of dropping the marker, so it must survive the hop to the harness - // evidence pool alongside ImageID. - DocType string -} - -// RetrievalRequest is the input to RetrievalService.Search. -type RetrievalRequest struct { - Query string - DatasetIDs []string - MemoryIDs []string - TopN int - RerankCandidatesCount int - TopK int - KeywordsSimilarityWeight *float64 - // VectorSimilarityWeight is the VECTOR leg's weight forwarded DIRECTLY (no - // 1-w inversion) by the agentic harness bridge. When set it wins over the - // canvas KeywordsSimilarityWeight inversion in nlpRequestFromRetrieval. - VectorSimilarityWeight *float64 - // DisableVectorLeg mirrors Python embd_mdl=None: the nlp layer runs the - // keyword-only branch with NO dense leg at all. - DisableVectorLeg bool - UseKG bool - SimilarityThreshold *float64 - AllowDenseFallback *bool - RerankID string - CrossLanguages []string - TOCEnhance bool - // RankFeature mirrors Python RAGTools.retrieve's rank_feature argument: - // question-type tags (from label_question) the retriever uses to boost - // matching chunks. The Go engine consumes it as a tag → weight map, so it - // is *map[string]float64. Nil means no tag boost (fall back to enhancer). - RankFeature *map[string]float64 - MetaDataFilter map[string]any - RetrievalFrom string - // DocScope restricts retrieval to a set of document ids (the doc_id list - // routed by the dataset_navigation_by_tree tool). Empty = no doc filter. - DocScope []string - // TenantID is the calling tenant (== user_id in RAGFlow's data model). - // It is used for dataset-name resolution and memory access. Reads from - // CanvasState.Sys["user_id"] when empty (set by the Begin component at - // internal/agent/component/begin.go:82). - TenantID string - // UserID optionally filters memory messages by the user_id they were - // recorded with (the Retrieval node's "User ID" field, e.g. resolved - // from sys.user_id). Empty = no user filter. Only meaningful for - // retrieval_from=memory. - UserID string - // ExcludeCompiled excludes compiled-product rows from plain retrieval - // (Python hybrid_search passes must_not={"exists": "compile_kwd"}, - // search.py:171). Compiled products have their own expansion step, so the - // base retrieval should surface ordinary document chunks only. - ExcludeCompiled bool -} - -// RetrievalService is the knowledge-base search interface used by the tool. -// The server installs NLPRetrievalAdapter during boot. -type RetrievalService interface { - Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) -} - -// MemoryRetrievalService is the memory-message retrieval surface used when -// retrieval_from=memory. It is separate from knowledge-base retrieval because -// memory messages live in different indices and have a different result shape. -type MemoryRetrievalService interface { - Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) -} - -// KGRetrievalService is the GraphRAG retrieval surface. The -// KBRetrieval service and the KGRetrieval service are kept -// separate on purpose: the kg backend's signature requires -// per-tenant chat + embedding model handles, while the nlp -// backend resolves them lazily through RetrievalRequest's -// EmbeddingModel field. Splitting the registries means each -// adapter can be tested in isolation and wired independently -// at boot. -type KGRetrievalService interface { - Search(ctx context.Context, db *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) -} +// Re-exported retrieval contracts (single owner: internal/agent/runtime). +type ( + RetrievalChunk = runtime.RetrievalChunk + RetrievalRequest = runtime.RetrievalRequest + RetrievalService = runtime.RetrievalService + MemoryRetrievalService = runtime.MemoryRetrievalService + KGRetrievalService = runtime.KGRetrievalService +) // ErrRetrievalServiceMissing is declared in retrieval.go so callers and the // default stub share the same sentinel. var ( - retrievalServiceMu sync.RWMutex - retrievalServiceImpl RetrievalService = stubRetrievalService{} + ErrRetrievalServiceMissing = runtime.ErrRetrievalServiceMissing + ErrMemoryRetrievalServiceMissing = runtime.ErrMemoryRetrievalServiceMissing + ErrKGRetrievalServiceMissing = runtime.ErrKGRetrievalServiceMissing ) -var ( - memoryRetrievalServiceMu sync.RWMutex - memoryRetrievalServiceImpl MemoryRetrievalService = stubMemoryRetrievalService{} -) +func SetRetrievalService(svc RetrievalService) { runtime.SetRetrievalService(svc) } +func GetRetrievalService() RetrievalService { return runtime.GetRetrievalService() } -func SetRetrievalService(svc RetrievalService) { - retrievalServiceMu.Lock() - defer retrievalServiceMu.Unlock() - if svc == nil { - retrievalServiceImpl = stubRetrievalService{} - return - } - retrievalServiceImpl = svc -} +func SetMemoryRetrievalService(svc MemoryRetrievalService) { runtime.SetMemoryRetrievalService(svc) } +func GetMemoryRetrievalService() MemoryRetrievalService { return runtime.GetMemoryRetrievalService() } -func GetRetrievalService() RetrievalService { - retrievalServiceMu.RLock() - defer retrievalServiceMu.RUnlock() - return retrievalServiceImpl -} - -func SetMemoryRetrievalService(svc MemoryRetrievalService) { - memoryRetrievalServiceMu.Lock() - defer memoryRetrievalServiceMu.Unlock() - if svc == nil { - memoryRetrievalServiceImpl = stubMemoryRetrievalService{} - return - } - memoryRetrievalServiceImpl = svc -} - -func GetMemoryRetrievalService() MemoryRetrievalService { - memoryRetrievalServiceMu.RLock() - defer memoryRetrievalServiceMu.RUnlock() - return memoryRetrievalServiceImpl -} - -type stubRetrievalService struct{} - -func (stubRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { - return nil, ErrRetrievalServiceMissing -} - -type stubMemoryRetrievalService struct{} - -func (stubMemoryRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { - return nil, ErrMemoryRetrievalServiceMissing -} - -// simpleRetrievalService is a deterministic test implementation that returns -// synthetic chunks based on the query. -type simpleRetrievalService struct{} - -func (simpleRetrievalService) Search(_ context.Context, _ *gorm.DB, req RetrievalRequest) ([]RetrievalChunk, error) { - if req.Query == "" { - return nil, nil - } - topN := req.TopN - if topN <= 0 { - topN = 8 - } - // Cap topN to a sane upper bound so a hostile canvas can't force - // a giant preallocation here. Real callers honor this cap; the - // production service has its own server-side limits as well. - const maxSimpleTopN = 1024 - if topN > maxSimpleTopN { - topN = maxSimpleTopN - } - // codeql[go/uncontrolled-allocation-size] False positive: topN - // is bounded to maxSimpleTopN (1024) above, so the resulting - // slice cannot exceed ~1 MiB (chunk items are small structs). - chunks := make([]RetrievalChunk, 0, topN) - for i := 0; i < topN && i < 3; i++ { - chunks = append(chunks, RetrievalChunk{ - ID: fmt.Sprintf("simple-%d", i), - Content: fmt.Sprintf("Chunk %d matching %q", i, req.Query), - DocumentID: "simple-doc", - Score: 0.9 - float64(i)*0.1, - }) - } - return chunks, nil -} +func SetKGRetrievalService(svc KGRetrievalService) { runtime.SetKGRetrievalService(svc) } +func GetKGRetrievalService() KGRetrievalService { return runtime.GetKGRetrievalService() } // SetSimpleRetrievalService installs deterministic synthetic retrieval for // tests and local demos. -func SetSimpleRetrievalService() { - SetRetrievalService(simpleRetrievalService{}) -} - -// ErrKGRetrievalServiceMissing is returned when the agent's -// RetrievalTool dispatches use_kg=true but no KGRetrievalService -// has been registered via SetKGRetrievalService. This is the -// expected "kg not yet wired" state — distinct from -// ErrRetrievalServiceMissing (which signals the nlp adapter is -// un-wired). -var ErrKGRetrievalServiceMissing = errors.New( - "GraphRAG (kg) retrieval service not yet wired — " + - "call tool.SetKGRetrievalService(tool.NewKGRetrievalAdapter(...)) at boot", -) - -var ErrMemoryRetrievalServiceMissing = errors.New( - "memory retrieval service not registered", -) - -var ( - kgRetrievalServiceMu sync.RWMutex - kgRetrievalServiceImpl KGRetrievalService = stubKGRetrievalService{} -) - -// SetKGRetrievalService installs the GraphRAG adapter. Passing -// nil reverts to the stub that returns ErrKGRetrievalServiceMissing. -// Idempotent: safe to call from cmd/server_main.go once at boot -// and from tests that want to swap the impl. -func SetKGRetrievalService(svc KGRetrievalService) { - kgRetrievalServiceMu.Lock() - defer kgRetrievalServiceMu.Unlock() - if svc == nil { - kgRetrievalServiceImpl = stubKGRetrievalService{} - return - } - kgRetrievalServiceImpl = svc -} - -// GetKGRetrievalService returns the registered KGRetrievalService. -// Always non-nil — defaults to the stub. -func GetKGRetrievalService() KGRetrievalService { - kgRetrievalServiceMu.RLock() - defer kgRetrievalServiceMu.RUnlock() - return kgRetrievalServiceImpl -} - -type stubKGRetrievalService struct{} - -func (stubKGRetrievalService) Search(_ context.Context, _ *gorm.DB, _ RetrievalRequest) ([]RetrievalChunk, error) { - return nil, ErrKGRetrievalServiceMissing -} +func SetSimpleRetrievalService() { runtime.SetSimpleRetrievalService() } diff --git a/internal/agent/tool/retrieval_wiring_test.go b/internal/agent/tool/retrieval_wiring_test.go index 347b5661e4..db8e5364bc 100644 --- a/internal/agent/tool/retrieval_wiring_test.go +++ b/internal/agent/tool/retrieval_wiring_test.go @@ -64,8 +64,12 @@ func TestRetrievalTool_SimpleServiceReturnsChunks(t *testing.T) { // TestSimpleRetrievalService_EmptyQuery: empty query → no chunks, // no error. func TestSimpleRetrievalService_EmptyQuery(t *testing.T) { + prev := GetRetrievalService() + t.Cleanup(func() { SetRetrievalService(prev) }) + SetSimpleRetrievalService() + ctx := t.Context() - chunks, err := simpleRetrievalService{}.Search(ctx, nil, RetrievalRequest{Query: ""}) + chunks, err := GetRetrievalService().Search(ctx, nil, RetrievalRequest{Query: ""}) if err != nil { t.Errorf("err: %v", err) } @@ -77,8 +81,12 @@ func TestSimpleRetrievalService_EmptyQuery(t *testing.T) { // TestSimpleRetrievalService_RespectsTopN: top_n=1 → exactly 1 // chunk. func TestSimpleRetrievalService_RespectsTopN(t *testing.T) { + prev := GetRetrievalService() + t.Cleanup(func() { SetRetrievalService(prev) }) + SetSimpleRetrievalService() + ctx := t.Context() - chunks, err := simpleRetrievalService{}.Search(ctx, nil, RetrievalRequest{ + chunks, err := GetRetrievalService().Search(ctx, nil, RetrievalRequest{ Query: "x", TopN: 1, }) diff --git a/internal/common/logger.go b/internal/common/logger.go index d18729f55e..185fe1b42a 100644 --- a/internal/common/logger.go +++ b/internal/common/logger.go @@ -18,6 +18,7 @@ package common import ( "bytes" + "context" "fmt" "io" "log" @@ -50,31 +51,37 @@ var ( // (stdout only). When Path is set, the file is written under ./logs/ // and rotated by lumberjack according to MaxSize / MaxBackups / MaxAge / Compress. // -// Numeric zero values are replaced with defaults (100 MB / 10 / 30 days) inside -// Init. Compress is left as the caller-provided value; the project default is -// applied by callers (see resolveCompress) so that "not set" can be distinguished -// from "explicitly false" via the *bool LogConfig.Compress field. +// Numeric zero values (MaxSize/MaxBackups/MaxAge) are replaced with defaults +// (100 MB / 10 / 30 days) inside Init. Compress is a *bool so that "not set" +// (nil) can be distinguished from "explicitly false"; when nil it defaults to +// DefaultLogCompress (true). type FileOutput struct { Filename string Path string MaxSize int MaxBackups int MaxAge int - Compress bool + Compress *bool } const ( - defaultMaxSizeMB = 100 - defaultMaxBackups = 10 - defaultMaxAgeDays = 30 - cyanLogMarker = "[[RAGFLOW_CYAN_LOG]]" - greenLogMarker = "[[RAGFLOW_GREEN_LOG]]" - redLogMarker = "[[RAGFLOW_RED_LOG]]" - resetLogMarker = "[[RAGFLOW_RESET_LOG]]" - ansiBrightCyan = "\x1b[96m" - ansiGreen = "\x1b[32m" - ansiRed = "\x1b[31m" - ansiReset = "\x1b[0m" + cyanLogMarker = "[[RAGFLOW_CYAN_LOG]]" + greenLogMarker = "[[RAGFLOW_GREEN_LOG]]" + redLogMarker = "[[RAGFLOW_RED_LOG]]" + resetLogMarker = "[[RAGFLOW_RESET_LOG]]" + ansiBrightCyan = "\x1b[96m" + ansiGreen = "\x1b[32m" + ansiRed = "\x1b[31m" + ansiReset = "\x1b[0m" + // DefaultLogMaxSizeMB is the default rotation threshold (lumberjack + // MaxSize is in MB, not bytes). + DefaultLogMaxSizeMB = 100 + // DefaultLogMaxBackups is the default number of rotated files retained. + DefaultLogMaxBackups = 10 + // DefaultLogMaxAgeDays is the default retention window for rotated files. + DefaultLogMaxAgeDays = 30 + // DefaultLogCompress is the project default for gzipping rotated files. + DefaultLogCompress = true ) type coloredLineWriteSyncer struct { @@ -176,45 +183,65 @@ func InitLogger(level string, file FileOutput, serviceName string) error { maxSize := file.MaxSize if maxSize <= 0 { - maxSize = defaultMaxSizeMB + maxSize = DefaultLogMaxSizeMB } maxBackups := file.MaxBackups if maxBackups <= 0 { - maxBackups = defaultMaxBackups + maxBackups = DefaultLogMaxBackups } maxAge := file.MaxAge if maxAge <= 0 { - maxAge = defaultMaxAgeDays + maxAge = DefaultLogMaxAgeDays } - ljLogger := &lumberjack.Logger{ - Filename: filepath.Join(file.Path, file.Filename), - MaxSize: maxSize, - MaxBackups: maxBackups, - MaxAge: maxAge, - Compress: file.Compress, - LocalTime: true, + compress := DefaultLogCompress + if file.Compress != nil { + compress = *file.Compress } stdoutSyncer := zapcore.AddSync(os.Stdout) - fileSyncer := zapcore.AddSync(ljLogger) + syncers := []zapcore.WriteSyncer{stdoutSyncer} + // File sink only when a destination is actually configured. The cmd/* + // entry points init the logger TWICE: a pre-config temporary logger + // (before the port is known, e.g. "api_server" → logs/api_server.log) + // and the real one after server.Init renames it (e.g. + // "api_server_9384" → logs/api_server_9384.log). The temporary pass now + // passes an empty FileOutput and stays stdout-only, so the pre-config + // startup window no longer litters a second, orphaned log file next to + // the real one (and ragflow-cli no longer drops a lumberjack file into + // os.TempDir()). + if file.Path != "" && file.Filename != "" { + ljLogger := &lumberjack.Logger{ + Filename: filepath.Join(file.Path, file.Filename), + MaxSize: maxSize, + MaxBackups: maxBackups, + MaxAge: maxAge, + Compress: compress, + LocalTime: true, + } + syncers = append(syncers, zapcore.AddSync(ljLogger)) + } + var core zapcore.Core if IsLLMDebugEnabled() { - core = zapcore.NewTee( + cores := []zapcore.Core{ zapcore.NewCore( zapcore.NewConsoleEncoder(encoderConfig), coloredLineWriteSyncer{WriteSyncer: stdoutSyncer, color: true}, atomicLevel, ), - zapcore.NewCore( + } + if len(syncers) > 1 { + cores = append(cores, zapcore.NewCore( zapcore.NewConsoleEncoder(encoderConfig), - coloredLineWriteSyncer{WriteSyncer: fileSyncer}, + coloredLineWriteSyncer{WriteSyncer: syncers[1]}, atomicLevel, - ), - ) + )) + } + core = zapcore.NewTee(cores...) } else { core = zapcore.NewCore( zapcore.NewConsoleEncoder(encoderConfig), - zap.CombineWriteSyncers(stdoutSyncer, fileSyncer), + zap.CombineWriteSyncers(syncers...), atomicLevel, ) } @@ -326,6 +353,81 @@ func (stdLogRouter) Write(p []byte) (int, error) { return len(p), nil } +// --- per-request log correlation ------------------------------------------- +// +// One conversation turn (a chat completion request) fans out into dozens of +// log lines across packages — agent, delivery gate, auditor, tools, retrieval, +// token usage — and concurrent benchmark questions interleave them. Attaching +// the session id to the request context and reading it back in the Ctx-variant +// log helpers below lets a single `grep session_id=` reconstruct one +// turn's full trail (the q71 postmortem had to reconstruct it from timestamps). + +type ctxKey int + +const sessionIDCtxKey ctxKey = iota + +// WithSessionID returns a context that tags every Ctx-variant log call with +// the conversation turn's session id. Empty ids are a no-op so callers do not +// need to guard. +func WithSessionID(ctx context.Context, sessionID string) context.Context { + if strings.TrimSpace(sessionID) == "" { + return ctx + } + return context.WithValue(ctx, sessionIDCtxKey, strings.TrimSpace(sessionID)) +} + +// SessionIDFromContext extracts the correlation id ("" when absent). +func SessionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + id, _ := ctx.Value(sessionIDCtxKey).(string) + return id +} + +func sessionFields(ctx context.Context) []zap.Field { + if id := SessionIDFromContext(ctx); id != "" { + return []zap.Field{zap.String("session_id", id)} + } + return nil +} + +// InfoCtx is Info plus the session_id correlation field when ctx carries one. +func InfoCtx(ctx context.Context, msg string, fields ...zap.Field) { + if Logger == nil { + return + } + Logger.Info(msg, append(sessionFields(ctx), fields...)...) +} + +// WarnCtx is Warn plus the session_id correlation field when ctx carries one. +func WarnCtx(ctx context.Context, msg string, fields ...zap.Field) { + if Logger == nil { + return + } + Logger.Warn(msg, append(sessionFields(ctx), fields...)...) +} + +// DebugCtx is Debug plus the session_id correlation field when ctx carries one. +func DebugCtx(ctx context.Context, msg string, fields ...zap.Field) { + if Logger == nil { + return + } + Logger.Debug(msg, append(sessionFields(ctx), fields...)...) +} + +// ErrorCtx is Error plus the session_id correlation field when ctx carries one. +func ErrorCtx(ctx context.Context, msg string, err error, fields ...zap.Field) { + if Logger == nil { + return + } + detail := fmt.Sprintf("%s, %v", msg, err) + if IsDebugEnabled() { + detail = fmt.Sprintf("%s, %+v", msg, err) + } + Logger.Error(detail, append(sessionFields(ctx), fields...)...) +} + // IsDebugEnabled returns true if debug logging is enabled. func IsDebugEnabled() bool { return atomicLevel.Enabled(zapcore.DebugLevel) diff --git a/internal/common/logger_ctx_test.go b/internal/common/logger_ctx_test.go new file mode 100644 index 0000000000..99c2dd71c9 --- /dev/null +++ b/internal/common/logger_ctx_test.go @@ -0,0 +1,115 @@ +// +// 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 common + +import ( + "context" + "testing" + + "go.uber.org/zap" + "go.uber.org/zap/zapcore" + "go.uber.org/zap/zaptest/observer" +) + +// withTestLogger swaps in an observer-backed logger and restores the old one. +// atomicLevel must be initialized alongside Logger: IsDebugEnabled reads it, +// and its zero value panics (in production InitLog always runs first, but a +// unit test starts from the zero package state). +func withTestLogger(t *testing.T) *observer.ObservedLogs { + t.Helper() + // Debug-level observer so the DebugCtx call in the assertions is captured + // too; production level filtering is orthogonal to field attachment. + core, logs := observer.New(zapcore.DebugLevel) + old := Logger + oldLevel := atomicLevel + Logger = zap.New(core) + atomicLevel = zap.NewAtomicLevelAt(zapcore.DebugLevel) + t.Cleanup(func() { Logger = old; atomicLevel = oldLevel }) + return logs +} + +// TestCtxLogOmitsSessionFieldWhenAbsent pins the nil-slice contract of +// sessionFields: a context without a session id must add NO session_id field +// at all — not an empty string, and certainly not the string "nil" (a nil +// []zap.Field appended into variadic fields is just the zero-length prefix, +// never a materialized field). +func TestCtxLogOmitsSessionFieldWhenAbsent(t *testing.T) { + logs := withTestLogger(t) + + InfoCtx(context.Background(), "no-session") + DebugCtx(WithSessionID(context.Background(), "abc"), "with-session") + WarnCtx(context.Background(), "bare-warn") + ErrorCtx(context.Background(), "bare-error", context.DeadlineExceeded) + InfoCtx(WithSessionID(context.Background(), "def"), "with-extra", zap.Int("n", 1)) + + entries := logs.All() + if len(entries) != 5 { + t.Fatalf("entries = %d, want 5", len(entries)) + } + + type tc struct { + msg string + wantID string // "" means the field must be absent + } + cases := []tc{ + {"no-session", ""}, + {"with-session", "abc"}, + {"bare-warn", ""}, + {"bare-error", ""}, + {"with-extra", "def"}, + } + for i, c := range cases { + entry := entries[i] + // ErrorCtx folds the error into the message, mirroring common.Error. + want := c.msg + if c.msg == "bare-error" { + want = "bare-error, context deadline exceeded" + } + if entry.Message != want { + t.Fatalf("entry %d message = %q, want %q", i, entry.Message, want) + } + got, present := entry.ContextMap()["session_id"] + if c.wantID == "" { + if present { + t.Errorf("entry %d (%q): session_id = %v, want the field absent", i, c.msg, got) + } + continue + } + if !present || got != c.wantID { + t.Errorf("entry %d (%q): session_id = (%v, %v), want %q", i, c.msg, got, present, c.wantID) + } + } +} + +// TestSessionIDFromContext covers the accessor's edge cases directly. +func TestSessionIDFromContext(t *testing.T) { + if got := SessionIDFromContext(nil); got != "" { + t.Errorf("nil context: got %q, want empty", got) + } + if got := SessionIDFromContext(context.Background()); got != "" { + t.Errorf("plain context: got %q, want empty", got) + } + // Blank and whitespace-only ids are a no-op, per WithSessionID's guard. + for _, blank := range []string{"", " ", "\t"} { + ctx := WithSessionID(context.Background(), blank) + if got := SessionIDFromContext(ctx); got != "" { + t.Errorf("WithSessionID(%q): got %q, want the ctx untouched", blank, got) + } + } + if got := SessionIDFromContext(WithSessionID(context.Background(), " s1 ")); got != "s1" { + t.Errorf("trimmed id: got %q, want %q", got, "s1") + } +} diff --git a/internal/common/model_assets.go b/internal/common/model_assets.go new file mode 100644 index 0000000000..e36e0e9a59 --- /dev/null +++ b/internal/common/model_assets.go @@ -0,0 +1,76 @@ +// +// 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 common + +import ( + "path/filepath" + "strings" +) + +// EnvModelAssetsDir points at a directory that plays the role of `ragflow_deps/`: +// downloaded model assets stored in the HuggingFace repo layout, i.e. +// `/huggingface.co//` (plus loose files that are not HF repos, such as +// the cl100k BPE table). It is how a deployment keeps model assets outside the checkout - +// a mounted volume, a read-only image layer, or a shared cache - instead of relying on +// the working directory the process happens to have. +// +// The name is deliberately not embedding-specific: the layout is shared by every +// downloaded model asset (embedding tokenizers, DeepDoc weights, and whatever comes +// next), so any loader that reads `huggingface.co//` should consult +// ModelAssetCandidates rather than inventing its own variable. +const EnvModelAssetsDir = "MODEL_ASSETS_DIR" + +// ModelAssetCandidates returns the paths to try, in priority order, for one asset when +// MODEL_ASSETS_DIR is set. `relative` is the shipped name - either +// "ragflow_deps/huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model" or a bare +// "cl100k_base.tiktoken" - and the candidates let an operator point at any of the three +// sensible roots: +// +// /huggingface.co// the ragflow_deps-like root +// // a huggingface.co-like root +// / a flat directory holding the files +// +// Returns nil when the variable is unset, leaving the caller's own search untouched. +func ModelAssetCandidates(relative string) []string { + root := strings.TrimSpace(GetEnv(EnvModelAssetsDir)) + if root == "" { + return nil + } + rel := strings.TrimPrefix(filepath.ToSlash(relative), "ragflow_deps/") + out := make([]string, 0, 3) + seen := make(map[string]struct{}, 3) + add := func(p string) { + if _, dup := seen[p]; dup { + return + } + seen[p] = struct{}{} + out = append(out, p) + } + add(filepath.Join(root, filepath.FromSlash(rel))) + if rest, ok := strings.CutPrefix(rel, "huggingface.co/"); ok { + add(filepath.Join(root, filepath.FromSlash(rest))) + // The flat layout documented above, for an operator who kept the files + // without their repo directories. It is last because a repo-shipped name + // like tokenizer.json collides across repos, so the repo layouts win. + add(filepath.Join(root, filepath.Base(rel))) + } else { + // A loose file (the cl100k table) has no repo directory: its root-level + // candidate is the one already added above, so this cannot add a second. + add(filepath.Join(root, filepath.Base(rel))) + } + return out +} diff --git a/internal/common/model_assets_test.go b/internal/common/model_assets_test.go new file mode 100644 index 0000000000..f45019f3aa --- /dev/null +++ b/internal/common/model_assets_test.go @@ -0,0 +1,69 @@ +// +// 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 common + +import ( + "path/filepath" + "testing" +) + +func TestModelAssetCandidatesOnlyAppliesWhenTheEnvIsSet(t *testing.T) { + t.Setenv(EnvModelAssetsDir, "") + if got := ModelAssetCandidates("ragflow_deps/huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model"); got != nil { + t.Fatalf("expected no candidates without %s, got %v", EnvModelAssetsDir, got) + } + t.Setenv(EnvModelAssetsDir, " ") + if got := ModelAssetCandidates("cl100k_base.tiktoken"); got != nil { + t.Fatalf("expected whitespace-only %s to be ignored, got %v", EnvModelAssetsDir, got) + } +} + +// The three layouts an operator may have in mind when pointing MODEL_ASSETS_DIR at a +// tree: the ragflow_deps-like root, a huggingface.co-like root, or a flat directory. +func TestModelAssetCandidatesCoverEverySensibleRoot(t *testing.T) { + root := filepath.Join(string(filepath.Separator)+"assets", "models") + t.Setenv(EnvModelAssetsDir, root) + + got := ModelAssetCandidates("ragflow_deps/huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model") + want := []string{ + filepath.Join(root, "huggingface.co", "BAAI", "bge-m3", "sentencepiece.bpe.model"), + filepath.Join(root, "BAAI", "bge-m3", "sentencepiece.bpe.model"), + // The flat layout: the operator kept the file without its repo directory. + filepath.Join(root, "sentencepiece.bpe.model"), + } + if len(got) != len(want) { + t.Fatalf("candidates = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Errorf("candidate %d = %q, want %q", i, got[i], want[i]) + } + } + + // The same asset named without the ragflow_deps/ prefix must resolve identically. + bare := ModelAssetCandidates("huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model") + if len(bare) != len(want) || bare[0] != want[0] { + t.Errorf("prefix form differs: %v vs %v", bare, want) + } + + // A file that is not an HF repo (the cl100k table) sits at the root itself, + // and its single candidate must not be repeated. + flat := ModelAssetCandidates("ragflow_deps/cl100k_base.tiktoken") + if len(flat) != 1 || flat[0] != filepath.Join(root, "cl100k_base.tiktoken") { + t.Errorf("flat asset candidates = %v, want exactly one root-level candidate", flat) + } +} diff --git a/internal/dao/ingestion_task_test.go b/internal/dao/ingestion_task_test.go index b914182335..3ead63809e 100644 --- a/internal/dao/ingestion_task_test.go +++ b/internal/dao/ingestion_task_test.go @@ -568,3 +568,50 @@ func TestIngestionTaskDAOHasDatasetStatusIndex(t *testing.T) { t.Fatal("expected composite dataset/status index on ingestion_task") } } + +// TestIngestionTaskDAOGetByDocumentID_ReturnsTask verifies the plain +// single-row lookup: with no task recorded for the document GetByDocumentID +// returns nil, and once one exists it returns that task. The test DB's +// AutoMigrate enforces the unique index on document_id, so this stays a +// one-row case — the ordering behaviour (newest first, id as tie-breaker) is +// covered by TestIngestionTaskDAOGetByDocumentIDUsesLatestTask and +// TestIngestionTaskDAOGetByDocumentIDUsesIDAsTieBreaker, both of which drop +// that index so they can create duplicate rows. +func TestIngestionTaskDAOGetByDocumentID_ReturnsTask(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + ctx := t.Context() + dao := NewIngestionTaskDAO() + + // No task yet. + task, err := dao.GetByDocumentID(ctx, db, "doc-1") + if err != nil { + t.Fatalf("GetByDocumentID empty: %v", err) + } + if task != nil { + t.Fatalf("expected nil for doc-1, got %+v", task) + } + + // Create one task. + task = &entity.IngestionTask{ + ID: "task-1", + UserID: "user-1", + DocumentID: "doc-1", + DatasetID: "kb-1", + Status: common.CREATED, + } + if err := db.Create(task).Error; err != nil { + t.Fatalf("create task: %v", err) + } + + got, err := dao.GetByDocumentID(ctx, db, "doc-1") + if err != nil { + t.Fatalf("GetByDocumentID: %v", err) + } + if got == nil || got.ID != "task-1" { + t.Fatalf("doc-1 task = %+v, want task-1", got) + } +} diff --git a/internal/dao/tenant_model.go b/internal/dao/tenant_model.go index 2a237cf195..6e9e685010 100644 --- a/internal/dao/tenant_model.go +++ b/internal/dao/tenant_model.go @@ -99,6 +99,25 @@ func (dao *TenantModelDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []stri return models, nil } +// GetActiveChatModelsByProviderAndInstanceIDs returns every ACTIVE model whose +// model_type bitmask covers `modelType`, scoped to the given provider / +// instance pairs. Used to enumerate a tenant's chat-capable models for the +// EinoChatModel failover chain: each returned row's ID doubles as the +// "model ref" that ModelProviderService.ResolveModelConfig accepts. +func (dao *TenantModelDAO) GetActiveModelsByProviderAndInstanceIDsAndType(ctx context.Context, db *gorm.DB, providerIDs, instanceIDs []string, modelType int) ([]*entity.TenantModel, error) { + var models []*entity.TenantModel + q := db.WithContext(ctx). + Where("provider_id IN ?", providerIDs). + Where("instance_id IN ?", instanceIDs). + Where("model_type & ? != 0", modelType). + Where("status = ?", "active") + err := q.Find(&models).Error + if err != nil { + return nil, err + } + return models, nil +} + func (dao *TenantModelDAO) GetModelByProviderIDAndInstanceIDAndModelName(ctx context.Context, db *gorm.DB, providerID, instanceID, modelName string) (*entity.TenantModel, error) { var model entity.TenantModel err := db.WithContext(ctx).Where("provider_id = ? AND instance_id = ? AND model_name = ?", providerID, instanceID, modelName).First(&model).Error diff --git a/internal/engine/elasticsearch/chunk.go b/internal/engine/elasticsearch/chunk.go index f2e85c35cb..c99fa737c3 100644 --- a/internal/engine/elasticsearch/chunk.go +++ b/internal/engine/elasticsearch/chunk.go @@ -153,7 +153,31 @@ func (e *Engine) CreateChunkStore(ctx context.Context, baseName, datasetID strin // InsertChunks inserts chunks into a chunk index // If a chunk with the same id + doc_id + kb_id already exists, it will be updated with the new value +// +// It waits for an index refresh before returning, so a caller that reads the +// chunks straight back (the chunk APIs, the debug endpoints) sees them right +// away. Ingestion must use InsertChunksNoRefresh instead - see that method for +// why the wait is the wrong trade there. func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string) ([]string, error) { + return e.insertChunks(ctx, chunks, baseName, datasetID, "wait_for") +} + +// InsertChunksNoRefresh inserts chunks without waiting for an index refresh: +// they become visible on the index's normal refresh cycle (the KB index ships +// with refresh_interval=1000ms), so the difference is at most a second of +// visibility - against a measured ~0.5s median (1.0s p90, 5.0s max) of pure +// latency per write when waiting. +// +// This is the ingestion path's contract, and it matches Python's: the Python +// ingestion inserts chunks with refresh=False (rag/svr/task_executor_refactor/ +// chunk_service.py:386 and :423), while its API default stays "wait_for". +func (e *Engine) InsertChunksNoRefresh(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string) ([]string, error) { + return e.insertChunks(ctx, chunks, baseName, datasetID, "") +} + +// insertChunks is the shared bulk-insert body. An empty refresh omits the +// parameter, which leaves the decision to the index (no forced refresh). +func (e *Engine) insertChunks(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string, refresh string) ([]string, error) { common.Info("ElasticsearchConnection.InsertChunks called", zap.String("index_name", baseName), zap.Int("chunkCount", len(chunks))) if len(chunks) == 0 { @@ -171,8 +195,83 @@ func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface } } - // Build bulk request body with index operations (upsert behavior: insert if not exists, update if exists) + // Bulk requests are executed in bounded batches so a single large document + // (thousands of chunks) never sends a payload above the Elasticsearch HTTP + // body limit (default http.max_content_length = 100MB), which surfaces as a + // 413 Request Entity Too Large. We flush a batch by chunk count AND byte + // size to stay well within the limit. + const ( + bulkBatchChunks = 500 + bulkBatchBytes = 20 * 1024 * 1024 // 20MB, comfortably under the 100MB default + ) + + flush := func(buf *bytes.Buffer) error { + if buf.Len() == 0 { + return nil + } + req := esapi.BulkRequest{ + Body: bytes.NewReader(buf.Bytes()), + Refresh: refresh, + } + res, err := req.Do(ctx, e.client) + if err != nil { + common.Error("Failed to execute bulk request", err) + return fmt.Errorf("failed to execute bulk request: %w", err) + } + defer res.Body.Close() + + if res.IsError() { + bodyBytes, _ := io.ReadAll(res.Body) + common.Sugar.Errorw("Elasticsearch bulk request returned error", "status", res.Status(), "body", string(bodyBytes)) + return fmt.Errorf("elasticsearch bulk request returned error: %s, body: %s", res.Status(), string(bodyBytes)) + } + + // Parse bulk response + var bulkResponse map[string]interface{} + if err := json.NewDecoder(res.Body).Decode(&bulkResponse); err != nil { + common.Error("Failed to parse bulk response", err) + return fmt.Errorf("failed to parse bulk response: %w", err) + } + + // Check for errors in bulk response. ES reports a 200 at the + // top level even when individual bulk items fail, so we must + // inspect the items and surface the reasons instead of returning + // success while the chunks were silently dropped. + if hasErrs, ok := bulkResponse["errors"].(bool); ok && hasErrs { + var reasons []string + if items, ok := bulkResponse["items"].([]interface{}); ok { + for _, it := range items { + im, ok := it.(map[string]interface{}) + if !ok { + continue + } + for _, op := range im { + om, ok := op.(map[string]interface{}) + if !ok { + continue + } + if errObj, ok := om["error"].(map[string]interface{}); ok { + if reason, ok := errObj["reason"].(string); ok && reason != "" { + reasons = append(reasons, reason) + } + } + } + } + } + if len(reasons) > 5 { + reasons = reasons[:5] + } + common.Error("Elasticsearch bulk request had item errors", + fmt.Errorf("index %s: %d failed items; first reasons: %v", baseName, len(reasons), reasons)) + return fmt.Errorf("elasticsearch bulk request had %d item errors (index %s); first reasons: %v", + len(reasons), baseName, reasons) + } + return nil + } + + // Build and execute bulk request bodies in bounded batches. var buf bytes.Buffer + batchChunkCount := 0 for _, doc := range chunks { if isMemoryIndex { // Memory messages arrive with logical fields; map them to @@ -209,66 +308,20 @@ func (e *Engine) InsertChunks(ctx context.Context, chunks []map[string]interface if err := jsonIterator.NewEncoder(&buf).Encode(docCopy); err != nil { return nil, fmt.Errorf("failed to encode document: %w", err) } - } - // Execute bulk request with refresh="wait_for" - req := esapi.BulkRequest{ - Body: bytes.NewReader(buf.Bytes()), - Refresh: "wait_for", - } - - res, err := req.Do(ctx, e.client) - if err != nil { - common.Error("Failed to execute bulk request", err) - return nil, fmt.Errorf("failed to execute bulk request: %w", err) - } - defer res.Body.Close() - - if res.IsError() { - bodyBytes, _ := io.ReadAll(res.Body) - common.Sugar.Errorw("Elasticsearch bulk request returned error", "status", res.Status(), "body", string(bodyBytes)) - return nil, fmt.Errorf("elasticsearch bulk request returned error: %s, body: %s", res.Status(), string(bodyBytes)) - } - - // Parse bulk response - var bulkResponse map[string]interface{} - if err := json.NewDecoder(res.Body).Decode(&bulkResponse); err != nil { - common.Error("Failed to parse bulk response", err) - return nil, fmt.Errorf("failed to parse bulk response: %w", err) - } - - // Check for errors in bulk response. ES reports a 200 at the - // top level even when individual bulk items fail, so we must - // inspect the items and surface the reasons instead of returning - // success while the chunks were silently dropped. - if errors, ok := bulkResponse["errors"].(bool); ok && errors { - var reasons []string - if items, ok := bulkResponse["items"].([]interface{}); ok { - for _, it := range items { - im, ok := it.(map[string]interface{}) - if !ok { - continue - } - for _, op := range im { - om, ok := op.(map[string]interface{}) - if !ok { - continue - } - if errObj, ok := om["error"].(map[string]interface{}); ok { - if reason, ok := errObj["reason"].(string); ok && reason != "" { - reasons = append(reasons, reason) - } - } - } + batchChunkCount++ + if batchChunkCount >= bulkBatchChunks || buf.Len() >= bulkBatchBytes { + // Flush the current batch before adding the next chunk. + if err := flush(&buf); err != nil { + return nil, err } + buf.Reset() + batchChunkCount = 0 } - if len(reasons) > 5 { - reasons = reasons[:5] - } - common.Error("Elasticsearch bulk request had item errors", - fmt.Errorf("index %s: %d failed items; first reasons: %v", baseName, len(reasons), reasons)) - return nil, fmt.Errorf("elasticsearch bulk request had %d item errors (index %s); first reasons: %v", - len(reasons), baseName, reasons) + } + // Flush the final partial batch. + if err := flush(&buf); err != nil { + return nil, err } common.Info("ElasticsearchConnection.InsertChunks result", zap.String("index_name", baseName), zap.Int("count", len(chunks))) @@ -1180,8 +1233,10 @@ func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.S // Build query body with text match and/or knn match queryBody := make(map[string]interface{}) + hasVectorMatch := matchDense != nil && len(matchDense.EmbeddingData) > 0 + if matchText != nil { - textQuery := buildQueryStringQuery(matchText, vectorSimilarityWeight, isSkillIndex, isMemoryIndex) + textQuery := buildQueryStringQuery(matchText, isSkillIndex, isMemoryIndex) if boolQuery != nil { if boolMap, ok := boolQuery["bool"].(map[string]interface{}); ok { if must, ok := boolMap["must"].([]interface{}); ok { @@ -1190,14 +1245,19 @@ func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.S } else { boolMap["must"] = []interface{}{textQuery} } - boolMap["boost"] = 1.0 - vectorSimilarityWeight + // Weighted-sum fusion scales the text leg by (1 - w) against the + // knn leg's w; a text-only query has no such pairing, so the + // boolean filter wrapper must not apply the constant boost — + // otherwise every BM25 score is silently multiplied by 0.5. + if hasVectorMatch { + boolMap["boost"] = 1.0 - vectorSimilarityWeight + } } } else { boolQuery = textQuery } } - hasVectorMatch := matchDense != nil && len(matchDense.EmbeddingData) > 0 if hasVectorMatch { if isMemoryIndex { if err := e.ensureMemoryMessageSearchVectorMappings(ctx, req.IndexNames, matchDense.VectorColumnName, len(matchDense.EmbeddingData)); err != nil { @@ -1237,7 +1297,13 @@ func (e *Engine) Search(ctx context.Context, req *types.SearchRequest) (*types.S } queryBody["knn"] = knnQuery - if boolQuery != nil { + // A DENSE-ONLY request (no text expression: the vector-only search and + // the knowledge-compile lookups) leaves the top-level query unset. The + // same scope conditions already ride inside knn.filter, so repeating + // them as a scored query adds one scoring pass per search AND blends + // that score into _score — which is exactly the number a vector-only + // caller reads as the cosine similarity. + if boolQuery != nil && matchText != nil { queryBody["query"] = boolQuery } } else if boolQuery != nil { @@ -1903,6 +1969,13 @@ func buildBoolQueryFromCondition(filter map[string]interface{}, kbIDs []string, map[string]interface{}{"terms": map[string]interface{}{"_id": listVal}}, ) } else if strListVal, ok := v.([]string); ok && len(strListVal) > 0 { + // A typed []string MUST be handled here: the generic loop + // below skips the "id" key, so without this branch a caller + // passing map[string]interface{}{"id": []string{...}} builds a + // query with NO id filter and the scoped read silently + // degenerates into "fetch up to Limit chunks of the document" + // (observed via list_chunks: an 11-chunk window returned 3.8MB + // because the whole 4171-chunk document was fetched instead). shouldClauses = append(shouldClauses, map[string]interface{}{"terms": map[string]interface{}{"id": strListVal}}, map[string]interface{}{"terms": map[string]interface{}{"_id": strListVal}}, @@ -1984,7 +2057,16 @@ func buildBoolQueryFromCondition(filter map[string]interface{}, kbIDs []string, // buildQueryStringQuery builds a query_string query from MatchTextExpr // When isSkillIndex is true, uses skill-specific fields (name_tks, tags_tks, etc.) // Otherwise uses document fields (title_tks, content_ltks, etc.) -func buildQueryStringQuery(matchText *types.MatchTextExpr, vectorSimilarityWeight float64, isSkillIndex, isMemoryIndex bool) map[string]interface{} { +// +// The MatchingText is lowercased: the *_tks/*_ltks fields it targets are +// tokenized with the whitespace analyzer and RAGFlow's ingest pipeline stores +// pre-lowercased tokens there, but the whitespace search analyzer does NOT +// lowercase the query. Capitalized terms (any proper noun: "Ross Feldner", +// "New Age Graphics") therefore silently match nothing, and a keyword query +// consisting only of proper nouns returns zero hits — observed live when a +// perfect pivot query ("Isabel Wood co-lead Ross Feldner ...") dead-ended a +// benchmark question that was answerable from the corpus (q156). +func buildQueryStringQuery(matchText *types.MatchTextExpr, isSkillIndex, isMemoryIndex bool) map[string]interface{} { if matchText == nil { return nil } @@ -2024,7 +2106,7 @@ func buildQueryStringQuery(matchText *types.MatchTextExpr, vectorSimilarityWeigh "query_string": map[string]interface{}{ "fields": fields, "type": "best_fields", - "query": matchText.MatchingText, + "query": strings.ToLower(matchText.MatchingText), "minimum_should_match": minimumShouldMatch, "boost": boost, }, diff --git a/internal/engine/elasticsearch/chunk_helpers_test.go b/internal/engine/elasticsearch/chunk_helpers_test.go index 0987e351bc..a69a475860 100644 --- a/internal/engine/elasticsearch/chunk_helpers_test.go +++ b/internal/engine/elasticsearch/chunk_helpers_test.go @@ -18,7 +18,7 @@ func TestBuildQueryStringQueryMapsSkillFieldsToTokenFields(t *testing.T) { query := buildQueryStringQuery(&types.MatchTextExpr{ MatchingText: "test", Fields: []string{"name^10", "tags^5", "description^3", "content^1"}, - }, 0, true, false) + }, true, false) queryString, ok := query["query_string"].(map[string]interface{}) if !ok { @@ -32,7 +32,7 @@ func TestBuildQueryStringQueryKeepsDocumentFieldsUnchanged(t *testing.T) { query := buildQueryStringQuery(&types.MatchTextExpr{ MatchingText: "test", Fields: []string{"name^10"}, - }, 0, false, false) + }, false, false) queryString, ok := query["query_string"].(map[string]interface{}) if !ok { @@ -41,6 +41,21 @@ func TestBuildQueryStringQueryKeepsDocumentFieldsUnchanged(t *testing.T) { assertEqual(t, queryString["fields"], []string{"name^10"}) } +func TestBuildQueryStringQueryLowercasesMatchingText(t *testing.T) { + // The *_tks/*_ltks fields are whitespace-analyzed and store lowercase + // tokens; a capitalized query term must not silently match nothing. + query := buildQueryStringQuery(&types.MatchTextExpr{ + MatchingText: "Isabel Wood co-lead Ross Feldner Bob Musil Bird Watch Wonder Program", + }, false, false) + + queryString, ok := query["query_string"].(map[string]interface{}) + if !ok { + t.Fatalf("query_string missing from %#v", query) + } + assertEqual(t, queryString["query"], + "isabel wood co-lead ross feldner bob musil bird watch wonder program") +} + func TestSearchUsesConfiguredKNNNumCandidates(t *testing.T) { if err := common.InitLogger("info", common.FileOutput{}, "elasticsearch_test"); err != nil { t.Fatalf("init logger: %v", err) diff --git a/internal/engine/elasticsearch/chunk_readback_integration_test.go b/internal/engine/elasticsearch/chunk_readback_integration_test.go index a24be66c46..dcef822925 100644 --- a/internal/engine/elasticsearch/chunk_readback_integration_test.go +++ b/internal/engine/elasticsearch/chunk_readback_integration_test.go @@ -37,12 +37,12 @@ func TestInsertChunks_ReadBackSuffixedFields(t *testing.T) { t.Skip("Skipping ES integration test; set ES_TEST=1 to run") } - engine, err := NewEngine(t.Context(), getESTestConfig()) + ctx := t.Context() + engine, err := NewEngine(ctx, getESTestConfig()) if err != nil { t.Fatalf("NewEngine: %v", err) } - ctx := t.Context() baseName := "ragflow_chunk_readback_test" datasetID := "kb-1" chunkID := "readback-chunk-1" diff --git a/internal/engine/elasticsearch/chunk_refresh_test.go b/internal/engine/elasticsearch/chunk_refresh_test.go new file mode 100644 index 0000000000..523d3ff34b --- /dev/null +++ b/internal/engine/elasticsearch/chunk_refresh_test.go @@ -0,0 +1,83 @@ +// +// 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 elasticsearch + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/elastic/go-elasticsearch/v8" +) + +const bulkOKBody = `{"took":1,"errors":false,"items":[{"index":{"_index":"ragflow_tenant","_id":"c1","status":201}}]}` + +// newBulkCaptureEngine returns an engine whose bulk requests are answered +// locally, recording the refresh parameter of every request it sees. +func newBulkCaptureEngine(t *testing.T, seen *[]string) *Engine { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + *seen = append(*seen, r.URL.Query().Get("refresh")) + w.Header().Set("X-Elastic-Product", "Elasticsearch") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(bulkOKBody)) + })) + t.Cleanup(server.Close) + + client, err := elasticsearch.NewClient(elasticsearch.Config{Addresses: []string{server.URL}}) + if err != nil { + t.Fatalf("new elasticsearch client: %v", err) + } + return &Engine{client: client} +} + +func oneChunk() []map[string]interface{} { + return []map[string]interface{}{{"id": "c1", "doc_id": "d1", "content_with_weight": "hello"}} +} + +// The chunk APIs read chunks straight back after writing them, so the plain +// inserter keeps waiting for a refresh. +func TestInsertChunksWaitsForRefresh(t *testing.T) { + var refresh []string + engine := newBulkCaptureEngine(t, &refresh) + + if _, err := engine.InsertChunks(t.Context(), oneChunk(), "ragflow_tenant", "kb-1"); err != nil { + t.Fatalf("InsertChunks: %v", err) + } + if len(refresh) != 1 || refresh[0] != "wait_for" { + t.Fatalf("refresh params = %v, want exactly one wait_for", refresh) + } +} + +// Ingestion must not pay the refresh wait: omitting the parameter leaves the +// decision to the index (refresh_interval=1000ms) and matches Python's +// ingestion, which inserts with refresh=False. Waiting here costs ~0.5s per +// write and buys nothing - the chunks are published on the index's own cycle. +func TestInsertChunksNoRefreshOmitsRefreshParam(t *testing.T) { + var refresh []string + engine := newBulkCaptureEngine(t, &refresh) + + if _, err := engine.InsertChunksNoRefresh(t.Context(), oneChunk(), "ragflow_tenant", "kb-1"); err != nil { + t.Fatalf("InsertChunksNoRefresh: %v", err) + } + if len(refresh) != 1 { + t.Fatalf("bulk requests = %d, want 1", len(refresh)) + } + if refresh[0] != "" { + t.Fatalf("refresh param = %q, want it omitted so the index decides", refresh[0]) + } +} diff --git a/internal/engine/elasticsearch/chunk_test.go b/internal/engine/elasticsearch/chunk_test.go index 61952e3fbc..8a3f233a64 100644 --- a/internal/engine/elasticsearch/chunk_test.go +++ b/internal/engine/elasticsearch/chunk_test.go @@ -435,6 +435,15 @@ func TestBuildBoolQueryFromConditionIDFilter(t *testing.T) { check("int_value", map[string]interface{}{ "id": 42, }, []string{"id", "_id"}) + + // A typed []string must be handled too: callers built from typed helpers + // (e.g. list_chunks' ChunkScope) pass []string, and the generic loop below + // skips the "id" key — without this branch the query carries NO id filter + // and a scoped read silently fetches the whole document (an 11-chunk + // window returned 3.8MB instead of ~17KB). + check("string_slice_value", map[string]interface{}{ + "id": []string{"a", "b", "c"}, + }, []string{"id", "_id"}) } // paginationGRID mirrors the (page_size, top) grid from @@ -518,7 +527,7 @@ func TestBuildQueryStringQueryMinimumShouldMatchHalfUp(t *testing.T) { query := buildQueryStringQuery(&types.MatchTextExpr{ MatchingText: "hello", ExtraOptions: map[string]interface{}{"minimum_should_match": tc.fraction}, - }, 0.5, false, false) + }, false, false) got := query["query_string"].(map[string]interface{})["minimum_should_match"].(string) if got != tc.want { t.Errorf("buildQueryStringQuery minimum_should_match for %g = %q, want %q", tc.fraction, got, tc.want) diff --git a/internal/engine/elasticsearch/kg_test.go b/internal/engine/elasticsearch/kg_test.go index 6792eae27b..b7c686a890 100644 --- a/internal/engine/elasticsearch/kg_test.go +++ b/internal/engine/elasticsearch/kg_test.go @@ -28,8 +28,8 @@ import ( // TestKGSearchSelectFields verifies that SelectFields overrides default output // columns when searching for knowledge graph entities. -// Requires a running Elasticsearch instance and KG data indexed by Python task executor. -// Set ES_TEST=1 to run. +// Requires a running Elasticsearch instance and KG data indexed by Python task +// executor. Gated by the integration build tag. func TestKGSearchSelectFields(t *testing.T) { if common.GetEnv(common.EnvESTest) != "1" { t.Skip("Skipping ES integration test; set ES_TEST=1 to run") diff --git a/internal/engine/nats/consumer_selfheal_test.go b/internal/engine/nats/consumer_selfheal_test.go new file mode 100644 index 0000000000..46d517497b --- /dev/null +++ b/internal/engine/nats/consumer_selfheal_test.go @@ -0,0 +1,143 @@ +// +// 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 nats + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +func TestIsConsumerGoneErr(t *testing.T) { + gone := []error{ + jetstream.ErrConsumerNotFound, + jetstream.ErrConsumerDeleted, + nats.ErrNoResponders, + fmt.Errorf("failed to fetch messages: %w", jetstream.ErrConsumerNotFound), + fmt.Errorf("failed to fetch messages: %w", nats.ErrNoResponders), + errors.New("failed to fetch messages: nats: no responders available for request"), + errors.New("failed to fetch messages: consumer not found"), + errors.New("failed to fetch messages: consumer deleted"), + } + for _, err := range gone { + if !isConsumerGoneErr(err) { + t.Errorf("isConsumerGoneErr(%v) = false, want true", err) + } + } + + notGone := []error{ + nil, + context.DeadlineExceeded, + context.Canceled, + errors.New("failed to fetch messages: nats: timeout"), + errors.New("NATS consumer is nil, engine not properly initialized"), + } + for _, err := range notGone { + if isConsumerGoneErr(err) { + t.Errorf("isConsumerGoneErr(%v) = true, want false", err) + } + } +} + +// A pull must repair itself when the consumer behind the stored handle is gone. +// Before this, the fetch loop answered "no responders available for request" +// forever: the workers kept logging errors, the queue stopped draining, and only +// restarting the process cleared it. +// +// Note on what recovery does and does not cover: the task stream is a +// WorkQueuePolicy stream, so messages that were pending for a deleted consumer +// are gone with it - recreating the consumer restores the pull plumbing, it does +// not resurrect those messages. Asserting that is the point of this test: after +// the repair, a pull must stop erroring and newly published tasks must flow +// again. +func TestPullMessages_RecreatesConsumerWhenItDisappears(t *testing.T) { + queue := newTaskQueue(t) + + publishTask(t, queue, "before-delete") + ctx, cancel := taskPullContext(t) + handle, err := queue.PullMessage(ctx) + cancel() + if err != nil { + t.Fatalf("PullMessage: %v", err) + } + if handle == nil || handle.GetMessage().TaskID != "before-delete" { + t.Fatalf("handle = %v, want before-delete", handle) + } + if ackErr := handle.Ack(); ackErr != nil { + t.Fatalf("ack: %v", ackErr) + } + + // Delete the consumer behind the engine's back; the stored handle now + // addresses something that no longer exists. + if err := queue.stream.DeleteConsumer(t.Context(), "RAGFLOW_CONSUMER"); err != nil { + t.Fatalf("delete consumer: %v", err) + } + + // The first pull after the deletion is the one that must self-heal: it may + // come back empty, but it must not come back with the no-responders error. + pullCtx, pullCancel := context.WithTimeout(t.Context(), 5*time.Second) + defer pullCancel() + if _, err := queue.PullMessage(pullCtx); err != nil { + t.Fatalf("PullMessage after the consumer was deleted: %v", err) + } + + // The consumer must exist again under the same durable name, so the next + // restart reattaches instead of orphaning a replacement. + if _, err := queue.stream.Consumer(t.Context(), "RAGFLOW_CONSUMER"); err != nil { + t.Fatalf("consumer was not recreated: %v", err) + } + + // And the recreated consumer must actually deliver: publish after the repair + // and pull it back. + publishTask(t, queue, "after-repair") + recovered, err := queue.PullMessage(pullCtx) + if err != nil { + t.Fatalf("PullMessage after the repair: %v", err) + } + if recovered == nil || recovered.GetMessage().TaskID != "after-repair" { + t.Fatalf("recovered handle = %v, want after-repair", recovered) + } +} + +// A healthy pull leaves the handle alone: the repair path must not run on the +// ordinary "nothing to deliver" outcome. +func TestPullMessages_KeepsHandleWhenNothingToDeliver(t *testing.T) { + queue := newTaskQueue(t) + before, err := queue.consumerHandle() + if err != nil { + t.Fatalf("consumerHandle: %v", err) + } + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + if _, err := queue.PullMessage(ctx); err != nil { + t.Fatalf("PullMessage on an empty queue: %v", err) + } + + after, err := queue.consumerHandle() + if err != nil { + t.Fatalf("consumerHandle: %v", err) + } + if before != after { + t.Fatal("the consumer handle was replaced although nothing went wrong") + } +} diff --git a/internal/engine/nats/knowledgecompile.go b/internal/engine/nats/knowledgecompile.go index 330a199f94..49b78f460f 100644 --- a/internal/engine/nats/knowledgecompile.go +++ b/internal/engine/nats/knowledgecompile.go @@ -60,7 +60,7 @@ func (n *NatsEngine) InitKnowledgeCompileStream() error { Storage: jetstream.FileStorage, Discard: jetstream.DiscardNew, MaxMsgs: 1024 * 1024, - MaxBytes: 1024 * 1024 * 64, + MaxBytes: 1024 * 1024 * 1024, }) if err != nil { return fmt.Errorf("knowledgecompile: create stream: %w", err) diff --git a/internal/engine/nats/nats.go b/internal/engine/nats/nats.go index 2657e4cc48..4560559009 100644 --- a/internal/engine/nats/nats.go +++ b/internal/engine/nats/nats.go @@ -31,13 +31,31 @@ import ( "github.com/nats-io/nats.go/jetstream" ) +// isStreamExistsErr reports whether err means the stream already exists, using +// the typed JetStream APIError (err_code 10058) rather than error-string +// matching. The APIError code is the canonical, stable signal from the server. +func isStreamExistsErr(err error) bool { + if errors.Is(err, nats.ErrStreamNameAlreadyInUse) { + return true + } + var apiErr *jetstream.APIError + if errors.As(err, &apiErr) { + return apiErr.ErrorCode == jetstream.JSErrCodeStreamNameInUse + } + return false +} + type NatsEngine struct { host string port int nc *nats.Conn jetStream jetstream.JetStream stream jetstream.Stream - consumer jetstream.Consumer + // consumer is the task consumer handle. Guarded by consumerMu: the pull loop + // and the admin inspection endpoint share one engine, and the handle is + // replaced when it is found to have outlived its consumer (see PullMessages). + consumerMu sync.Mutex + consumer jetstream.Consumer // dataset-level compile consumer (§11) state. knowledgeCompileStream jetstream.Stream @@ -80,19 +98,35 @@ func (n *NatsEngine) Init() error { Retention: jetstream.WorkQueuePolicy, Storage: jetstream.FileStorage, Discard: jetstream.DiscardNew, - MaxMsgs: 1024 * 128, - MaxBytes: 1024 * 1024 * 64, - // Server-side dedup window. Inert for task publishes: PublishTask - // intentionally sends no MsgID (see below — dedup would swallow - // retry republishes of a reused task_id). It only takes effect for - // a publisher that opts into MsgIDs. - Duplicates: 10 * time.Minute, + MaxMsgs: 1024 * 1024, + MaxBytes: 1024 * 1024 * 1024, } n.stream, err = ensureStreamConfig(ctx, n.jetStream, streamCfg) if err != nil { - n.nc.Close() - return fmt.Errorf("fail to create stream at %s: %w", natsURL, err) + if !isStreamExistsErr(err) { + n.nc.Close() + return fmt.Errorf("fail to create stream at %s: %w", natsURL, err) + } + + common.Info("NATS stream already exists, use existing stream") + n.stream, err = n.jetStream.Stream(ctx, "RAGFLOW_TASKS") + if err != nil { + n.nc.Close() + return fmt.Errorf("fail to get existing stream at %s: %w", natsURL, err) + } + // Reconcile the running stream with the configured limits so an existing + // stream created under older settings (e.g. DiscardOld + small MaxBytes) + // picks up DiscardNew and the larger MaxMsgs/MaxBytes. CreateStream never + // touches an already-existing stream, so UpdateStream is required here. + n.stream, err = n.jetStream.UpdateStream(ctx, streamCfg) + if err != nil { + n.nc.Close() + return fmt.Errorf("fail to update existing stream at %s: %w", natsURL, err) + } + common.Info("NATS stream updated with current config (discard=new, larger limits)") + } else { + common.Info(fmt.Sprintf("NATS stream create successfully at %s", natsURL)) } common.Info(fmt.Sprintf("NATS stream RAGFLOW_TASKS ready at %s", natsURL)) @@ -212,13 +246,21 @@ func (n *NatsEngine) ListMessages(messageType string, pending bool) ([]map[strin } func (n *NatsEngine) InitConsumer(subject string) error { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + return n.ensureConsumer(ctx) +} + +// ensureConsumer creates or updates RAGFLOW_CONSUMER and stores the handle. +// It is idempotent, so it doubles as the repair path when a stored handle is +// found to have outlived its consumer. The consumer name is a fixed, durable +// name on purpose: a restart must reattach to the same consumer (and its +// pending messages) instead of leaving an orphan behind. +func (n *NatsEngine) ensureConsumer(ctx context.Context) error { if n.stream == nil { return fmt.Errorf("NATS stream is nil, engine not properly initialized") } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - var err error // Explicit redelivery schedule: BackOff paces successive redeliveries // (5s/15s/30s, then 60s repeated) so an unsettled message (crash, slow // DB) is retried with breathing room instead of the broker default. The @@ -233,7 +275,7 @@ func (n *NatsEngine) InitConsumer(subject string) error { // applied. When MaxWaiting matches (the default, since this config omits // it), the other three update in place. The fallback below handles the // MaxWaiting mismatch by keeping the existing consumer. - n.consumer, err = n.stream.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + consumer, err := n.stream.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ Name: "RAGFLOW_CONSUMER", AckPolicy: jetstream.AckExplicitPolicy, MaxDeliver: 16, @@ -246,7 +288,7 @@ func (n *NatsEngine) InitConsumer(subject string) error { // MaxWaiting is immutable after consumer creation (AckWait/BackOff/MaxAckPending remain mutable when MaxWaiting matches; the update is atomic). // If the consumer already exists, fall back to fetching it. if strings.Contains(err.Error(), "max waiting can not be updated") { - n.consumer, err = n.stream.Consumer(ctx, "RAGFLOW_CONSUMER") + consumer, err = n.stream.Consumer(ctx, "RAGFLOW_CONSUMER") if err != nil { return fmt.Errorf("failed to get existing consumer: %w", err) } @@ -254,26 +296,128 @@ func (n *NatsEngine) InitConsumer(subject string) error { return fmt.Errorf("failed to create Consumer: %w", err) } } + n.consumerMu.Lock() + n.consumer = consumer + n.consumerMu.Unlock() return nil } +// consumerHandle returns the task consumer handle, or an error when the engine +// has not created one yet. +func (n *NatsEngine) consumerHandle() (jetstream.Consumer, error) { + n.consumerMu.Lock() + defer n.consumerMu.Unlock() + if n.consumer == nil { + return nil, errors.New("NATS consumer is nil, engine not properly initialized") + } + return n.consumer, nil +} + +// recreateConsumer rebuilds RAGFLOW_CONSUMER and returns the fresh handle. +// Consumer creation is a JetStream API round trip, so it gets its own timeout +// rather than borrowing the (deliberately short) pull deadline. +func (n *NatsEngine) recreateConsumer(ctx context.Context) (jetstream.Consumer, error) { + createCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + if err := n.ensureConsumer(createCtx); err != nil { + return nil, err + } + return n.consumerHandle() +} + +// consumerRepairPullAttempts / consumerRepairPullDelay bound the retry that +// follows a consumer rebuild: the pull subject of a freshly created consumer can +// take a moment to become addressable, and a pull in that window answers the +// very same "no responders" error the repair exists to clear. +const ( + consumerRepairPullAttempts = 5 + consumerRepairPullDelay = 200 * time.Millisecond +) + +// isConsumerGoneErr reports whether err means the stored consumer handle no +// longer addresses a live consumer: it was deleted, never existed, or the +// JetStream context went stale. Typed JetStream/NATS errors are checked first - +// they are the stable signal - with a string fallback for wrapped variants. +func isConsumerGoneErr(err error) bool { + if err == nil { + return false + } + if errors.Is(err, jetstream.ErrConsumerNotFound) || + errors.Is(err, jetstream.ErrConsumerDeleted) || + errors.Is(err, nats.ErrNoResponders) { + return true + } + msg := strings.ToLower(err.Error()) + return strings.Contains(msg, "no responders") || + strings.Contains(msg, "consumer not found") || + strings.Contains(msg, "consumer deleted") +} + // PullMessages fetches up to messageCount messages before ctx expires. func (n *NatsEngine) PullMessages(ctx context.Context, messageCount int) ([]common.TaskHandle, error) { if messageCount < 1 || messageCount > common.MaxManualPullMessages { return nil, fmt.Errorf("message count must be between 1 and %d", common.MaxManualPullMessages) } - if n.consumer == nil { - return nil, errors.New("NATS consumer is nil, engine not properly initialized") + consumer, err := n.consumerHandle() + if err != nil { + return nil, err } if _, ok := ctx.Deadline(); !ok { return nil, errors.New("pull messages context must have a deadline") } - resultMessages := make([]common.TaskHandle, 0, messageCount) - messages, err := n.consumer.Fetch(messageCount, jetstream.FetchContext(ctx)) - if err != nil { - return nil, fmt.Errorf("failed to fetch messages: %w", err) + resultMessages, pullErr := n.fetchOnce(ctx, consumer, messageCount) + if !isConsumerGoneErr(pullErr) { + if pullErr != nil { + return nil, fmt.Errorf("failed to fetch messages: %w", pullErr) + } + return resultMessages, nil } + + // The handle outlived its consumer (deleted behind our back, or the + // JetStream context went stale across a reconnect). Retrying against the + // dead handle answers "no responders available for request" forever: the + // worker logs errors, the queue stops draining, and only a process restart + // clears it. Rebuild the consumer and retry, bounded, so a persistent failure + // still surfaces instead of spinning. + common.Warn(fmt.Sprintf("NATS consumer went missing (%v); recreating RAGFLOW_CONSUMER", pullErr)) + consumer, recreateErr := n.recreateConsumer(ctx) + if recreateErr != nil { + return nil, fmt.Errorf("failed to fetch messages: %w (recreate consumer: %v)", pullErr, recreateErr) + } + // A pull issued in the instant right after the consumer is created can still + // land before the server registers the new consumer's pull subject - the same + // "no responders" answer, briefly - so retry a few times with a short pause + // instead of declaring the repair failed. + for attempt := 1; attempt <= consumerRepairPullAttempts; attempt++ { + resultMessages, pullErr = n.fetchOnce(ctx, consumer, messageCount) + if !isConsumerGoneErr(pullErr) || attempt == consumerRepairPullAttempts { + break + } + select { + case <-ctx.Done(): + return nil, fmt.Errorf("failed to fetch messages: %w", pullErr) + case <-time.After(consumerRepairPullDelay): + } + } + if pullErr != nil { + return nil, fmt.Errorf("failed to fetch messages: %w", pullErr) + } + return resultMessages, nil +} + +// fetchOnce performs a single pull and folds the batch error into the returned +// error. A pull deadline means "nothing was delivered" and is not an error - the +// caller's loop relies on that - while any other batch failure nacks whatever did +// arrive so the messages stay available. Errors can surface either from Fetch +// itself or from the returned batch, which is why the no-responders repair above +// has to look at both. +func (n *NatsEngine) fetchOnce(ctx context.Context, consumer jetstream.Consumer, messageCount int) ([]common.TaskHandle, error) { + messages, err := consumer.Fetch(messageCount, jetstream.FetchContext(ctx)) + if err != nil { + return nil, err + } + resultMessages := make([]common.TaskHandle, 0, messageCount) for message := range messages.Messages() { resultMessages = append(resultMessages, NewNatsMessageHandle(message)) } @@ -286,7 +430,7 @@ func (n *NatsEngine) PullMessages(ctx context.Context, messageCount int) ([]comm common.Error("nack message after failed pull", nackErr) } } - return nil, fmt.Errorf("failed to fetch messages: %w", batchErr) + return nil, batchErr } return resultMessages, nil } diff --git a/internal/engine/nats/stream_config_test.go b/internal/engine/nats/stream_config_test.go index d6baf63875..be3cb8a58c 100644 --- a/internal/engine/nats/stream_config_test.go +++ b/internal/engine/nats/stream_config_test.go @@ -104,10 +104,10 @@ func TestPublishTaskDeliversRepeatedTaskIDs(t *testing.T) { } // TestInitMigratesLegacyStreamConfig: a stream created by an older deployment -// (no Duplicates, 1MB MaxBytes) must be migrated in place by Init instead of -// being left stale behind an "already exists" error. The server-side config is -// the merge base: fields the helper does not own (Subjects, Retention) must -// survive the update. +// (1MB MaxBytes, server-default Duplicates) must be migrated in place by Init +// instead of being left stale behind an "already exists" error. The server-side +// config is the merge base: fields the wanted config does not own (Subjects, +// Retention, Duplicates) must survive the update. func TestInitMigratesLegacyStreamConfig(t *testing.T) { host, port := newEmbeddedNatsServer(t) @@ -144,9 +144,6 @@ func TestInitMigratesLegacyStreamConfig(t *testing.T) { if legacyInfo.Config.MaxBytes != int64(1024*1024) { t.Fatalf("precondition: legacy MaxBytes = %d, want 1MB", legacyInfo.Config.MaxBytes) } - if legacyInfo.Config.Duplicates == 10*time.Minute { - t.Fatalf("precondition: legacy Duplicates already migrated (%v)", legacyInfo.Config.Duplicates) - } if legacyInfo.Config.Discard != jetstream.DiscardOld { t.Fatalf("precondition: legacy Discard = %v, want DiscardOld", legacyInfo.Config.Discard) } @@ -160,11 +157,15 @@ func TestInitMigratesLegacyStreamConfig(t *testing.T) { if err != nil { t.Fatalf("stream info after migration: %v", err) } - if got := info.Config.MaxBytes; got != int64(1024*1024*64) { - t.Fatalf("MaxBytes after migration = %d, want %d", got, int64(1024*1024*64)) + if got := info.Config.MaxBytes; got != int64(1024*1024*1024) { + t.Fatalf("MaxBytes after migration = %d, want %d", got, int64(1024*1024*1024)) } - if got := info.Config.Duplicates; got != 10*time.Minute { - t.Fatalf("Duplicates after migration = %v, want 10m", got) + // Duplicates is deliberately NOT part of the wanted config (RAGFLOW_TASKS + // publishes without a MsgID, so its dedup window never applies), which makes it + // a non-owned field like Retention/Subjects below: migration must leave the + // legacy value exactly as it was. + if got := info.Config.Duplicates; got != legacyInfo.Config.Duplicates { + t.Fatalf("Duplicates after migration = %v, want the legacy %v (must not be reset)", got, legacyInfo.Config.Duplicates) } if got := info.Config.Discard; got != jetstream.DiscardNew { t.Fatalf("Discard after migration = %v, want DiscardNew", got) diff --git a/internal/engine/nats/syncer.go b/internal/engine/nats/syncer.go index e2253b900c..fd00447119 100644 --- a/internal/engine/nats/syncer.go +++ b/internal/engine/nats/syncer.go @@ -68,12 +68,34 @@ func (n *NatsEngine) initSyncerStreamLocked() error { Retention: jetstream.WorkQueuePolicy, Storage: jetstream.FileStorage, Discard: jetstream.DiscardNew, - MaxMsgs: 1024 * 128, - MaxBytes: 1024 * 1024 * 64, + MaxMsgs: 1024 * 1024, + MaxBytes: 1024 * 1024 * 1024, Duplicates: 10 * time.Minute, }) if err != nil { - return fmt.Errorf("syncer: create stream: %w", err) + if !isStreamExistsErr(err) { + return fmt.Errorf("syncer: create stream: %w", err) + } + stream, err = n.jetStream.Stream(ctx, syncerStreamName) + if err != nil { + return fmt.Errorf("syncer: get existing stream: %w", err) + } + // Reconcile an existing stream created under older settings with the + // current config (DiscardNew + larger MaxMsgs/MaxBytes). CreateStream + // never touches an already-existing stream, so UpdateStream is required. + stream, err = n.jetStream.UpdateStream(ctx, jetstream.StreamConfig{ + Name: syncerStreamName, + Subjects: []string{syncerSubjectPattern}, + Retention: jetstream.WorkQueuePolicy, + Storage: jetstream.FileStorage, + Discard: jetstream.DiscardNew, + MaxMsgs: 1024 * 1024, + MaxBytes: 1024 * 1024 * 1024, + Duplicates: 10 * time.Minute, + }) + if err != nil { + return fmt.Errorf("syncer: update existing stream: %w", err) + } } n.syncerStream = stream return nil diff --git a/internal/entity/models/302ai.go b/internal/entity/models/302ai.go index f63aabad1a..b0d999a3be 100644 --- a/internal/entity/models/302ai.go +++ b/internal/entity/models/302ai.go @@ -117,7 +117,7 @@ func (a *AI302Model) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (a *AI302Model) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/aliyun.go b/internal/entity/models/aliyun.go index 388dc3170d..3abe2344e3 100644 --- a/internal/entity/models/aliyun.go +++ b/internal/entity/models/aliyun.go @@ -98,7 +98,7 @@ func (a *AliyunModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/astraflow.go b/internal/entity/models/astraflow.go index 25e49f212b..b913450d06 100644 --- a/internal/entity/models/astraflow.go +++ b/internal/entity/models/astraflow.go @@ -93,7 +93,7 @@ func (a *AstraflowModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender opens the SSE chat-completions diff --git a/internal/entity/models/avian.go b/internal/entity/models/avian.go index 457c1bb050..bae804df52 100644 --- a/internal/entity/models/avian.go +++ b/internal/entity/models/avian.go @@ -83,7 +83,7 @@ func (a *AvianModel) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (a *AvianModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/azure_openai.go b/internal/entity/models/azure_openai.go index 104348c88d..baa4fed4ac 100644 --- a/internal/entity/models/azure_openai.go +++ b/internal/entity/models/azure_openai.go @@ -110,7 +110,7 @@ func (a *AzureOpenAIModel) ChatWithMessages(ctx context.Context, modelName strin return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response via the diff --git a/internal/entity/models/baichuan.go b/internal/entity/models/baichuan.go index 737bd5da0b..a7ebe6785f 100644 --- a/internal/entity/models/baichuan.go +++ b/internal/entity/models/baichuan.go @@ -69,7 +69,7 @@ func (b *BaichuanModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (b *BaichuanModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/baidu.go b/internal/entity/models/baidu.go index 6f66a041b2..3187125b13 100644 --- a/internal/entity/models/baidu.go +++ b/internal/entity/models/baidu.go @@ -116,7 +116,7 @@ func (b *BaiduModel) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (b *BaiduModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/base_model.go b/internal/entity/models/base_model.go index 1768a5ab02..20d000363c 100644 --- a/internal/entity/models/base_model.go +++ b/internal/entity/models/base_model.go @@ -39,6 +39,19 @@ const ( maxLoggedVectorFloats = 3 ) +// APIStatusError is a provider HTTP failure with its status code preserved, so +// callers can act on the status (failover cooldown, retry) instead of matching +// the error text. The shared request helpers return it; the message is byte-for- +// byte what plain fmt.Errorf produced before, so existing assertions still hold. +type APIStatusError struct { + Status int + Body string +} + +func (e *APIStatusError) Error() string { + return fmt.Sprintf("API request failed with status %d: %s", e.Status, e.Body) +} + type BaseModel struct { BaseURL map[string]string URLSuffix URLSuffix @@ -225,7 +238,7 @@ func (b *BaseModel) doRequest(ctx context.Context, url string, apiConfig *APICon if err != nil { return nil, fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err) } - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + return nil, &APIStatusError{Status: resp.StatusCode, Body: string(body)} } body, err := readModelResponseBody(resp.Body) @@ -260,7 +273,7 @@ func (b *BaseModel) doGetRequest(ctx context.Context, url string, apiConfig *API if err != nil { return nil, fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err) } - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + return nil, &APIStatusError{Status: resp.StatusCode, Body: string(body)} } body, err := readModelResponseBody(resp.Body) @@ -293,7 +306,7 @@ func (b *BaseModel) doStreamRequest(ctx context.Context, url string, apiConfig * if err != nil { return fmt.Errorf("API request failed with status %d; failed to read error response: %w", resp.StatusCode, err) } - return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) + return &APIStatusError{Status: resp.StatusCode, Body: string(body)} } return handler(resp.Body) diff --git a/internal/entity/models/chat_tools.go b/internal/entity/models/chat_tools.go index f84eccdf24..481ed0e0ca 100644 --- a/internal/entity/models/chat_tools.go +++ b/internal/entity/models/chat_tools.go @@ -26,18 +26,31 @@ import ( "ragflow/internal/tokenizer" ) -// recordUsageFromResponse records one chat call's token usage to both -// cm.LastUsage and the context-level run sink (if installed). Callers -// should invoke this after each ChatWithMessages / ChatStreamlyWithSender -// call so the canvas-layer aggregator and Langfuse both see the split. -func recordUsageFromResponse(ctx context.Context, cm *ChatModel) { +// recordCallUsage feeds ONE chat call's token split to the context-level run sink (if +// installed). The caller passes the numbers it already holds: the ChatModel is shared +// by every call on it, so a per-call field would let a concurrent call's tokens be +// attributed to this one. The canvas component gets its own copy from the message +// metadata (ResponseMeta.Usage) for the same reason. +func recordCallUsage(ctx context.Context, cm *ChatModel, prompt, completion, total int) { if cm == nil { return } - if cm.LastUsage == nil { + model := "" + if cm.ModelName != nil { + model = *cm.ModelName + } + recordUsage(ctx, model, &TokenUsage{PromptTokens: prompt, CompletionTokens: completion, TotalTokens: total}) +} + +// recordUsage records one chat call's token usage to the context-level run sink (if +// installed). The caller passes the usage its own call returned: nothing about one +// call's tokens is parked on the shared *ChatModel, where a concurrent call could +// replace them between the write and the read and attribute them to the wrong request. +func recordUsage(ctx context.Context, model string, usage *TokenUsage) { + if usage == nil { return } - tokenizer.RecordRunTokenUsage(ctx, cm.LastUsage.PromptTokens, cm.LastUsage.CompletionTokens, cm.LastUsage.TotalTokens) + tokenizer.RecordRunTokenUsageFor(ctx, model, usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens) } const ( @@ -96,9 +109,6 @@ func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsLis // Mirrors Python PR #16420 fix: previously the total was overwritten // each round; now we accumulate so multi-round tool conversations // report the correct grand total. - // Reset stale per-call usage from a previous call so a response - // without usage doesn't leak the prior call's data. - cm.LastUsage = nil var totalTokens int aggUsage := &TokenUsage{} @@ -113,10 +123,7 @@ func runToolLoop(ctx context.Context, cm *ChatModel, history []Message, toolsLis totalTokens = aggUsage.TotalTokens // Store per-round delta (not cumulative) so RecordRunTokenUsage // records each round's contribution exactly once. - cm.LastUsage = &TokenUsage{ - PromptTokens: u.PromptTokens, CompletionTokens: u.CompletionTokens, TotalTokens: u.TotalTokens, - } - recordUsageFromResponse(ctx, cm) + recordCallUsage(ctx, cm, u.PromptTokens, u.CompletionTokens, u.TotalTokens) } for round := 0; round <= maxRounds; round++ { @@ -253,8 +260,6 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to // Aggregate token counts across every tool-calling round (each round is a // separate provider request). Committing per round avoids the previous // bug where a later round's total overwrote earlier rounds. - // Reset stale per-call usage from a previous call. - cm.LastUsage = nil var totalTokens int aggUsage := &TokenUsage{} @@ -276,10 +281,7 @@ func runStreamToolLoop(ctx context.Context, cm *ChatModel, history []Message, to aggUsage.TotalTokens += deltaTotal } totalTokens = aggUsage.TotalTokens - cm.LastUsage = &TokenUsage{ - PromptTokens: deltaPrompt, CompletionTokens: deltaCompletion, TotalTokens: deltaTotal, - } - recordUsageFromResponse(ctx, cm) + recordCallUsage(ctx, cm, deltaPrompt, deltaCompletion, deltaTotal) } for round := 0; round <= maxRounds; round++ { diff --git a/internal/entity/models/cometapi.go b/internal/entity/models/cometapi.go index 4760447dd5..f8842fae4d 100644 --- a/internal/entity/models/cometapi.go +++ b/internal/entity/models/cometapi.go @@ -127,7 +127,7 @@ func (c *CometAPIModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response diff --git a/internal/entity/models/deepinfra.go b/internal/entity/models/deepinfra.go index fa1c568912..8177dca5c6 100644 --- a/internal/entity/models/deepinfra.go +++ b/internal/entity/models/deepinfra.go @@ -93,7 +93,7 @@ func (d *DeepInfraModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (d *DeepInfraModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index d90c888123..5bf8808975 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -138,7 +138,7 @@ func (d *DeepSeekModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/futurmix.go b/internal/entity/models/futurmix.go index 9a9dfde737..a6d62732e5 100644 --- a/internal/entity/models/futurmix.go +++ b/internal/entity/models/futurmix.go @@ -91,7 +91,7 @@ func (f *FuturMixModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends a streaming chat completion diff --git a/internal/entity/models/gitee.go b/internal/entity/models/gitee.go index c1a9f14faf..09ff023765 100644 --- a/internal/entity/models/gitee.go +++ b/internal/entity/models/gitee.go @@ -89,7 +89,7 @@ func (g *GiteeModel) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/gpustack.go b/internal/entity/models/gpustack.go index 5043f7bce6..610670dd2b 100644 --- a/internal/entity/models/gpustack.go +++ b/internal/entity/models/gpustack.go @@ -76,7 +76,7 @@ func (g *GPUStackModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender streams the response via the sender. diff --git a/internal/entity/models/greenpt.go b/internal/entity/models/greenpt.go index 309ea41ba5..4f570b5dd1 100644 --- a/internal/entity/models/greenpt.go +++ b/internal/entity/models/greenpt.go @@ -103,7 +103,7 @@ func (m *GreenPTModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function diff --git a/internal/entity/models/groq.go b/internal/entity/models/groq.go index e8294b21b8..27e0d4701b 100644 --- a/internal/entity/models/groq.go +++ b/internal/entity/models/groq.go @@ -128,7 +128,7 @@ func (g *GroqModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (g *GroqModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/huaweicloud.go b/internal/entity/models/huaweicloud.go index 139200c679..8346e782e2 100644 --- a/internal/entity/models/huaweicloud.go +++ b/internal/entity/models/huaweicloud.go @@ -191,7 +191,7 @@ func (h *HuaweiCloudModel) ChatWithMessages(ctx context.Context, modelName strin return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (h *HuaweiCloudModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/huggingface.go b/internal/entity/models/huggingface.go index 62efdab69f..041e1f7515 100644 --- a/internal/entity/models/huggingface.go +++ b/internal/entity/models/huggingface.go @@ -84,7 +84,7 @@ func (h *HuggingFaceModel) ChatWithMessages(ctx context.Context, modelName strin return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (h *HuggingFaceModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/hunyuan.go b/internal/entity/models/hunyuan.go index 0ac018b997..951830d9a7 100644 --- a/internal/entity/models/hunyuan.go +++ b/internal/entity/models/hunyuan.go @@ -89,7 +89,7 @@ func (h *HunyuanModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender opens the SSE chat-completions diff --git a/internal/entity/models/jiekouai.go b/internal/entity/models/jiekouai.go index 436c365d78..5207da35b7 100644 --- a/internal/entity/models/jiekouai.go +++ b/internal/entity/models/jiekouai.go @@ -129,7 +129,7 @@ func (j *JieKouAIModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (j *JieKouAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/jina.go b/internal/entity/models/jina.go index f7a54a3dec..2181804f1f 100644 --- a/internal/entity/models/jina.go +++ b/internal/entity/models/jina.go @@ -127,7 +127,7 @@ func (j *JinaModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (j *JinaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/llm.go b/internal/entity/models/llm.go index 7ec073e960..72a551b8a4 100644 --- a/internal/entity/models/llm.go +++ b/internal/entity/models/llm.go @@ -33,23 +33,88 @@ package models import ( "context" "encoding/json" + "errors" "fmt" + "net" + "net/http" + "strings" "sync" + "time" "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" + "go.uber.org/zap" + + "ragflow/internal/common" ) -// EinoChatModel adapts a RAGFlow *ChatModel to eino's chat model interfaces. -// It is safe for concurrent use: all per-request state lives in the -// receiver's fields which are only mutated through WithTools (which returns -// a new instance, never mutating in place — see eino's +// EinoChatModel adapts one or more RAGFlow *ChatModels to eino's chat model +// interfaces. With more than one model it forms a FAILOVER chain: Generate / +// Stream start at the STICKY CURSOR and, when a call ends in ANY terminal +// error (provider quota walls, rate limits, outages, auth failures), move to +// the next entry — only a full sweep of the chain reports failure. The cursor +// stays on the model that last served a call, so a dead primary is not re-hit +// on every single Generate of a long conversation. +// +// It is safe for concurrent use: the cursor is mutex-guarded, and the +// per-request tool state is only mutated through WithTools (which returns a +// new instance, never mutating in place — see eino's // components/model/interface.go:84-99 for the rationale). type EinoChatModel struct { - inner *ChatModel + chain []*ChatModel + labels []string // human-readable tag per chain entry (model @ instance), logs only + // sweep caches the outcome of the last FULL-CHAIN failure. A provider + // quota wall does not recover within one question, and an agentic turn + // issues tens of Generate calls — without this, one dead roster costs + // (calls × models) doomed round-trips per question. Inside the cooldown a + // new Generate returns the cached error immediately instead of sweeping + // the chain again. + sweep struct { + failedAt time.Time + err error + } + chatCfg *ChatConfig tools []*schema.ToolInfo toolChoice *string + + mu sync.Mutex // guards cursor + sweep + cursor int // chain index to try first on the next call +} + +// failoverCooldown is how long a full-chain failure short-circuits later +// Generate calls. Long enough to absorb an agentic turn (which would otherwise +// re-sweep the roster on every ReAct step), short enough that a replenished +// plan is picked up without a restart. +const failoverCooldown = 30 * time.Second + +// cacheableSweepFailure reports whether a terminal chain failure may be cached +// to short-circuit later Generate/Stream calls for failoverCooldown. +// +// Only provider-WIDE failures qualify. A request-level failure (invalid +// messages, an unknown tool, a bad parameter) also fails on every chain entry, +// but it says nothing about the next request — caching it would answer every +// valid turn with that stale error for the whole cooldown. +func cacheableSweepFailure(err error) bool { + if err == nil { + return false + } + // A cancelled context is the caller's doing, not the provider's health. + if errors.Is(err, context.Canceled) { + return false + } + // Transport failures (timeout, DNS, refused connection) hit every request. + var netErr net.Error + if errors.As(err, &netErr) { + return true + } + // A provider status: rate limiting and 5xx mean the provider cannot serve + // anything right now; any other 4xx is about THIS request. + var statusErr *APIStatusError + if errors.As(err, &statusErr) { + return statusErr.Status == http.StatusTooManyRequests || statusErr.Status >= 500 + } + return false } // NewEinoChatModel wraps an existing RAGFlow *ChatModel so it can be passed @@ -59,19 +124,63 @@ type EinoChatModel struct { // Driver is taken from cm.ModelDriver, model name from cm.ModelName, and // API key / region from cm.APIConfig. These are fixed for the lifetime of // the wrapper; per-request variations belong in WithTools / a new instance. +// For a model chain with automatic failover, see NewFailoverEinoChatModel. func NewEinoChatModel(cm *ChatModel, chatConfig *ChatConfig) *EinoChatModel { return &EinoChatModel{ - inner: cm, + chain: []*ChatModel{cm}, chatCfg: chatConfig, } } -// name returns the underlying model name (best-effort; nil-safe). +// NewFailoverEinoChatModel wraps a chain of RAGFlow *ChatModels. models[0] is +// the primary; on a terminal Generate/Stream error the next entry is tried, +// and the chain is swept at most once per call before the last error is +// returned. Use NewEinoChatModel for a single model. +func NewFailoverEinoChatModel(models []*ChatModel, chatConfig *ChatConfig) (*EinoChatModel, error) { + return NewFailoverEinoChatModelWithLabels(models, nil, chatConfig) +} + +// NewFailoverEinoChatModelWithLabels is NewFailoverEinoChatModel plus a +// human-readable label per entry (e.g. "MiniMax-M3 @ zyf"). Chain entries are +// frequently the same model name on different provider instances, so the +// label is what makes failover logs attributable. A missing label falls back +// to the model name. +func NewFailoverEinoChatModelWithLabels(models []*ChatModel, labels []string, chatConfig *ChatConfig) (*EinoChatModel, error) { + chain := make([]*ChatModel, 0, len(models)) + tags := make([]string, 0, len(models)) + for i, cm := range models { + if cm == nil || cm.ModelDriver == nil { + continue + } + chain = append(chain, cm) + label := "" + if i < len(labels) { + label = labels[i] + } + if label == "" { + label = modelNameOf(cm) + } + tags = append(tags, label) + } + if len(chain) == 0 { + return nil, fmt.Errorf("models: NewFailoverEinoChatModel: no usable chat model in chain") + } + return &EinoChatModel{ + chain: chain, + labels: tags, + chatCfg: chatConfig, + }, nil +} + +// name returns the primary model's name (best-effort; nil-safe). func (m *EinoChatModel) name() string { - if m == nil || m.inner == nil || m.inner.ModelName == nil { + // chain[0] itself can be nil: NewEinoChatModel takes the caller's *ChatModel + // as-is (unlike the failover constructor, which filters nil entries), so a + // nil model must not turn Name() into a nil dereference. + if m == nil || len(m.chain) == 0 || m.chain[0] == nil || m.chain[0].ModelName == nil { return "" } - return *m.inner.ModelName + return *m.chain[0].ModelName } // toInternalMessages converts eino's []schema.Message into the existing @@ -182,17 +291,112 @@ func fromInternalResponse(resp *ChatResponse) *schema.Message { if len(resp.ToolCalls) > 0 { msg.ToolCalls = toolCallsFromInternal(resp.ToolCalls) } + if resp.Usage != nil { + // The call's token split travels ON THE MESSAGE (eino's own per-response + // metadata), so a caller that reports a node's cost reads it from the value it + // was handed instead of a field shared by every call on the ChatModel - where a + // concurrent call could replace it between the write and the read. + msg.ResponseMeta = &schema.ResponseMeta{ + Usage: &schema.TokenUsage{ + PromptTokens: resp.Usage.PromptTokens, + CompletionTokens: resp.Usage.CompletionTokens, + TotalTokens: resp.Usage.TotalTokens, + }, + } + } return msg } // Generate blocks until the model returns a complete response. Mirrors -// eino's model.BaseChatModel.Generate. +// eino's model.BaseChatModel.Generate. With a failover chain, a terminal +// error moves the call to the next entry (wrapping around once), and the +// sticky cursor stays on the entry that served the call. func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { - if m == nil || m.inner == nil || m.inner.ModelDriver == nil { + if m == nil || len(m.chain) == 0 { + return nil, fmt.Errorf("models: EinoChatModel: empty model chain") + } + // A full sweep that just failed stays failed for the cooldown: replaying + // the whole roster on every ReAct step of one question is what turns a + // dead plan into a wall of identical provider errors. + m.mu.Lock() + start := m.cursor + if waited := time.Since(m.sweep.failedAt); m.sweep.err != nil && waited < failoverCooldown { + err := m.sweep.err + m.mu.Unlock() + common.DebugCtx(ctx, "models: eino generate short-circuited by failover cooldown", + zap.Duration("waited", waited), zap.Error(err)) + return nil, err + } + m.mu.Unlock() + + var lastErr error + // cacheable stays true only while EVERY failed attempt was a provider-wide + // failure: a sweep that mixes a request-specific 400 with a transient 503 must not + // be remembered as "the whole roster is down", or the next request is rejected on + // the cooldown without trying a provider that could have served it. + cacheable := true + for i := 0; i < len(m.chain); i++ { + idx := (start + i) % len(m.chain) + cm := m.chain[idx] + if err := ctx.Err(); err != nil { + // The shared budget is spent: every remaining model would fail + // identically, so don't burn the chain on a dead context. + return nil, err + } + resp, err := m.generateOnce(ctx, cm, msgs, opts...) + if err == nil { + m.mu.Lock() + m.cursor = idx + m.sweep.failedAt = time.Time{} + m.sweep.err = nil + m.mu.Unlock() + return resp, nil + } + lastErr = err + cacheable = cacheable && cacheableSweepFailure(err) + next := (idx + 1) % len(m.chain) + common.WarnCtx(ctx, "models: eino generate failed, failing over to next model", + zap.String("failed_model", m.labelOf(idx)), + zap.String("next_model", m.labelOf(next)), + zap.Error(err)) + } + + // Terminal: every entry failed. Rotate the cursor so the next attempt starts + // on a different entry instead of always spending the primary's error budget + // first, and cache the outcome only when EVERY attempt was a provider-wide + // failure (see cacheable above), not merely the last one. + cached := lastErr != nil && cacheable + m.mu.Lock() + m.cursor = (start + 1) % len(m.chain) + if cached { + m.sweep.failedAt = time.Now() + m.sweep.err = lastErr + } + m.mu.Unlock() + common.WarnCtx(ctx, "models: eino generate failed on every model in the chain", + zap.Int("models", len(m.chain)), zap.Duration("cooldown", failoverCooldown), + zap.Bool("cached", cached), zap.Error(lastErr)) + return nil, lastErr +} + +// labelOf returns the human-readable chain entry tag (model @ instance). +func (m *EinoChatModel) labelOf(idx int) string { + if idx < 0 || idx >= len(m.chain) { + return "(out of range)" + } + if idx < len(m.labels) && m.labels[idx] != "" { + return m.labels[idx] + } + return modelNameOf(m.chain[idx]) +} + +// generateOnce runs one Generate attempt against a single chain entry. +func (m *EinoChatModel) generateOnce(ctx context.Context, cm *ChatModel, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + if cm == nil || cm.ModelDriver == nil { return nil, fmt.Errorf("models: EinoChatModel: nil inner ModelDriver") } internal := toInternalMessages(msgs) - if m.inner.ModelName == nil { + if cm.ModelName == nil { return nil, fmt.Errorf("models: EinoChatModel: nil model name") } // ChatWithMessages does not take a context.Context today — Phase 0 kept @@ -201,53 +405,120 @@ func (m *EinoChatModel) Generate(ctx context.Context, msgs []*schema.Message, op if err := ctx.Err(); err != nil { return nil, err } - // Reset stale per-call usage before the call so that a response - // without a usage block doesn't leak the previous call's data. - // Mirrors Python's LLMBundle._reset_last_usage(). - m.inner.LastUsage = nil chatCfg, err := m.chatConfigForGenerate() if err != nil { return nil, err } - if containsToolResult(internal) { + // eino's ChatModelAgent binds tools via the per-call model.WithTools option + // (not by calling WithTools). Merge those into the config so the model + // actually emits tool_calls; otherwise the ReAct loop would have no tools. + chatCfg, err = m.chatConfigWithOptsTools(chatCfg, opts) + if err != nil { + return nil, err + } + // Once a tool result is in context the model must be free to answer in + // prose. Code-exec agents pin tool_choice to "required"/execute_code, and + // keeping that pin on the follow-up turn would loop forever instead of + // returning the answer. + if chatCfg != nil && containsToolResult(internal) { choice := "auto" chatCfg.ToolChoice = &choice chatCfg.ToolChoiceValue = nil } - resp, err := m.inner.ModelDriver.ChatWithMessages(ctx, *m.inner.ModelName, internal, m.inner.APIConfig, chatCfg, nil) + common.Debug("models: eino generate request", + zap.String("model", *cm.ModelName), + zap.Int("messages", len(internal)), + zap.Int("tools", toolCount(chatCfg)), + ) + common.Debug("models: eino generate message skeleton", + zap.String("skeleton", describeInternalMessages(internal))) + resp, err := cm.ModelDriver.ChatWithMessages(ctx, *cm.ModelName, internal, cm.APIConfig, chatCfg, nil) if err != nil { - return nil, fmt.Errorf("models: EinoChatModel.Generate(%s): %w", *m.inner.ModelName, err) + return nil, fmt.Errorf("models: EinoChatModel.Generate(%s): %w", *cm.ModelName, err) } // Record the per-call token usage so the canvas-level aggregator (and // Langfuse) can compute the run total. Mirrors Python's // LLMBundle._report_usage() / self.mdl.last_usage pattern. if resp != nil && resp.Usage != nil { - m.inner.LastUsage = &TokenUsage{ + // The run sink gets this call's own split. The canvas component gets the same + // numbers from the message it is handed (ResponseMeta.Usage, set in + // fromInternalResponse): nothing about one call's usage lives on the shared + // ChatModel, where a concurrent call could replace it between write and read. + recordUsage(ctx, *cm.ModelName, &TokenUsage{ PromptTokens: resp.Usage.PromptTokens, CompletionTokens: resp.Usage.CompletionTokens, TotalTokens: resp.Usage.TotalTokens, - } - recordUsageFromResponse(ctx, m.inner) + }) } + // Guard the debug log against a nil resp: some drivers may return (nil, nil) + // on an aborted/empty completion, and len(resp.ToolCalls) would panic. + toolCalls := 0 + if resp != nil { + toolCalls = len(resp.ToolCalls) + } + // METADATA only: the answer body is user-visible (and often retrieved private) + // content, so a debug line must not copy it into the log file (CWE-532). + // Length + shape still separate an empty completion from a truncated one. + common.Debug("models: eino generate response", + zap.String("model", *cm.ModelName), + zap.Int("answer_bytes", len(answerHead(resp))), + zap.Int("tool_calls", toolCalls), + zap.Int("completion_tokens", usageCompletion(resp)), + ) return fromInternalResponse(resp), nil } -func containsToolResult(messages []Message) bool { - for _, message := range messages { - if message.Role == "tool" { - return true - } +// modelNameOf returns the model's display name (best-effort; nil-safe). +func modelNameOf(cm *ChatModel) string { + if cm == nil || cm.ModelName == nil { + return "(nil)" + } + return *cm.ModelName +} + +// answerHead returns a short preview of the response answer for log lines. +func answerHead(resp *ChatResponse) string { + if resp == nil || resp.Answer == nil { + return "" + } + return *resp.Answer +} + +func usageCompletion(resp *ChatResponse) int { + if resp == nil || resp.Usage == nil { + return 0 + } + return resp.Usage.CompletionTokens +} + +// toolCount returns the number of tools on a config, tolerating the +// interface{} storage type. +func toolCount(cfg *ChatConfig) int { + if cfg == nil { + return 0 + } + switch t := cfg.Tools.(type) { + case []map[string]any: + return len(t) + case nil: + return 0 + default: + return 0 } - return false } func (m *EinoChatModel) chatConfigForGenerate() (*ChatConfig, error) { - if len(m.tools) == 0 { - return m.chatCfg, nil - } + // Always hand back a COPY. Both callers (generateOnce, Stream) release the + // tool_choice on the config when the turn carries a tool result, and m.chatCfg + // is the shared base every WithTools/WithToolChoice instance carries: writing + // through it would persist one turn's "auto" into later turns (losing a + // required choice) and race concurrent calls on the same fields. cfg := &ChatConfig{} if m.chatCfg != nil { cp := *m.chatCfg cfg = &cp } + if len(m.tools) == 0 { + return cfg, nil + } tools, err := openAIToolsFromEino(m.tools) if err != nil { return nil, err @@ -303,6 +574,30 @@ func toolChoiceBody(choice string) map[string]any { } } +// chatConfigWithOptsTools overlays tools supplied via the per-call +// model.WithTools option onto base. eino's ChatModelAgent binds tools this way, +// so Generate/Stream must honor opts or the model will never emit tool_calls. +// When opts carries no tools, base is returned unchanged. +func (m *EinoChatModel) chatConfigWithOptsTools(base *ChatConfig, opts []model.Option) (*ChatConfig, error) { + co := model.GetCommonOptions(nil, opts...) + if co == nil || len(co.Tools) == 0 { + return base, nil + } + tools, err := openAIToolsFromEino(co.Tools) + if err != nil { + return nil, err + } + cfg := &ChatConfig{} + if base != nil { + cp := *base + cfg = &cp + } + cfg.Tools = tools + choice := "auto" + cfg.ToolChoice = &choice + return cfg, nil +} + func openAIToolsFromEino(infos []*schema.ToolInfo) ([]map[string]any, error) { tools := make([]map[string]any, 0, len(infos)) for _, info := range infos { @@ -330,6 +625,37 @@ func openAIToolsFromEino(infos []*schema.ToolInfo) ([]map[string]any, error) { return tools, nil } +// describeInternalMessages renders the message sequence's tool-call skeleton — +// "sys|user|asst[tc=ID1,ID2]|tool(ID1)|tool(ID2)|asst|user" — the exact shape a +// provider's tool-id validator sees, so a rejection like MiniMax's +// `tool result's tool id(X) not found` can be matched against the replayed +// sequence directly from the log instead of being reproduced blind. +func describeInternalMessages(msgs []Message) string { + var b strings.Builder + for i, m := range msgs { + if i > 0 { + b.WriteByte('|') + } + switch m.Role { + case "assistant": + b.WriteString("asst") + if len(m.ToolCalls) > 0 { + ids := make([]string, 0, len(m.ToolCalls)) + for _, tc := range m.ToolCalls { + id, _ := tc["id"].(string) + ids = append(ids, id) + } + b.WriteString(fmt.Sprintf("[tc=%s]", strings.Join(ids, ","))) + } + case "tool": + b.WriteString("tool(" + m.ToolCallID + ")") + default: + b.WriteString(m.Role) + } + } + return b.String() +} + func toolCallsToInternal(calls []schema.ToolCall) []map[string]interface{} { out := make([]map[string]interface{}, 0, len(calls)) for _, call := range calls { @@ -372,9 +698,17 @@ func toolCallsFromInternal(calls []map[string]interface{}) []schema.ToolCall { } } } + // Index MUST be set: stream consumers merge tool-call chunks by Index + // (nil reads as 0), and this message reaches them as one complete + // chunk carrying EVERY parallel call. With nil indexes the whole batch + // collapses into a single call — the last ID wins, the others' results + // are orphaned, and a provider replay (gate repair turn) rejects the + // sequence with `tool result's tool id(X) not found`. + idx := i out = append(out, schema.ToolCall{ - ID: id, - Type: callType, + ID: id, + Type: callType, + Index: &idx, Function: schema.FunctionCall{ Name: fnName, Arguments: fnArgs, @@ -384,19 +718,31 @@ func toolCallsFromInternal(calls []map[string]interface{}) []schema.ToolCall { return out } +// containsToolResult reports whether any message already carries a tool +// result. ReAct turns that contain results must keep tool_choice free: the +// model is expected to answer, not to fire another tool. +func containsToolResult(messages []Message) bool { + for _, message := range messages { + if message.Role == "tool" { + return true + } + } + return false +} + // Stream returns a schema.StreamReader that yields message chunks // incrementally. Uses the existing ChatStreamlyWithSender pathway; the // sender callback pushes the streamed delta into the StreamReader. +// With a failover chain, a stream that fails BEFORE any delta was emitted +// moves to the next model; once deltas have reached the client the call +// cannot be replayed, so the error is surfaced as-is. func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { - if m == nil || m.inner == nil || m.inner.ModelDriver == nil { - return nil, fmt.Errorf("models: EinoChatModel: nil inner ModelDriver") + if m == nil || len(m.chain) == 0 { + return nil, fmt.Errorf("models: EinoChatModel: empty model chain") } if err := ctx.Err(); err != nil { return nil, err } - if m.inner.ModelName == nil { - return nil, fmt.Errorf("models: EinoChatModel: nil model name") - } internalMessage := toInternalMessages(msgs) // Some OpenAI-compatible providers (including the configured MiniMax // endpoint) stream tool intent as ordinary prose. Use the provider's @@ -419,14 +765,30 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts if err != nil { return nil, err } - if containsToolResult(internalMessage) { + // Same tool-binding fix as Generate: honor the per-call model.WithTools + // option so the ReAct loop's model requests actually carry tool definitions. + chatCfg, err = m.chatConfigWithOptsTools(chatCfg, opts) + if err != nil { + return nil, err + } + // Same tool_choice release as generateOnce: a turn carrying tool results + // must be able to produce the final text instead of being forced to call + // a tool again. + if chatCfg != nil && containsToolResult(internalMessage) { choice := "auto" chatCfg.ToolChoice = &choice chatCfg.ToolChoiceValue = nil } + common.Debug("models: eino stream request", + zap.String("model", modelNameOf(m.chain[0])), + zap.Int("messages", len(internalMessage)), + zap.Int("tools", toolCount(chatCfg)), + zap.Int("chain", len(m.chain)), + ) sr, sw := schema.Pipe[*schema.Message](1) var sendMu sync.Mutex + var sentAny bool sender := func(content *string, reasoning *string) error { sendMu.Lock() defer sendMu.Unlock() @@ -449,21 +811,120 @@ func (m *EinoChatModel) Stream(ctx context.Context, msgs []*schema.Message, opts if closed := sw.Send(msg, nil); closed { return fmt.Errorf("models: stream closed before send completed") } + sentAny = true return nil } go func() { defer sw.Close() - if err := m.inner.ModelDriver.ChatStreamlyWithSender(ctx, *m.inner.ModelName, internalMessage, m.inner.APIConfig, chatCfg, nil, sender); err != nil { + m.mu.Lock() + start := m.cursor + if waited := time.Since(m.sweep.failedAt); m.sweep.err != nil && waited < failoverCooldown { + err := m.sweep.err + m.mu.Unlock() _ = sw.Send(nil, err) return } - if chatCfg != nil && chatCfg.ToolCallsResult != nil && len(*chatCfg.ToolCallsResult) > 0 { - msg := &schema.Message{ - Role: schema.Assistant, - ToolCalls: toolCallsFromInternal(*chatCfg.ToolCallsResult), + m.mu.Unlock() + var lastErr error + // Same rule as Generate: cache the sweep only when every failed attempt was + // provider-wide, never on the strength of the last error alone. + cacheable := true + for i := 0; i < len(m.chain); i++ { + idx := (start + i) % len(m.chain) + cm := m.chain[idx] + if cm == nil || cm.ModelDriver == nil { + // generateOnce reports this entry as a failed attempt ("nil inner + // ModelDriver"); mirror that here, and run it through the same + // cacheable rule. Skipping it silently left lastErr nil, so a + // chain of unusable entries reported success with no message. + lastErr = fmt.Errorf("models: EinoChatModel: nil inner ModelDriver on chain entry %d", idx) + cacheable = cacheable && cacheableSweepFailure(lastErr) + continue } - _ = sw.Send(msg, nil) + if cm.ModelName == nil { + // generateOnce reports this entry as a failed attempt; the + // streaming path used to dereference it, and a panic raised here + // runs inside a goroutine, which kills the process. It is a chain + // misconfiguration, not provider health, so it must not be cached + // as a sweep-wide failure either. + lastErr = fmt.Errorf("models: EinoChatModel: nil model name on chain entry %d", idx) + cacheable = cacheable && cacheableSweepFailure(lastErr) + continue + } + if err := ctx.Err(); err != nil { + _ = sw.Send(nil, err) + return + } + sendMu.Lock() + sentBefore := sentAny + sendMu.Unlock() + // Attempt-local config: the driver writes ToolCallsResult/UsageResult + // INTO the config it is handed, so reusing one instance would let a + // failed attempt's leftovers be attributed to the model that actually + // served the turn. + attemptCfg := &ChatConfig{} + if chatCfg != nil { + cp := *chatCfg + attemptCfg = &cp + } + attemptCfg.ToolCallsResult = nil + attemptCfg.UsageResult = nil + err := cm.ModelDriver.ChatStreamlyWithSender(ctx, *cm.ModelName, internalMessage, cm.APIConfig, attemptCfg, nil, sender) + if err == nil { + // Streamed turns report their token usage through the config + // (stream_options.include_usage), not through the nil modelUsage + // argument, so the run-level accumulator has to be fed from here — + // otherwise every streamed turn's tokens are missing from the total. + if attemptCfg.UsageResult != nil && attemptCfg.UsageResult.TotalTokens > 0 { + recordUsage(ctx, *cm.ModelName, attemptCfg.UsageResult) + } + if attemptCfg.ToolCallsResult != nil && len(*attemptCfg.ToolCallsResult) > 0 { + common.Debug("models: eino stream tool calls", + zap.String("model", *cm.ModelName), + zap.Int("tool_calls", len(*attemptCfg.ToolCallsResult))) + msg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: toolCallsFromInternal(*attemptCfg.ToolCallsResult), + } + _ = sw.Send(msg, nil) + } + m.mu.Lock() + m.cursor = idx + m.sweep.failedAt = time.Time{} + m.sweep.err = nil + m.mu.Unlock() + return + } + lastErr = err + cacheable = cacheable && cacheableSweepFailure(err) + sendMu.Lock() + sentAfter := sentAny + sendMu.Unlock() + if sentAfter != sentBefore || sentAfter { + // Deltas already reached the client: the stream cannot be + // replayed on another model, so fail as-is. + _ = sw.Send(nil, err) + return + } + next := (idx + 1) % len(m.chain) + common.WarnCtx(ctx, "models: eino stream failed before first delta, failing over to next model", + zap.String("failed_model", m.labelOf(idx)), + zap.String("next_model", m.labelOf(next)), + zap.Error(err)) } + // Terminal: sweep exhausted. Rotate, and cache only when EVERY attempt was a + // provider-wide failure - mirroring Generate. + cached := lastErr != nil && cacheable + m.mu.Lock() + m.cursor = (start + 1) % len(m.chain) + if cached { + m.sweep.failedAt = time.Now() + m.sweep.err = lastErr + } + m.mu.Unlock() + common.Debug("models: eino stream response error", + zap.String("model", modelNameOf(m.chain[0])), zap.Bool("cached", cached), zap.Error(lastErr)) + _ = sw.Send(nil, lastErr) }() return sr, nil } @@ -482,7 +943,18 @@ func (m *EinoChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingCh if m == nil { return nil, fmt.Errorf("models: EinoChatModel.WithTools: nil receiver") } - cp := *m + m.mu.Lock() + cursor := m.cursor + m.mu.Unlock() + // Field-by-field copy: EinoChatModel contains a sync.Mutex, which must + // never be copied (go vet copylocks). + cp := EinoChatModel{ + chain: m.chain, + labels: m.labels, + chatCfg: m.chatCfg, + toolChoice: m.toolChoice, + cursor: cursor, + } cp.tools = append([]*schema.ToolInfo(nil), tools...) return &cp, nil } @@ -493,7 +965,19 @@ func (m *EinoChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingCh // operates on the RAGFlow wrapper so the choice reaches the driver's request // body via chatConfigForGenerate. func (m *EinoChatModel) WithToolChoice(choice string) *EinoChatModel { - cp := *m + m.mu.Lock() + cursor := m.cursor + m.mu.Unlock() + // Field-by-field copy for the same reason as WithTools: cp := *m would copy + // the sync.Mutex (go vet copylocks) and could hand back a model whose mutex + // is permanently locked, deadlocking its next Generate/Stream/WithTools. + cp := EinoChatModel{ + chain: m.chain, + labels: m.labels, + chatCfg: m.chatCfg, + tools: append([]*schema.ToolInfo(nil), m.tools...), + cursor: cursor, + } if choice == "" { cp.toolChoice = nil return &cp @@ -511,14 +995,14 @@ func (m *EinoChatModel) Tools() []*schema.ToolInfo { return append([]*schema.ToolInfo(nil), m.tools...) } -// Inner exposes the wrapped *ChatModel for callers that need direct +// Inner exposes the primary wrapped *ChatModel for callers that need direct // access (e.g. to read token usage from the response after a custom // Generate call). Not part of any eino interface. func (m *EinoChatModel) Inner() *ChatModel { - if m == nil { + if m == nil || len(m.chain) == 0 { return nil } - return m.inner + return m.chain[0] } // Name returns the wrapped model name (used by tools / debugging). diff --git a/internal/entity/models/llm_test.go b/internal/entity/models/llm_test.go index 7200cb32c2..b437e5407e 100644 --- a/internal/entity/models/llm_test.go +++ b/internal/entity/models/llm_test.go @@ -3,12 +3,15 @@ package models import ( "context" "errors" + "fmt" "io" - "ragflow/internal/common" "reflect" "testing" + einomodel "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" + + "ragflow/internal/common" ) func TestEinoChatModelRequiresExecuteCodeTool(t *testing.T) { @@ -218,6 +221,59 @@ func TestEinoChatModelGenerateSendsBoundTools(t *testing.T) { } } +// TestEinoChatModelGenerateHonorsOptsTools verifies the per-call +// model.WithTools option (how eino's ChatModelAgent binds tools) reaches the +// driver. Without this, the ReAct loop's model requests would carry no tool +// definitions and the model would never emit tool_calls. +func TestEinoChatModelGenerateHonorsOptsTools(t *testing.T) { + apiKey := "key" + modelName := "chat" + driver := &captureToolDriver{ + resp: &ChatResponse{ + ToolCalls: []map[string]interface{}{ + { + "id": "call-1", "type": "function", + "function": map[string]interface{}{ + "name": "search_my_dateset", "arguments": `{"query":"hello"}`, + }, + }, + }, + }, + } + base := NewChatModel(driver, &modelName, &APIConfig{ApiKey: &apiKey}) + // NOTE: no WithTools binding — tools arrive only via the call option. + m := NewEinoChatModel(base, nil) + + toolInfo := &schema.ToolInfo{ + Name: "search_my_dateset", + Desc: "Search datasets.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: schema.String, Required: true}, + }), + } + msg, err := m.Generate(context.Background(), + []*schema.Message{schema.UserMessage("hello")}, + einomodel.WithTools([]*schema.ToolInfo{toolInfo}), + ) + if err != nil { + t.Fatalf("Generate: %v", err) + } + if driver.lastConfig == nil || driver.lastConfig.Tools == nil { + t.Fatal("Generate did not send opts-provided tools to driver") + } + tools, ok := driver.lastConfig.Tools.([]map[string]any) + if !ok || len(tools) != 1 { + t.Fatalf("driver tools = %#v, want one OpenAI-style tool", driver.lastConfig.Tools) + } + fn, _ := tools[0]["function"].(map[string]any) + if fn["name"] != "search_my_dateset" { + t.Fatalf("tool function name = %#v, want search_my_dateset", fn["name"]) + } + if len(msg.ToolCalls) != 1 || msg.ToolCalls[0].Function.Name != "search_my_dateset" { + t.Fatalf("msg tool calls = %#v, want search_my_dateset", msg.ToolCalls) + } +} + func TestEinoChatModelStreamWithToolsYieldsToolCalls(t *testing.T) { apiKey := "key" modelName := "chat" @@ -523,3 +579,29 @@ func TestToInternalMessagesUnsupportedPartsFallBackToString(t *testing.T) { t.Fatalf("Content = %#v, want string %q", internal[0].Content, "plain") } } + +// TestToolCallsFromInternalSetsIndex: the streamed tool-calls message reaches +// stream consumers as ONE complete chunk carrying every parallel call, and +// those consumers merge chunks by Index (nil reads as 0). Without an Index, +// a whole parallel batch collapses into a single call — the last ID wins and +// the other calls' results are orphaned on replay (MiniMax: +// `tool result's tool id(X) not found`). +func TestToolCallsFromInternalSetsIndex(t *testing.T) { + calls := []map[string]interface{}{ + {"id": "call_a_1", "type": "function", "function": map[string]interface{}{"name": "t1", "arguments": "{}"}}, + {"id": "call_a_2", "type": "function", "function": map[string]interface{}{"name": "t2", "arguments": "{}"}}, + {"id": "call_a_3", "type": "function", "function": map[string]interface{}{"name": "t3", "arguments": "{}"}}, + } + out := toolCallsFromInternal(calls) + if len(out) != 3 { + t.Fatalf("len = %d, want 3", len(out)) + } + for i, tc := range out { + if tc.Index == nil || *tc.Index != i { + t.Fatalf("call %d Index = %v, want %d", i, tc.Index, i) + } + if tc.ID != fmt.Sprintf("call_a_%d", i+1) { + t.Fatalf("call %d ID = %q", i, tc.ID) + } + } +} diff --git a/internal/entity/models/lmstudio.go b/internal/entity/models/lmstudio.go index 967394b005..eb1fdef3a0 100644 --- a/internal/entity/models/lmstudio.go +++ b/internal/entity/models/lmstudio.go @@ -95,7 +95,7 @@ func (l *LmStudioModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/localai.go b/internal/entity/models/localai.go index fb6ff32a08..cffb4a7289 100644 --- a/internal/entity/models/localai.go +++ b/internal/entity/models/localai.go @@ -89,7 +89,7 @@ func (l *LocalAIModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response via the diff --git a/internal/entity/models/longcat.go b/internal/entity/models/longcat.go index 3a48cc7d11..d381b1deba 100644 --- a/internal/entity/models/longcat.go +++ b/internal/entity/models/longcat.go @@ -102,7 +102,7 @@ func (l *LongCatModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response via the diff --git a/internal/entity/models/minimax.go b/internal/entity/models/minimax.go index 57057bfc04..15b77646e0 100644 --- a/internal/entity/models/minimax.go +++ b/internal/entity/models/minimax.go @@ -158,7 +158,7 @@ func (m *MinimaxModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, fmt.Errorf("minimax API error: %s", errMsg) } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/mistral.go b/internal/entity/models/mistral.go index 4cf94553f0..c6fa8bb9e6 100644 --- a/internal/entity/models/mistral.go +++ b/internal/entity/models/mistral.go @@ -153,7 +153,7 @@ func (m *MistralModel) ChatWithMessages(ctx context.Context, modelName string, m // to the flat string shape the shared handler understands. body = normalizeMistralStructuredContent(body) - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response diff --git a/internal/entity/models/model.go b/internal/entity/models/model.go index ff965070c1..5d360bfca4 100644 --- a/internal/entity/models/model.go +++ b/internal/entity/models/model.go @@ -170,11 +170,16 @@ type Model struct { URL string `json:"url"` MaxDimension *int `json:"max_dimension"` // used by embedding models MaxBatchSize *int `json:"max_batch_size"` // used by embedding models - Dimensions []int `json:"dimensions"` - BatchSize *int `json:"batch_size"` // max texts per Embed request; used by embedding models - Alias []string `json:"alias"` - Rank *int `json:"rank"` - ModelTypeMap map[string]bool + // Tokenizer names the tokenizer family that counts this embedding model's + // input (see internal/tokenizer: "xlmr-spm", "bert-wordpiece", "cl100k_base", + // ...). Empty or unknown means "we cannot count this model exactly"; callers + // then count with cl100k plus a calibrated safety ratio. + Tokenizer *string `json:"tokenizer"` + Dimensions []int `json:"dimensions"` + BatchSize *int `json:"batch_size"` // max texts per Embed request; used by embedding models + Alias []string `json:"alias"` + Rank *int `json:"rank"` + ModelTypeMap map[string]bool } // Provider represents an LLM provider @@ -482,6 +487,55 @@ func GetEmbeddingBatchSize(modelName string) int { return DefaultEmbeddingBatchSize } +// GetEmbeddingTokenizer returns the tokenizer id declared for the named model in +// the provider catalog (conf/all_models.json "tokenizer"), or "" when the model +// is unknown or declares none. The TOKENIZER_EMBEDDING_TOKENIZER env var overrides +// everything, so an operator can pin a family without editing the catalog. +// +// The returned value is an id understood by internal/tokenizer; whether an id whose asset +// is missing may fall back to a calibrated cl100k count is the CALLER's decision - the +// ingest embedder refuses (internal/ingestion/task/embedder.go), while a model that +// declares no tokenizer at all is calibrated by design. +func GetEmbeddingTokenizer(modelName string) string { + if v := strings.TrimSpace(os.Getenv("TOKENIZER_EMBEDDING_TOKENIZER")); v != "" { + return v + } + if modelName != "" { + if pm := GetProviderManager(); pm != nil { + if m := pm.GetModelByNameOrAlias(modelName); m != nil && m.Tokenizer != nil { + return strings.TrimSpace(*m.Tokenizer) + } + } + } + return "" +} + +// GetEmbeddingMaxTokens returns the input window declared for the named model: +// its own max_tokens first, then the provider catalog's context_length. It +// returns 0 when the catalog knows nothing, leaving the caller to apply its +// default — deliberately, because a wrong non-zero value here is a rejected +// request (an oversized input) rather than a shorter one. +func GetEmbeddingMaxTokens(modelName string) int { + if v := strings.TrimSpace(os.Getenv("TOKENIZER_EMBEDDING_MAX_TOKENS")); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + if modelName != "" { + if pm := GetProviderManager(); pm != nil { + if m := pm.GetModelByNameOrAlias(modelName); m != nil { + if m.MaxTokens != nil && *m.MaxTokens > 0 { + return *m.MaxTokens + } + if m.ContextLength != nil && *m.ContextLength > 0 { + return *m.ContextLength + } + } + } + } + return 0 +} + // 2. Show specific provider information (including base_url) func (pm *ProviderManager) GetProviderByName(providerName string) (map[string]interface{}, error) { diff --git a/internal/entity/models/model_test.go b/internal/entity/models/model_test.go index b3caf44365..d7c6c7b483 100644 --- a/internal/entity/models/model_test.go +++ b/internal/entity/models/model_test.go @@ -17,13 +17,21 @@ package models import ( + "bytes" "context" + "encoding/json" + "errors" + "io" + "io/fs" "net/http" "net/http/httptest" "os" "path/filepath" + "sort" "strings" "testing" + + "ragflow/internal/tokenizer" ) // joinModelNames extracts model names from a ListModelResponse slice and @@ -455,3 +463,230 @@ func TestSiliconFlowProviderConfigLoadsLatestProModels(t *testing.T) { t.Errorf("GLM-5.1 max_output=%d context_length=%d", *glm51.MaxOutput, *glm51.ContextLength) } } + +// TestAllModelsCatalogHasNoDuplicateKeys walks conf/all_models.json as a token stream and +// fails on any object key that occurs twice in the same object. +// +// encoding/json keeps only the last occurrence of a duplicated key, so a catalog that +// writes a field twice loads fine and behaves exactly as if it were written once - the +// duplicate is invisible to every test that goes through the parsed structs, which is how +// thirteen duplicated "tokenizer" fields got into this file and survived review. Only the +// raw token stream can see them, so the check lives here rather than in a schema validator. +func TestAllModelsCatalogHasNoDuplicateKeys(t *testing.T) { + // Control: the detector has to report a duplicate on input that has one, otherwise the + // catalog check below is a no-op that passes on anything. That is not hypothetical - + // the first version of the detector drove its state machine off the commas between + // members, and json.Decoder.Token does not emit them, so it recognised at most the + // first key of each object and reported zero duplicates for this very sample. + // The sample also pins that array elements are not mistaken for keys ("p", "p"). + control := []byte(`{"a": 1, "a": 2, "b": {"c": "x", "c": "y"}, "d": ["p", "p", {"e": 1, "e": 2}]}`) + if dups := duplicateJSONKeys(t, control); len(dups) != 3 { + t.Fatalf("detector control: got %d duplicates (%v), want 3 (a, c, e)", len(dups), dups) + } + + target := filepath.Join(findRepoRoot(), "conf", "all_models.json") + data, err := os.ReadFile(target) + if err != nil { + t.Fatalf("read %s: %v", target, err) + } + for _, dup := range duplicateJSONKeys(t, data) { + t.Errorf("%s:%d: object key %q appears twice (byte %d); encoding/json keeps only the last one, so the file means something other than it says", + target, dup.line, dup.key, dup.offset) + } +} + +type duplicateKey struct { + key string + offset int64 + line int +} + +// duplicateJSONKeys returns every key that occurs twice inside the same JSON object. +func duplicateJSONKeys(t *testing.T, data []byte) []duplicateKey { + t.Helper() + + // One frame per open object or array. expectKey means "the next string in this object + // is a key": inside an object Decoder.Token yields key, value, key, value ... because + // the commas are not tokens, so after every value one has to expect a key again. + type frame struct { + object bool + expectKey bool + keys map[string]bool + } + var ( + stack []frame + dups []duplicateKey + ) + + // afterValue marks that the frame on top of the stack has just received a value. + afterValue := func() { + if n := len(stack); n > 0 { + stack[n-1].expectKey = true + } + } + + dec := json.NewDecoder(bytes.NewReader(data)) + for { + tok, err := dec.Token() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + t.Fatalf("decode catalog: %v", err) + } + switch v := tok.(type) { + case json.Delim: + switch v { + case '{', '[': + // The enclosing object just received a value: whatever follows it is a key, + // not a value. The nested frame tracks its own keys from here. + if n := len(stack); n > 0 { + stack[n-1].expectKey = false + } + stack = append(stack, frame{object: v == '{', expectKey: v == '{', keys: map[string]bool{}}) + case '}', ']': + if len(stack) > 0 { + stack = stack[:len(stack)-1] + } + afterValue() + } + continue + case string: + if n := len(stack); n > 0 && stack[n-1].object && stack[n-1].expectKey { + top := &stack[n-1] + if top.keys[v] { + offset := dec.InputOffset() + dups = append(dups, duplicateKey{key: v, offset: offset, line: lineAtOffset(data, offset)}) + } + top.keys[v] = true + top.expectKey = false + continue + } + } + afterValue() + } + return dups +} + +// lineAtOffset reports the 1-based line that holds a byte offset. +func lineAtOffset(data []byte, offset int64) int { + if offset < 0 { + offset = 0 + } + if offset > int64(len(data)) { + offset = int64(len(data)) + } + return 1 + bytes.Count(data[:offset], []byte{'\n'}) +} + +// TestConfFilesHaveNoDuplicateKeys applies the duplicate-key check to every JSON file +// under conf/. +// +// A definition written into a file that already had one does not win: the parser keeps the +// last occurrence, so the older copy stays in force and the newer one is silently ignored. +// That is not hypothetical - conf/infinity_mapping.json carried two deleted_doc_id +// definitions, and the duplicate key meant the column the doc-delete fix (#17685) added - +// with its analyzer - never took effect. +func TestConfFilesHaveNoDuplicateKeys(t *testing.T) { + confDir := filepath.Join(findRepoRoot(), "conf") + var files []string + err := filepath.WalkDir(confDir, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + if !entry.IsDir() && strings.HasSuffix(path, ".json") { + files = append(files, path) + } + return nil + }) + if err != nil { + t.Fatalf("walk %s: %v", confDir, err) + } + if len(files) == 0 { + t.Fatalf("no JSON file under %s; the check would pass vacuously", confDir) + } + for _, path := range files { + data, err := os.ReadFile(path) + if err != nil { + t.Errorf("read %s: %v", path, err) + continue + } + for _, dup := range duplicateJSONKeys(t, data) { + t.Errorf("%s:%d: object key %q appears twice (byte %d); encoding/json keeps only the last one, so the file means something other than it says", + path, dup.line, dup.key, dup.offset) + } + } + t.Logf("checked %d JSON files under conf/", len(files)) +} + +// TestAllModelsCatalogTokenizerTagsAreKnown checks the catalog's half of the tokenizer +// link: every "tokenizer" value it declares has to name a counter the Go side knows. +// +// A tag that names nothing fails nowhere. ResolveCounter falls back to cl100k_base, the +// ingest path counts with the calibrated estimate, and the only trace is a lower-precision +// count - which is the degradation this field exists to prevent. So a typo in a +// hand-edited catalog is invisible unless the two lists are compared, which is this. +// +// The known ids come from CounterStatuses, not from CounterByID: the bool CounterByID +// returns means "this counter is available", so on a checkout that has not downloaded the +// tokenizer assets it is false even for a valid tag. Keying this test off it would fail on +// every machine without the assets - for the wrong reason. CounterStatuses enumerates the +// ids whether or not their assets are present. +func TestAllModelsCatalogTokenizerTagsAreKnown(t *testing.T) { + known := map[string]bool{} + for _, status := range tokenizer.CounterStatuses() { + known[status.ID] = true + } + if len(known) == 0 { + t.Fatal("tokenizer.CounterStatuses reports no counters; the check would pass vacuously") + } + + target := filepath.Join(findRepoRoot(), "conf", "all_models.json") + data, err := os.ReadFile(target) + if err != nil { + t.Fatalf("read %s: %v", target, err) + } + var catalog map[string]any + if err := json.Unmarshal(data, &catalog); err != nil { + t.Fatalf("parse %s: %v", target, err) + } + + knownIDs := sortedKeys(known) + tags, declared := map[string]int{}, 0 + for _, section := range catalog { + entries, ok := section.([]any) + if !ok { + continue + } + for _, entry := range entries { + model, ok := entry.(map[string]any) + if !ok { + continue + } + tag, ok := model["tokenizer"].(string) + if !ok || tag == "" { + continue + } + declared++ + tags[tag]++ + if !known[tag] { + t.Errorf("%s: model %v declares tokenizer %q, which is not a counter id %v; ingest will silently count it with the calibrated estimate", + target, model["name"], tag, knownIDs) + } + } + } + if declared == 0 { + t.Fatalf("%s declares no tokenizer for any model; the check would pass vacuously", target) + } + t.Logf("catalog tokenizer tags: %v", tags) +} + +// sortedKeys returns the keys of a set, sorted, for a deterministic message. +func sortedKeys(set map[string]bool) []string { + out := make([]string, 0, len(set)) + for key := range set { + out = append(out, key) + } + sort.Strings(out) + return out +} diff --git a/internal/entity/models/modelscope.go b/internal/entity/models/modelscope.go index 1ff6302ab0..344f561a0c 100644 --- a/internal/entity/models/modelscope.go +++ b/internal/entity/models/modelscope.go @@ -90,7 +90,7 @@ func (m *ModelScopeModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender. diff --git a/internal/entity/models/moonshot.go b/internal/entity/models/moonshot.go index dc72cc7157..933d3e687e 100644 --- a/internal/entity/models/moonshot.go +++ b/internal/entity/models/moonshot.go @@ -119,7 +119,7 @@ func (m *MoonshotModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/mws.go b/internal/entity/models/mws.go index 595a639932..ab5f26677d 100644 --- a/internal/entity/models/mws.go +++ b/internal/entity/models/mws.go @@ -178,7 +178,7 @@ func (m *MWSModel) ChatWithMessages(ctx context.Context, modelName string, messa if err != nil { return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends a streaming MWS Chat Completions request. diff --git a/internal/entity/models/n1n.go b/internal/entity/models/n1n.go index 08f0d46e13..766fb939f1 100644 --- a/internal/entity/models/n1n.go +++ b/internal/entity/models/n1n.go @@ -120,7 +120,7 @@ func (n *N1NModel) ChatWithMessages(ctx context.Context, modelName string, messa return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends a streaming chat completion. diff --git a/internal/entity/models/novita.go b/internal/entity/models/novita.go index 90a2533478..f3a83636d5 100644 --- a/internal/entity/models/novita.go +++ b/internal/entity/models/novita.go @@ -80,7 +80,7 @@ func (n *NovitaModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response via diff --git a/internal/entity/models/nvidia.go b/internal/entity/models/nvidia.go index cefe2312c8..99fbc6079a 100644 --- a/internal/entity/models/nvidia.go +++ b/internal/entity/models/nvidia.go @@ -117,7 +117,7 @@ func (n *NvidiaModel) ChatWithMessages(ctx context.Context, modelName string, me if err != nil { return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (n *NvidiaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/openai.go b/internal/entity/models/openai.go index a682716052..3f5ffc9813 100644 --- a/internal/entity/models/openai.go +++ b/internal/entity/models/openai.go @@ -90,7 +90,7 @@ func (o *OpenAIModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response @@ -138,6 +138,7 @@ func (o *OpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName stri type openaiEmbeddingResponse struct { Data []openaiEmbeddingData `json:"data"` + ID string `json:"id"` Model string `json:"model"` Object string `json:"object"` Usage openaiUsage `json:"usage"` @@ -230,6 +231,15 @@ func (o *OpenAIModel) Embed(ctx context.Context, modelName *string, request Embe embeddings = append(embeddings, embeddingData) } + // The provider's own token count is the oracle the ingest path calibrates + // its tokenizer against (see internal/tokenizer: ObserveUsage). Reporting it + // also fixes token accounting, which previously saw 0 tokens for every + // embedding call on this driver. + recordResponseUsage(modelUsage, parsed.ID, &TokenUsage{ + PromptTokens: parsed.Usage.PromptTokens, + TotalTokens: parsed.Usage.TotalTokens, + }, "embedding") + return embeddings, nil } diff --git a/internal/entity/models/openai_api_compatible.go b/internal/entity/models/openai_api_compatible.go index 764b45be6c..3ff7117a7a 100644 --- a/internal/entity/models/openai_api_compatible.go +++ b/internal/entity/models/openai_api_compatible.go @@ -104,7 +104,7 @@ func (m *OpenAIAPICompatibleModel) ChatWithMessages(ctx context.Context, modelNa return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function diff --git a/internal/entity/models/openrouter.go b/internal/entity/models/openrouter.go index dc7300c090..3e9330dd81 100644 --- a/internal/entity/models/openrouter.go +++ b/internal/entity/models/openrouter.go @@ -82,7 +82,7 @@ func (o *OpenRouterModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/orcarouter.go b/internal/entity/models/orcarouter.go index 7f9e80590d..cecd2f9153 100644 --- a/internal/entity/models/orcarouter.go +++ b/internal/entity/models/orcarouter.go @@ -74,7 +74,7 @@ func (o *OrcaRouterModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (o *OrcaRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/perplexity.go b/internal/entity/models/perplexity.go index 4c83f86f74..848ddf492f 100644 --- a/internal/entity/models/perplexity.go +++ b/internal/entity/models/perplexity.go @@ -89,7 +89,7 @@ func (p *PerplexityModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (p *PerplexityModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/ppio.go b/internal/entity/models/ppio.go index 9667cc6b2f..f0b3548b6c 100644 --- a/internal/entity/models/ppio.go +++ b/internal/entity/models/ppio.go @@ -82,7 +82,7 @@ func (p *PPIOModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (p *PPIOModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/qiniu.go b/internal/entity/models/qiniu.go index f858100c28..029486d75f 100644 --- a/internal/entity/models/qiniu.go +++ b/internal/entity/models/qiniu.go @@ -163,7 +163,7 @@ func (q *QiniuModel) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (q *QiniuModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/ragcon.go b/internal/entity/models/ragcon.go index a3176b7339..5cf32b8bd9 100644 --- a/internal/entity/models/ragcon.go +++ b/internal/entity/models/ragcon.go @@ -92,7 +92,7 @@ func (r *RAGconModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (r *RAGconModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/response_handler.go b/internal/entity/models/response_handler.go index ea537ccb1e..8f624f35f4 100644 --- a/internal/entity/models/response_handler.go +++ b/internal/entity/models/response_handler.go @@ -17,6 +17,7 @@ package models import ( + "context" "encoding/json" "fmt" "io" @@ -27,7 +28,7 @@ import ( // HandleNonStreamingResponse processes a complete non-streaming chat // response using the ParserConfig's ResponseParser to extract usage. -func HandleNonStreamingResponse(body []byte, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig) (*ChatResponse, error) { +func HandleNonStreamingResponse(ctx context.Context, body []byte, modelUsage *common.ModelUsage, chatConfig *ChatConfig, cfg *ParserConfig) (*ChatResponse, error) { var result map[string]any if err := json.Unmarshal(body, &result); err != nil { return nil, fmt.Errorf("failed to parse response: %w", err) @@ -56,7 +57,7 @@ func HandleNonStreamingResponse(body []byte, modelUsage *common.ModelUsage, chat if chatConfig != nil { chatConfig.UsageResult = usage model, _ := result["model"].(string) - common.Info("NonStreamUsage", zap.String("model", model), zap.Int("prompt", usage.PromptTokens), zap.Int("completion", usage.CompletionTokens), zap.Int("total", usage.TotalTokens)) + common.InfoCtx(ctx, "NonStreamUsage", zap.String("model", model), zap.Int("prompt", usage.PromptTokens), zap.Int("completion", usage.CompletionTokens), zap.Int("total", usage.TotalTokens)) } } diff --git a/internal/entity/models/siliconflow.go b/internal/entity/models/siliconflow.go index 8d51a1b8c8..7cf2d4ecb0 100644 --- a/internal/entity/models/siliconflow.go +++ b/internal/entity/models/siliconflow.go @@ -87,7 +87,7 @@ func (s *SiliconflowModel) ChatWithMessages(ctx context.Context, modelName strin if err != nil { return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) @@ -211,6 +211,14 @@ func (s *SiliconflowModel) Embed(ctx context.Context, modelName *string, request } if resp.StatusCode != http.StatusOK { + // Carry Retry-After in the message when the provider sends one: a + // TPM-limited provider knows when the quota window refills, and the + // retry layer in internal/ingestion/task/embedder.go would otherwise + // have to guess a delay that is either too short (another 429) or far + // longer than necessary. + if retryAfter := strings.TrimSpace(resp.Header.Get("Retry-After")); retryAfter != "" { + return nil, fmt.Errorf("SILICONFLOW API error: %s, retry-after: %s, body: %s", resp.Status, retryAfter, string(body)) + } return nil, fmt.Errorf("SILICONFLOW API error: %s, body: %s", resp.Status, string(body)) } diff --git a/internal/entity/models/stepfun.go b/internal/entity/models/stepfun.go index ffcdf77a0c..16485de0e1 100644 --- a/internal/entity/models/stepfun.go +++ b/internal/entity/models/stepfun.go @@ -75,7 +75,7 @@ func (s *StepFunModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response diff --git a/internal/entity/models/togetherai.go b/internal/entity/models/togetherai.go index dce137aaea..cdddb68a97 100644 --- a/internal/entity/models/togetherai.go +++ b/internal/entity/models/togetherai.go @@ -106,7 +106,7 @@ func (t *TogetherAIModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (t *TogetherAIModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/tokenhub.go b/internal/entity/models/tokenhub.go index 4eaa0d4238..c9223c3cea 100644 --- a/internal/entity/models/tokenhub.go +++ b/internal/entity/models/tokenhub.go @@ -79,7 +79,7 @@ func (t *TokenHubModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (t *TokenHubModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/tokenpony.go b/internal/entity/models/tokenpony.go index 5d9644784e..1f9f305de8 100644 --- a/internal/entity/models/tokenpony.go +++ b/internal/entity/models/tokenpony.go @@ -72,7 +72,7 @@ func (t *TokenPonyModel) ChatWithMessages(ctx context.Context, modelName string, return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender opens the SSE chat-completions diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 50da4a8331..07f1aab413 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -93,7 +93,11 @@ type TokenUsage struct { type EmbeddingData struct { Embedding []float64 `json:"embedding"` Index int `json:"index"` - // FIXME: add implementation + // TokenCount is what this input cost, taken from the provider's reported + // usage. Embedding APIs report usage per *request*, not per input, so the + // ingest path distributes the request total across the inputs it sent + // (internal/ingestion/task/embedder.go). It stays 0 for providers that + // report no usage at all — by design, rather than inventing a number. TokenCount int `json:"token_count"` } @@ -294,6 +298,39 @@ func (m *EmbeddingModel) ResolveBatchSize() int { return GetEmbeddingBatchSize(name) } +// ResolveMaxTokens is ResolveBatchSize's counterpart for the input window: the +// model's own declaration wins, then the provider catalog's context_length, then +// 0, which tells the caller to apply its own default. Deliberately not 8192: +// the catalog has embedding models with 512-token windows, and overshooting a +// window is a rejected request while undershooting only truncates. +func (m *EmbeddingModel) ResolveMaxTokens() int { + if m == nil { + return 0 + } + if m.MaxTokens > 0 { + return m.MaxTokens + } + var name string + if m.ModelName != nil { + name = *m.ModelName + } + return GetEmbeddingMaxTokens(name) +} + +// ResolveTokenizerID returns the tokenizer family declared for this model, or "" +// when the model's own tokenizer is unknown (the caller then counts with cl100k +// and a calibrated ratio). +func (m *EmbeddingModel) ResolveTokenizerID() string { + if m == nil { + return "" + } + var name string + if m.ModelName != nil { + name = *m.ModelName + } + return GetEmbeddingTokenizer(name) +} + // RerankModel wraps a ModelDriver with rerank-specific configuration type RerankModel struct { ModelDriver ModelDriver @@ -369,10 +406,6 @@ type ChatModel struct { ModelName *string APIConfig *APIConfig ToolConfig *ToolConfig - // LastUsage holds the token usage (prompt/completion/total) of the most - // recent chat call. Consumed by callers for accurate Langfuse reporting - // and per-run token aggregation. Reset before each call. - LastUsage *TokenUsage } // NewChatModel creates a new ChatModel diff --git a/internal/entity/models/upstage.go b/internal/entity/models/upstage.go index 5a43ca0af0..20b7818832 100644 --- a/internal/entity/models/upstage.go +++ b/internal/entity/models/upstage.go @@ -80,7 +80,7 @@ func (u *UpstageModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response diff --git a/internal/entity/models/vllm.go b/internal/entity/models/vllm.go index 514e3f0c40..391310f278 100644 --- a/internal/entity/models/vllm.go +++ b/internal/entity/models/vllm.go @@ -117,7 +117,7 @@ func (v *VllmModel) ChatWithMessages(ctx context.Context, modelName string, mess return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/volcengine.go b/internal/entity/models/volcengine.go index e4e89533da..a98a267de0 100644 --- a/internal/entity/models/volcengine.go +++ b/internal/entity/models/volcengine.go @@ -153,7 +153,7 @@ func (v *VolcEngine) ChatWithMessages(ctx context.Context, modelName string, mes return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/entity/models/xai.go b/internal/entity/models/xai.go index 00ae4b2b20..9757dabae0 100644 --- a/internal/entity/models/xai.go +++ b/internal/entity/models/xai.go @@ -85,7 +85,7 @@ func (x *XAIModel) ChatWithMessages(ctx context.Context, modelName string, messa return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams the response diff --git a/internal/entity/models/xiaomi.go b/internal/entity/models/xiaomi.go index 2413c0d352..9121a5807b 100644 --- a/internal/entity/models/xiaomi.go +++ b/internal/entity/models/xiaomi.go @@ -101,7 +101,7 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/xinference.go b/internal/entity/models/xinference.go index d0ec2f1a7c..b072e44f7e 100644 --- a/internal/entity/models/xinference.go +++ b/internal/entity/models/xinference.go @@ -94,7 +94,7 @@ func (x *XinferenceModel) ChatWithMessages(ctx context.Context, modelName string return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender. diff --git a/internal/entity/models/xunfei.go b/internal/entity/models/xunfei.go index f08a1c33a0..5cd35cb409 100644 --- a/internal/entity/models/xunfei.go +++ b/internal/entity/models/xunfei.go @@ -126,7 +126,7 @@ func (x *XunFeiModel) ChatWithMessages(ctx context.Context, modelName string, me return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } func (x *XunFeiModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 8e2993ed4c..a99ad0fad5 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -92,7 +92,7 @@ func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, m return nil, err } - return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) + return HandleNonStreamingResponse(ctx, body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/handler/chat_audio.go b/internal/handler/chat_audio.go index e08f3061e9..c281ab0562 100644 --- a/internal/handler/chat_audio.go +++ b/internal/handler/chat_audio.go @@ -279,7 +279,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) { streamMode := strings.ToLower(c.PostForm("stream")) == "true" if streamMode { - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") diff --git a/internal/handler/chat_session.go b/internal/handler/chat_session.go index 7af16f44e5..99cd0bce98 100644 --- a/internal/handler/chat_session.go +++ b/internal/handler/chat_session.go @@ -247,7 +247,7 @@ func (h *ChatSessionHandler) ChatCompletions(c *gin.Context) { } if streamMode { - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") @@ -276,6 +276,13 @@ func (h *ChatSessionHandler) ChatCompletions(c *gin.Context) { return true }) } else { + // The non-stream path computes the whole answer before its first + // (and only) write. Agentic runs can take minutes — well past the + // server's http.Server.WriteTimeout, which starts when the request + // headers are read and is not reset by later writes. Clear the write + // deadline so the final JSON write is not aborted with + // "write tcp ... i/o timeout" after a long agent run. + clearResponseWriteDeadline(c) var result map[string]interface{} result, err = h.chatSessionService.ChatCompletions( c.Request.Context(), userID, diff --git a/internal/handler/openai_chat.go b/internal/handler/openai_chat.go index 495684050e..a4aff9bc60 100644 --- a/internal/handler/openai_chat.go +++ b/internal/handler/openai_chat.go @@ -133,5 +133,8 @@ func (h *OpenAIChatHandler) OpenAIChatCompletions(c *gin.Context) { // All early-rejection checks passed. Delegate to the service for the // actual LLM call. + // Long agentic runs stream for minutes or compute before their single + // write; clear http.Server.WriteTimeout so neither is cut off mid-response. + clearResponseWriteDeadline(c) h.svc.OpenAIChatCompletions(c, user.ID, chatID, bodyBytes) } diff --git a/internal/handler/providers.go b/internal/handler/providers.go index 6b42dd7864..7a57a2f540 100644 --- a/internal/handler/providers.go +++ b/internal/handler/providers.go @@ -986,7 +986,7 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) { // Check if it's a stream request if req.Stream { // Set SSE headers - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") @@ -1270,7 +1270,7 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) { // Check if it's a stream request if req.Stream { // Set SSE headers - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") @@ -1378,7 +1378,7 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) { // Check if it's a stream request if req.Stream { // Set SSE headers - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") diff --git a/internal/handler/search.go b/internal/handler/search.go index 12b6cebfb9..7677558111 100644 --- a/internal/handler/search.go +++ b/internal/handler/search.go @@ -408,7 +408,7 @@ func (h *SearchHandler) Completion(c *gin.Context) { return } - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream; charset=utf-8") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") diff --git a/internal/handler/searchbot.go b/internal/handler/searchbot.go index e1ee4beb14..c6ab68819e 100644 --- a/internal/handler/searchbot.go +++ b/internal/handler/searchbot.go @@ -283,7 +283,7 @@ func (h *SearchBotHandler) Ask(c *gin.Context) { return } - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Header("Cache-Control", "no-cache") c.Header("Connection", "keep-alive") diff --git a/internal/handler/streaming.go b/internal/handler/streaming.go index b59f8e30a6..e7bffc1925 100644 --- a/internal/handler/streaming.go +++ b/internal/handler/streaming.go @@ -23,6 +23,10 @@ import ( "github.com/gin-gonic/gin" ) -func disableWriteDeadlineForSSE(c *gin.Context) { +// clearResponseWriteDeadline removes http.Server.WriteTimeout from the current +// response so handlers that stream over minutes or compute before their first +// write (SSE, agentic chats) are not cut off mid-response by the connection +// write deadline. +func clearResponseWriteDeadline(c *gin.Context) { _ = http.NewResponseController(c.Writer).SetWriteDeadline(time.Time{}) } diff --git a/internal/handler/streaming_test.go b/internal/handler/streaming_test.go index 1237a7abb5..6e0b1422bd 100644 --- a/internal/handler/streaming_test.go +++ b/internal/handler/streaming_test.go @@ -27,12 +27,12 @@ import ( "github.com/gin-gonic/gin" ) -func TestDisableWriteDeadlineForSSEAllowsLongLivedStream(t *testing.T) { +func TestClearWriteDeadlineAllowsLongLivedStream(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.GET("/stream", func(c *gin.Context) { - disableWriteDeadlineForSSE(c) + clearResponseWriteDeadline(c) c.Header("Content-Type", "text/event-stream") c.Writer.WriteHeader(http.StatusOK) c.Writer.Flush() @@ -77,3 +77,41 @@ func TestDisableWriteDeadlineForSSEAllowsLongLivedStream(t *testing.T) { } } } + +// TestClearWriteDeadlineAllowsLateSingleWrite pins the non-stream counterpart: +// a handler that computes past http.Server.WriteTimeout before writing its one +// JSON response must still deliver the body once the deadline is cleared — the +// exact failure mode of long agentic chat runs ("write tcp ... i/o timeout"). +func TestClearWriteDeadlineAllowsLateSingleWrite(t *testing.T) { + gin.SetMode(gin.TestMode) + + router := gin.New() + router.POST("/complete", func(c *gin.Context) { + clearResponseWriteDeadline(c) + + time.Sleep(120 * time.Millisecond) + + c.JSON(http.StatusOK, gin.H{"answer": "done"}) + }) + + server := httptest.NewUnstartedServer(router) + server.Config.WriteTimeout = 30 * time.Millisecond + server.Start() + defer server.Close() + + client := server.Client() + client.Timeout = time.Second + resp, err := client.Post(server.URL+"/complete", "application/json", strings.NewReader(`{}`)) + if err != nil { + t.Fatalf("post complete: %v", err) + } + defer resp.Body.Close() + + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + if got := string(bodyBytes); !strings.Contains(got, `"answer":"done"`) { + t.Fatalf("response body missing answer after exceeded WriteTimeout: %q", got) + } +} diff --git a/internal/harness/events/nats.go b/internal/harness/events/nats.go index 1e5d94d7e0..9d5593981a 100644 --- a/internal/harness/events/nats.go +++ b/internal/harness/events/nats.go @@ -48,21 +48,31 @@ func NewNATSEventStore(conn *nats.Conn, stream string) (*NATSEventStore, error) return nil, fmt.Errorf("jetstream init: %w", err) } + streamCfg := jetstream.StreamConfig{ + Name: stream, + Subjects: []string{fmt.Sprintf("%s.>", defaultNATSPrefix)}, + MaxAge: 7 * 24 * time.Hour, // 7 days retention + Storage: jetstream.FileStorage, + Retention: jetstream.LimitsPolicy, + Discard: jetstream.DiscardNew, + MaxMsgs: 1024 * 1024, + MaxBytes: 1024 * 1024 * 1024, + } // Ensure the stream exists. _, err = js.Stream(ctxForInit(), stream) if err != nil { // Create the stream if it doesn't exist. - _, err = js.CreateStream(ctxForInit(), jetstream.StreamConfig{ - Name: stream, - Subjects: []string{fmt.Sprintf("%s.>", defaultNATSPrefix)}, - MaxAge: 7 * 24 * time.Hour, // 7 days retention - Storage: jetstream.FileStorage, - Retention: jetstream.LimitsPolicy, - Discard: jetstream.DiscardNew, - }) + _, err = js.CreateStream(ctxForInit(), streamCfg) if err != nil { return nil, fmt.Errorf("create jetstream stream: %w", err) } + } else { + // Reconcile an existing stream created under older settings with the + // current config (DiscardNew + larger MaxMsgs/MaxBytes). CreateStream + // never touches an already-existing stream, so UpdateStream is required. + if _, err = js.UpdateStream(ctxForInit(), streamCfg); err != nil { + return nil, fmt.Errorf("update jetstream stream: %w", err) + } } return &NATSEventStore{ diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 051075cd85..29fffb6d8a 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -1507,7 +1507,8 @@ func defaultChatModelRef(ctx context.Context, db *gorm.DB, tenantID string) stri } // extractorContextFitBudget returns 97% of the model's context window as the -// fitting budget, mirroring the agent component's contextFitBudget. The +// fitting budget, mirroring agent/chat's ContextFitBudget (the canonical +// helper) plus the clamp below. The // margin leaves headroom for the difference between the cl100k tokenizer used // for counting and the model's own tokenizer, plus per-message formatting // overhead, so a fitted prompt stays inside the provider's real context limit diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index 527047283e..62ab750b34 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -154,6 +154,43 @@ type Embedder interface { Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error) } +// EmbedderTrimmer is an OPTIONAL extension of Embedder, implemented by embedders +// that can count with the embedding model's own tokenizer (or a calibrated upper +// bound of it) instead of cl100k_base. When present the component uses it for +// every truncation, because cl100k and a model's own tokenizer disagree by up to +// ~2% and the sign of the disagreement depends on the content: a chunk that +// cl100k scores just under the limit can be over it in the model's tokenizer, +// which the provider answers with a 400 and, before this, the loss of the whole +// document. See internal/tokenizer/embedding_token_limits.md. +type EmbedderTrimmer interface { + Trim(text string) (trimmed string, tokens int) +} + +// EmbedderMaxResolver is the optional counterpart that exposes the input window +// to honour: the model's declared value, then the provider catalog's +// context_length. Without it the component has to guess, and guessing 8192 +// overshoots the window of every model with a smaller one. +type EmbedderMaxResolver interface { + ResolveMaxTokens() int +} + +// trimForEmbedding trims `text` the best way the embedder allows. +func trimForEmbedding(embedder Embedder, text string) string { + if trimmer, ok := embedder.(EmbedderTrimmer); ok { + trimmed, _ := trimmer.Trim(text) + return trimmed + } + return truncateForEmbedding(text, resolveEmbedderMaxTokens(embedder)) +} + +// resolveEmbedderMaxTokens is the window the embedder will actually honour. +func resolveEmbedderMaxTokens(embedder Embedder) int { + if resolver, ok := embedder.(EmbedderMaxResolver); ok { + return resolver.ResolveMaxTokens() + } + return embedder.MaxTokens() +} + // EmbedderResolver resolves the embedder and its dataset-bound embedding-model // id for one tokenizer invocation. The resolver derives the model exclusively // from the knowledgebase's configured embd_id (see internal/ingestion/task/ @@ -444,7 +481,7 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na // stale embedding for up to the cache TTL. embdID must be non-empty or // every model would collapse onto one key and served vectors could come // from a different model. - trunc := truncateForEmbedding(txt, embedder.MaxTokens()) + trunc := trimForEmbedding(embedder, txt) if chunkID, ok := ck.GetExtraString("id"); ok && embdID != "" && store != nil { if cached, hit := chunkcache.Get(ctx, store, chunkcache.Key("emb", embdID, chunkID, trunc)); hit { var vec []float64 @@ -559,34 +596,31 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, na } // defaultEmbeddingTokenLimit is the safe fallback used when an embedder reports -// no token limit (maxTokens <= 0). It both prevents empty embedding inputs and -// keeps truncation active for every path instead of passing the full text through. -const defaultEmbeddingTokenLimit = 8192 +// no token limit at all. It is deliberately conservative: overshooting a model's +// window is a rejected request, while undershooting only truncates. +const defaultEmbeddingTokenLimit = tokenizer.EmbeddingTokenLimitDefault -// truncateForEmbedding keeps the first maxTokens tokens of text so it fits the -// embedding model's limit. +// truncateForEmbedding keeps the first `limit` cl100k tokens of text so it fits +// the embedding model's limit. // -// For a positive maxTokens it mirrors Python common/token_utils.py:183-185 -// `truncate(string, max_len)` (keep the first max_len tokens). +// This is the LEGACY path: it is used only when the embedder cannot trim for +// itself (i.e. it does not implement EmbedderTrimmer), such as a test stub. The +// production embedder counts with the model's own tokenizer instead, because +// cl100k alone is not the model's tokenizer and the difference is +// content-dependent. // -// An unconfigured embedder reports maxTokens <= 0. Rather than mirror Python's -// behaviour of returning "" (which would make the embeddings API reject the whole -// batch with "inputs cannot be empty"), Go clamps the limit to a safe default -// (defaultEmbeddingTokenLimit = 8192). This both prevents empty inputs AND keeps -// truncation active for every path (Builtin and generic) instead of silently -// passing the full, untruncated text when no limit is configured. +// Two changes from the original implementation, both fixing the same incident: +// +// - the margin is proportional (tokenizer.EmbeddingTokenLimit) instead of a +// flat 10 tokens, which was 0.12% of an 8192-token window — an order of +// magnitude smaller than the disagreement between two tokenizers; +// - the fallback limit is the resolved window (declared value, then the +// catalog's context_length, then a conservative default) rather than a +// hard-coded 8192, which overshoots the window of every smaller model. func truncateForEmbedding(text string, maxTokens int) string { - if maxTokens <= 0 { - maxTokens = defaultEmbeddingTokenLimit - } - // Keep a 10-token safety margin, mirroring Python's embedding path - // (rag/svr/task_executor.py uses `mdl.max_length - 10`). Only apply it - // when the limit is large enough; for small limits (<=10) keep the full - // value so the result stays non-empty instead of collapsing to "". - if maxTokens > 10 { - maxTokens -= 10 - } - return tokenizer.TrimContentToTokenLimit(text, maxTokens) + declared := tokenizer.ResolveEmbeddingMaxTokens(maxTokens, 0) + limit := tokenizer.EmbeddingTokenLimit(declared) + return tokenizer.TrimContentToTokenLimit(text, limit) } func mergeEmbeddingVectors(titleVec, contentVec []float64, titleWeight float64) ([]float64, error) { diff --git a/internal/ingestion/service/doc_run_convergence_test.go b/internal/ingestion/service/doc_run_convergence_test.go new file mode 100644 index 0000000000..001cc38b95 --- /dev/null +++ b/internal/ingestion/service/doc_run_convergence_test.go @@ -0,0 +1,85 @@ +// +// 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 service + +import ( + "testing" + + "ragflow/internal/common" + "ragflow/internal/entity" +) + +// TestDocRunFromTasks_Empty covers the "no tasks" case: the doc is UNSTART. +func TestDocRunFromTasks_Empty(t *testing.T) { + if got := docRunFromTasks(nil); got != string(entity.TaskStatusUnstart) { + t.Fatalf("empty task set run = %q, want %q", got, entity.TaskStatusUnstart) + } +} + +// TestDocRunFromTasks_LatestWins covers the single-task mapping for every +// status. tasks are newest-first; the first element is authoritative. +func TestDocRunFromTasks_LatestWins(t *testing.T) { + cases := []struct { + name string + status string + want entity.TaskStatus + }{ + {"created", common.CREATED, entity.TaskStatusRunning}, + {"running", common.RUNNING, entity.TaskStatusRunning}, + {"stopping", common.STOPPING, entity.TaskStatusRunning}, + {"failed", common.FAILED, entity.TaskStatusFail}, + {"stopped", common.STOPPED, entity.TaskStatusCancel}, + {"completed", common.COMPLETED, entity.TaskStatusDone}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + tasks := []*entity.IngestionTask{{ID: "task-1", DocumentID: "doc-1", Status: tc.status}} + if got := docRunFromTasks(tasks); got != string(tc.want) { + t.Fatalf("run = %q, want %q", got, tc.want) + } + }) + } +} + +// TestDocRunFromTasks_HistoricalIgnored covers the multi-task case: only the +// newest task (first element) determines the run label; older terminal tasks +// from previous parse rounds must not affect it. +func TestDocRunFromTasks_HistoricalIgnored(t *testing.T) { + cases := []struct { + name string + latest string + history []string + want entity.TaskStatus + }{ + {"newest-failed-over-old-completed", common.FAILED, []string{common.COMPLETED}, entity.TaskStatusFail}, + {"newest-completed-over-old-failed", common.COMPLETED, []string{common.FAILED, common.STOPPED}, entity.TaskStatusDone}, + {"newest-running-over-old-failed", common.RUNNING, []string{common.FAILED}, entity.TaskStatusRunning}, + {"newest-stopped-over-old-done", common.STOPPED, []string{common.COMPLETED}, entity.TaskStatusCancel}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + // Build newest-first: latest task first, then older history tasks. + tasks := []*entity.IngestionTask{{ID: "task-latest", DocumentID: "doc-1", Status: tc.latest}} + for i, st := range tc.history { + tasks = append(tasks, &entity.IngestionTask{ID: string(rune('a' + i)), DocumentID: "doc-1", Status: st}) + } + if got := docRunFromTasks(tasks); got != string(tc.want) { + t.Fatalf("run = %q, want %q (latest=%q)", got, tc.want, tc.latest) + } + }) + } +} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index 13938256df..a5b0675236 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -958,9 +958,30 @@ func (e *Ingestor) settleMessage(ctx context.Context, taskCtx *taskpkg.TaskConte // (handleAndExecute Ack-skips an already-FAILED task); Nack for // redelivery. The broker's redelivery limit handles deterministic // poison messages. - common.Error(fmt.Sprintf("task %s panicked: %v", taskCtx.IngestionTask.ID, r), fmt.Errorf("%v", r)) - e.markFailed(ctx, taskCtx.IngestionTask.ID) - e.recordTerminalPipelineLog(ctx, taskCtx.IngestionTask, string(entity.TaskStatusFail), fmt.Sprintf("Task panicked: %v", r)) + // + // The bookkeeping below reaches the DB, so it is guarded in turn: a + // panic raised inside a deferred function is unrecoverable and would + // take the worker process down — the exact outcome this handler + // exists to prevent. taskCtx is read defensively as well, since a + // panic is proof that an assumption upstream already broke. + var task *entity.IngestionTask + if taskCtx != nil { + task = taskCtx.IngestionTask + } + taskID := "" + if task != nil { + taskID = task.ID + } + func() { + defer func() { + if r2 := recover(); r2 != nil { + common.Error(fmt.Sprintf("task %s: failure handler panicked: %v", taskID, r2), fmt.Errorf("%v", r2)) + } + }() + common.Error(fmt.Sprintf("task %s panicked: %v", taskID, r), fmt.Errorf("%v", r)) + e.markFailed(ctx, taskID) + e.recordTerminalPipelineLog(ctx, task, string(entity.TaskStatusFail), fmt.Sprintf("Task panicked: %v", r)) + }() terminal = false } if e.leaseAbandoned(hb) { @@ -1079,6 +1100,32 @@ func (e *Ingestor) pollCancel(taskID string, cancel context.CancelFunc, done <-c } } +// docRunFromTasks maps the doc's latest ingestion task to the document-level +// run label. tasks is a single-element slice holding the newest task (as +// returned by IngestionTaskDAO.GetByDocumentID, ordered by create_time DESC); +// tasks[0] is the current parse round and its status is authoritative. A +// document can be parsed multiple times over its lifetime, but document.run +// reflects the latest parse round, so historical tasks are ignored. An empty +// task set means UNSTART. +func docRunFromTasks(tasks []*entity.IngestionTask) string { + if len(tasks) == 0 { + return string(entity.TaskStatusUnstart) + } + latest := tasks[0] + switch latest.Status { + case common.CREATED, common.RUNNING, common.STOPPING: + return string(entity.TaskStatusRunning) + case common.FAILED: + return string(entity.TaskStatusFail) + case common.STOPPED: + return string(entity.TaskStatusCancel) + case common.COMPLETED: + return string(entity.TaskStatusDone) + default: + return string(entity.TaskStatusUnstart) + } +} + // claimTask registers a worker claim on a task ID. Returns false if another // worker has already claimed it (e.g. MQ redelivery), true on first claim. // The claim is released by releaseTask when the worker finishes, so a future diff --git a/internal/ingestion/service/nil_database_guard_test.go b/internal/ingestion/service/nil_database_guard_test.go new file mode 100644 index 0000000000..50d6104e0f --- /dev/null +++ b/internal/ingestion/service/nil_database_guard_test.go @@ -0,0 +1,54 @@ +// +// 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 service + +import ( + "context" + "testing" + + "ragflow/internal/entity" + "ragflow/internal/ingestion/testutil" +) + +// dao.DB is a process-wide handle that test helpers swap in and restore, while a +// worker goroutine is free to still be processing a task. Both the worker's +// terminal bookkeeping and settleMessage's panic handler read it, so a nil +// handle has to degrade to a logged error: a panic raised inside a deferred +// function is unrecoverable, so it escapes the handler that exists to keep the +// worker alive and kills the process with it. Regression: CI run 35325530403, +// where that turned a torn-down test database into a crashed test binary. +func TestMarkFailedToleratesNilDatabase(t *testing.T) { + defer testutil.ReplaceDBForTest(t, nil)() + + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) + if terminal := ingestor.markFailed(context.Background(), "task-1"); terminal { + t.Fatal("markFailed reported a durable FAILED write without a database") + } +} + +// TestRunTaskToleratesNilDatabase covers the other half of the same path: the +// terminal pipeline-log/reset bookkeeping that runs after the pipeline itself. +func TestRunTaskToleratesNilDatabase(t *testing.T) { + defer testutil.ReplaceDBForTest(t, nil)() + + runID := "run-1" + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(context.Context, *entity.IngestionTask) error { return nil } + ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", PipelineLogID: &runID, + }) +} diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go index 2bde20cb4c..1f2f224031 100644 --- a/internal/ingestion/task/embedder.go +++ b/internal/ingestion/task/embedder.go @@ -18,17 +18,42 @@ package task import ( "context" + "crypto/sha256" + "errors" "fmt" + "math" + "math/rand" + "net" + "net/http" + "regexp" + "strconv" "strings" + "sync" + "syscall" + "time" + "go.uber.org/zap" + + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity/models" componentpkg "ragflow/internal/ingestion/component" "ragflow/internal/service" + "ragflow/internal/tokenizer" ) type embedder struct { model *models.EmbeddingModel + + // limiter is resolved once per embedder (one embedder is built per tokenizer + // invocation, i.e. per document) and binds the model's tokenizer, the + // declared input window and the calibration for this provider instance. + limiterOnce sync.Once + limiterVal tokenizer.Limiter + // limiterErr records the model declaring a tokenizer whose asset is not on disk. + // It is not a warning: the counter that would be used instead is a calibrated + // estimate of a *different* tokenizer, so the embedder refuses to run (limiter()). + limiterErr error } func (e *embedder) MaxTokens() int { @@ -38,6 +63,16 @@ func (e *embedder) MaxTokens() int { return e.model.MaxTokens } +// ResolveMaxTokens is the input window to actually honour: the model's declared +// value, then the provider catalog's context_length, then a conservative +// default. MaxTokens() stays as the raw declaration for compatibility. +func (e *embedder) ResolveMaxTokens() int { + if e == nil || e.model == nil { + return tokenizer.EmbeddingTokenLimitDefault + } + return tokenizer.ResolveEmbeddingMaxTokens(e.model.ResolveMaxTokens(), 0) +} + func (e *embedder) BatchSize() int { if e == nil || e.model == nil { return models.DefaultEmbeddingBatchSize @@ -45,20 +80,617 @@ func (e *embedder) BatchSize() int { return e.model.ResolveBatchSize() } +// limiter returns the counter+margin+calibration for this embedding model. +// +// A model that declares a tokenizer family whose asset is not on disk is an ERROR, not +// a warning: the calibrated cl100k fallback counts a different tokenizer (cl100k +// under-counts XLM-R on some content, and an under-count is what makes a provider answer +// 400), so carrying on would trade away the exactness these counters exist for. Trim() +// cannot report it - the EmbedderTrimmer seam returns no error - so it surfaces on the +// first Encode(), which every ingest reaches. +func (e *embedder) limiter() tokenizer.Limiter { + e.limiterOnce.Do(func() { + id := "" + if e.model != nil { + id = e.model.ResolveTokenizerID() + } + if id != "" && !tokenizer.CounterExact(id) { + e.limiterErr = fmt.Errorf( + "embedding tokenizer %q is declared for %s but its asset is unavailable (check that ragflow_deps/huggingface.co is present; run `uv run ragflow_deps/download_go_deps.py`): refusing to count with the calibrated estimate", + id, e.quotaKey()) + } + e.limiterVal = tokenizer.LimiterFor(id, string(e.quotaKey()), tokenizer.DefaultCalibration()) + }) + return e.limiterVal +} + +// Trim cuts text down to what this model accepts. This implements the +// component's optional EmbedderTrimmer seam, which is what makes the ingest path +// count with the model's own tokenizer (or a calibrated bound of it) instead of +// cl100k alone. +func (e *embedder) Trim(text string) (string, int) { + limiter := e.limiter() + return limiter.Trim(text, e.ResolveMaxTokens()) +} + +// Encode embeds texts, keeping three properties the previous implementation did +// not have: +// +// 1. It never lets one oversized input fail the whole document: an over-limit +// rejection shrinks the budget and retries, then isolates the offending +// input, instead of returning the error to the caller; +// 2. It reports the provider's own token usage, which doubles as the token +// count for accounting and as the oracle that calibrates our counting; +// 3. It records an over-limit rejection against the calibration, so the next +// attempt - and the next document - starts from a tighter bound. func (e *embedder) Encode(ctx context.Context, texts []string) ([]componentpkg.EmbeddingResult, error) { if e.model.ModelDriver == nil { return nil, fmt.Errorf("embedder: embedding model driver is nil for model %v", e.model.ModelName) } + if len(texts) == 0 { + return []componentpkg.EmbeddingResult{}, nil + } + limiter := e.limiter() + if e.limiterErr != nil { + return nil, e.limiterErr + } + maxTokens := e.ResolveMaxTokens() + cal, calKey := limiter.Calibration() + + var lastErr error + for _, limit := range overLimitLadder(limiter.Limit(maxTokens)) { + candidate := trimAll(texts, limiter, maxTokens, limit) + embeds, usage, err := e.embedWithRetry(ctx, candidate) + if err == nil { + e.recordUsage(cal, calKey, limiter, candidate, usage) + return distributeTokenCount(embeds, usage, candidate, limiter), nil + } + if !isOverLimitErr(err) { + return nil, err + } + // An over-limit rejection is the only signal that reveals the true + // ratio for a model we cannot count exactly; use it before retrying. The + // window applies per input, so the observation must be the LARGEST + // input's own count: a batch total is expected to exceed the window and + // would imply a ratio below 1, i.e. teach us nothing. + if cal != nil { + cal.ObserveOverLimit(calKey, ownTokenMax(candidate, limiter), maxTokens) + } + common.Warn("embedding input over the model limit; shrinking and retrying", + zap.String("model", derefString(e.model.ModelName)), + zap.Int("limit", limit), + zap.String("counter", limiter.Counter().ID()), + zap.Float64("ratio", limiter.Ratio())) + lastErr = err + } + + // The shrink ladder was not enough, which means the batch probably contains + // one pathological input next to healthy ones. Isolate instead of failing. + vecs, err := e.encodeIsolating(ctx, texts, limiter, maxTokens) + if err == nil { + return vecs, nil + } + if lastErr != nil { + return nil, lastErr + } + return nil, err +} + +// overLimitLadder is the sequence of budgets tried, as fractions of the first +// one, oldest first. Trimming is a prefix operation, so each step simply sends a +// shorter prefix than the previous one. +// +// 1 -> 0.75 -> 0.5 -> 0.25 -> 0.125: the first two steps are gentle, because an +// over-limit rejection usually means the input is a little past the window (a +// tokenizer that under-counts by a percent or two), and cutting a quarter of the +// content for that would lose text for no reason. The later steps are for the +// genuinely oversized input (a table the chunker never split), where only a large +// cut helps. +// +// The ladder deliberately stops above the floor: reaching it here would trim EVERY +// input of the batch to 64 tokens because one input is pathological, while the +// per-input isolation pass (overLimitLadderToFloor, which Encode falls through to) +// lets each healthy input keep its own budget. A provider that rejects even the last +// proportional step therefore still gets one attempt at the floor - per input, on the +// input that actually needs it. +func overLimitLadder(budget int) []int { + if budget <= 0 { + return []int{0} + } + limits := []int{budget} + for _, factor := range []float64{0.75, 0.5, 0.25, 0.125} { + next := int(float64(budget) * factor) + if next < overLimitFloorTokens { + next = overLimitFloorTokens + } + if next < limits[len(limits)-1] { + limits = append(limits, next) + } + } + return limits +} + +// overLimitLadderToFloor is overLimitLadder plus the last-resort floor, for the +// per-input isolation path: an input that no proportional step fits may still fit +// at the floor, and trying it there is what turns "this document cannot be +// embedded" into an embedded (truncated) chunk. +// +// The batch loop deliberately stops above the floor. Reaching it there would trim +// EVERY input of the batch to 64 tokens just because one input is pathological, +// while isolating lets each healthy input keep its own budget. +func overLimitLadderToFloor(budget int) []int { + limits := overLimitLadder(budget) + if limits[len(limits)-1] > overLimitFloorTokens { + limits = append(limits, overLimitFloorTokens) + } + return limits +} + +// overLimitFloorTokens is the smallest budget worth trying before declaring the +// input genuinely broken rather than merely long. +const overLimitFloorTokens = 64 + +// trimAll re-trims every text to `limit` tokens. Texts handed in by the caller +// are already trimmed to the full budget, and trimming a prefix again just makes +// it shorter, so this needs no access to the untrimmed input. +func trimAll(texts []string, limiter tokenizer.Limiter, maxTokens, limit int) []string { + out := make([]string, len(texts)) + for i, t := range texts { + out[i] = limiter.Counter().TrimToLimit(t, limit) + } + return out +} + +// ownTokenTotal estimates what the batch costs in our own counter, used for the +// usage-based calibration. +func ownTokenTotal(texts []string, limiter tokenizer.Limiter) int { + total := 0 + for _, t := range texts { + total += limiter.Counter().Count(t) + } + return total +} + +// ownTokenMax is the largest single input's cost in our own counter. The +// provider's window bounds each input individually, so this - not the batch +// total - is what an over-limit rejection tells us about. +func ownTokenMax(texts []string, limiter tokenizer.Limiter) int { + max := 0 + for _, t := range texts { + if n := limiter.Counter().Count(t); n > max { + max = n + } + } + return max +} + +// embedWithRetry performs one embedding call, retrying rate limits exactly as +// before (a TPM-limited provider refills on a one-minute window, and the shared +// cooldown damps all workers at once). It returns the provider's token usage for +// the request. +func (e *embedder) embedWithRetry(ctx context.Context, texts []string) ([]models.EmbeddingData, int, error) { config := &models.EmbeddingConfig{Dimension: 0} - embeds, err := e.model.ModelDriver.Embed(ctx, e.model.ModelName, models.EmbedRequest{Texts: texts}, e.model.APIConfig, config, nil) - if err != nil { - return nil, err + req := models.EmbedRequest{Texts: texts} + // The cooldown is scoped to this provider instance: workers embedding + // through the same credentials+endpoint back off together, everything else + // keeps running. + quota := e.quotaKey() + + var ( + embeds []models.EmbeddingData + err error + ) + for attempt := 1; attempt <= maxEncodeAttempts; attempt++ { + // Sit out any cooldown that another worker's 429 started. The provider + // quota is shared, so a worker that keeps firing while the window is + // exhausted only earns more 429s - and the whole point of the cooldown + // is that one worker's rejection damps all of them. + if cerr := waitOutCooldown(ctx, quota); cerr != nil { + return nil, 0, cerr + } + usage := &common.ModelUsage{} + embeds, err = e.model.ModelDriver.Embed(ctx, e.model.ModelName, req, e.model.APIConfig, config, usage) + if err == nil { + return embeds, usage.InputTokens, nil + } + if !isRetryableErr(err) { + return nil, 0, err + } + if attempt == maxEncodeAttempts { + return nil, 0, err + } + wait := encodeBackoff(attempt, err) + startCooldown(quota, wait) + select { + case <-ctx.Done(): + return nil, 0, ctx.Err() + case <-time.After(wait): + } } + return nil, 0, err +} + +// distributeTokenCount attaches a token count to every input. Providers report +// usage per request, so the total is spread over the inputs in proportion to our +// own counts of exactly the texts we sent; the rounding remainder goes to the +// largest input so the sum matches the provider's number. +func distributeTokenCount(embeds []models.EmbeddingData, usageTotal int, texts []string, limiter tokenizer.Limiter) []componentpkg.EmbeddingResult { vecs := make([]componentpkg.EmbeddingResult, len(embeds)) - for i, v := range embeds { - vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: v.TokenCount} + own := make([]int, len(texts)) + total := 0 + for i, t := range texts { + own[i] = limiter.Counter().Count(t) + total += own[i] } - return vecs, nil + if usageTotal <= 0 || total <= 0 { + // No usage reported by the provider: fall back to our own count rather + // than to 0, which is what the previous implementation always reported. + for i, v := range embeds { + tc := 0 + if i < len(own) { + tc = own[i] + } + vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: tc} + } + return vecs + } + largest, assigned := 0, 0 + for i := range own { + if own[i] > own[largest] { + largest = i + } + } + for i, v := range embeds { + tc := 0 + if i < len(own) { + tc = usageTotal * own[i] / total + assigned += tc + } + vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: tc} + } + if reminder := usageTotal - assigned; reminder > 0 && largest < len(vecs) { + vecs[largest].TokenCount += reminder + } + return vecs +} + +// recordUsage feeds the calibration: the provider's total for this batch against +// our own count of the same texts. +func (e *embedder) recordUsage(cal *tokenizer.Calibration, key string, limiter tokenizer.Limiter, texts []string, usageTotal int) { + if cal == nil || usageTotal <= 0 { + return + } + own := ownTokenTotal(texts, limiter) + if own <= 0 { + return + } + cal.ObserveUsage(key, own, usageTotal) +} + +// encodeIsolating is the last resort: it embeds inputs one at a time so a single +// pathological input cannot drag the rest down, and falls back to the floor +// budget for inputs that still do not fit. It returns an error only when an +// input fails at the floor, i.e. for reasons unrelated to size. +func (e *embedder) encodeIsolating(ctx context.Context, texts []string, limiter tokenizer.Limiter, maxTokens int) ([]componentpkg.EmbeddingResult, error) { + out := make([]componentpkg.EmbeddingResult, len(texts)) + cal, calKey := limiter.Calibration() + var lastErr error + for i, text := range texts { + if err := ctx.Err(); err != nil { + return nil, err + } + embedded := false + for _, limit := range overLimitLadderToFloor(limiter.Limit(maxTokens)) { + candidate := []string{limiter.Counter().TrimToLimit(text, limit)} + embeds, usage, err := e.embedWithRetry(ctx, candidate) + if err == nil { + e.recordUsage(cal, calKey, limiter, candidate, usage) + out[i] = distributeTokenCount(embeds, usage, candidate, limiter)[0] + embedded = true + break + } + if !isOverLimitErr(err) { + return nil, err + } + if cal != nil { + cal.ObserveOverLimit(calKey, limiter.Counter().Count(candidate[0]), maxTokens) + } + lastErr = err + } + if !embedded { + return nil, fmt.Errorf("embedder: input %d does not fit the model window even at %d tokens: %w", i, overLimitFloorTokens, lastErr) + } + } + return out, nil +} + +// overLimitMarkers are the provider wordings that mean "this input is longer +// than the model accepts". 20015 is SiliconFlow's code for exactly that (it +// answers a generic "The parameter is invalid" message); the rest are the usual +// phrasings elsewhere. +var overLimitMarkers = []string{ + "20015", + "too long", + "too many tokens", + "maximum context", + "context length", + "context_length", + "input length", + "token limit", + "reduce the length", + "maximum allowed", +} + +// isOverLimitErr reports whether err is an over-limit rejection rather than a +// rate limit or a genuine failure. Only 4xx rejections qualify: a 5xx is the +// provider's problem and shrinking the input would not help. +func isOverLimitErr(err error) bool { + if err == nil { + return false + } + msg := strings.ToLower(err.Error()) + if !strings.Contains(msg, "400") && !strings.Contains(msg, "413") && !strings.Contains(msg, "422") { + return false + } + for _, marker := range overLimitMarkers { + if strings.Contains(msg, marker) { + return true + } + } + return false +} + +// Retry policy for rate-limited embedding calls. +// +// A TPM-limited provider (SiliconFlow bge-m3 answers "429 ... TPM limit +// reached") refills its quota on a one-minute window, so a retry schedule that +// exhausts itself inside a single window cannot succeed: the 2s/4s/8s/16s of the +// original implementation always ran out while the quota was still spent, and +// every concurrent worker kept hammering the same exhausted quota, turning a +// transient throttle into a FAILED document whose chunks had already been +// deleted. The schedule below spans roughly two and a half windows, and the +// shared cooldown makes the workers back off together instead of independently. +// +// These are variables rather than constants so tests can shrink them. +var ( + maxEncodeAttempts = 6 + encodeBackoffBase = 5 * time.Second + encodeBackoffMax = 120 * time.Second + encodeBackoffJitter = 0.2 +) + +// providerQuotaKey identifies the thing a TPM quota actually belongs to: one +// provider instance (endpoint + credentials + model). See quotaKey. +type providerQuotaKey string + +// rateLimitCooldowns parks callers per provider instance - and only per provider +// instance. +// +// The unit is the provider instance, not the process: a TPM quota is owned by +// the credentials, so a worker that got throttled on provider A must not stall +// ingestion that talks to provider B, and must not stall a different dataset +// embedding through a different model. It cannot be the embedder value either: +// newEmbedderResolver builds a fresh embedder for every tokenizer invocation +// (one per document), so state hung off the value would never be shared by the +// workers that are actually competing for the same quota. A keyed registry +// gives exactly the required scope - all workers embedding through the same +// provider instance back off together, everyone else is untouched. +var rateLimitCooldowns = struct { + sync.Mutex + until map[providerQuotaKey]time.Time +}{until: make(map[providerQuotaKey]time.Time)} + +// startCooldown parks callers sharing `key` for at least d, extending an active +// cooldown but never shortening it. +func startCooldown(key providerQuotaKey, d time.Duration) { + if d <= 0 { + return + } + deadline := time.Now().Add(d) + rateLimitCooldowns.Lock() + defer rateLimitCooldowns.Unlock() + if existing, ok := rateLimitCooldowns.until[key]; !ok || deadline.After(existing) { + rateLimitCooldowns.until[key] = deadline + } + // Keys are bounded by the number of configured provider instances, so this + // is housekeeping rather than a leak fix. + if len(rateLimitCooldowns.until) > 64 { + for k, until := range rateLimitCooldowns.until { + if !until.After(time.Now()) { + delete(rateLimitCooldowns.until, k) + } + } + } +} + +// cooldownRemaining reports how long the cooldown for `key` still has to run. +func cooldownRemaining(key providerQuotaKey) time.Duration { + rateLimitCooldowns.Lock() + defer rateLimitCooldowns.Unlock() + return time.Until(rateLimitCooldowns.until[key]) +} + +// waitOutCooldown blocks until the cooldown for `key` expires, honoring ctx. +func waitOutCooldown(ctx context.Context, key providerQuotaKey) error { + remaining := cooldownRemaining(key) + if remaining <= 0 { + return nil + } + timer := time.NewTimer(remaining) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// quotaKey derives the provider-instance key the cooldown is scoped to. The API +// key decides who owns the quota, so it is part of the key but hashed: the key +// must never carry a plaintext secret. +func (e *embedder) quotaKey() providerQuotaKey { + if e == nil || e.model == nil { + return providerQuotaKey("embedder:unknown") + } + var baseURL, region, apiKey string + if cfg := e.model.APIConfig; cfg != nil { + baseURL = derefString(cfg.BaseURL) + region = derefString(cfg.Region) + apiKey = derefString(cfg.ApiKey) + } + sum := sha256.Sum256([]byte(apiKey)) + return providerQuotaKey(fmt.Sprintf("%s|%s|%s|%x", baseURL, region, derefString(e.model.ModelName), sum[:8])) +} + +// derefString safely dereferences an optional string. +func derefString(p *string) string { + if p == nil { + return "" + } + return *p +} + +// encodeBackoff returns how long to wait before retry `attempt` (1-based) of a +// call that failed with `err`. A provider-supplied Retry-After wins; otherwise +// the delay grows exponentially with jitter so concurrent workers spread out +// instead of retrying in lockstep. +func encodeBackoff(attempt int, err error) time.Duration { + if hint, ok := retryAfterFromError(err); ok { + if hint > encodeBackoffMax { + return encodeBackoffMax + } + return hint + } + delay := float64(encodeBackoffBase) * math.Pow(2, float64(attempt-1)) + if delay > float64(encodeBackoffMax) { + delay = float64(encodeBackoffMax) + } + if encodeBackoffJitter > 0 { + spread := delay * encodeBackoffJitter + delay += (rand.Float64()*2 - 1) * spread + } + if delay < float64(time.Second) { + delay = float64(time.Second) + } + return time.Duration(delay) +} + +// retryAfterPattern matches the seconds form of Retry-After surfaced by the +// model drivers ("retry-after: 30", "retry_after=30", "retry after 30 seconds"). +var retryAfterPattern = regexp.MustCompile(`(?i)retry[-_ ]?after["'\s:=]+(\d+(\.\d+)?)`) + +// retryAfterFromError extracts a Retry-After hint in seconds from an error +// message, if the driver carried one. +func retryAfterFromError(err error) (time.Duration, bool) { + if err == nil { + return 0, false + } + m := retryAfterPattern.FindStringSubmatch(err.Error()) + if m == nil { + return 0, false + } + seconds, perr := strconv.ParseFloat(m[1], 64) + if perr != nil || seconds <= 0 { + return 0, false + } + return time.Duration(seconds * float64(time.Second)), true +} + +// retryableStatus reports whether an HTTP status means "the provider could not +// serve this request now" rather than "this request is bad": 429, 529 (Anthropic's +// overloaded status) and the provider-side 5xx family. +func retryableStatus(status int) bool { + return status == http.StatusTooManyRequests || status == 529 || status >= 500 +} + +// statusBefore/statusAfter pick the HTTP status a driver wrote into its own +// error message ("... error: 429 ...", "... status 503 ...", "500 Internal +// Server Error"). The status WORDING is required so a bare three-digit number in +// a response body (a token count, a document id) cannot be read as a status. +var ( + statusBefore = regexp.MustCompile(`(?i)\b(?:http|status|code|error)\b[^0-9]{0,24}(\d{3})`) + statusAfter = regexp.MustCompile(`(?i)(\d{3})[^0-9]{0,24}\b(?:http|status|code|error)\b`) +) + +// statusFromMessage extracts the status a driver embedded in its message. +func statusFromMessage(msg string) (int, bool) { + for _, re := range []*regexp.Regexp{statusBefore, statusAfter} { + m := re.FindStringSubmatch(msg) + if m == nil { + continue + } + code, err := strconv.Atoi(m[1]) + if err != nil || code < 100 || code >= 600 { + continue + } + return code, true + } + return 0, false +} + +// isRetryableErr reports whether a failed Embed call is worth retrying: the +// provider is rate limiting us, failed on its side (5xx/529), or the request +// never completed (transport failure). +// +// Structured evidence decides first — models.APIStatusError carries the status +// seen by the shared HTTP helpers, and net.Error covers timeouts/DNS/refused +// connections. Only the drivers that build their own HTTP errors have to be read +// from the message (every embedding driver does, unlike the chat path, whose +// helpers return APIStatusError); there an explicit status decides in BOTH +// directions, so a permanent 4xx can no longer buy six retries plus a shared +// cooldown just because its body also mentions rate limits. Provider wording is +// consulted only when no status is present at all. +func isRetryableErr(err error) bool { + if err == nil { + return false + } + var statusErr *models.APIStatusError + if errors.As(err, &statusErr) { + return retryableStatus(statusErr.Status) + } + if errors.Is(err, context.DeadlineExceeded) { + return true + } + var netErr net.Error + if errors.As(err, &netErr) { + return isRetryableNetErr(err, netErr) + } + msg := strings.ToLower(err.Error()) + if code, ok := statusFromMessage(msg); ok { + return retryableStatus(code) + } + return strings.Contains(msg, "too many requests") || + strings.Contains(msg, "rate limit") || + strings.Contains(msg, "rate limiting") || + strings.Contains(msg, "tpm limit") +} + +// isRetryableNetErr separates transient transport failures from permanent ones. A +// net.Error is NOT automatically transient: net.DNSError reports NXDOMAIN ("no such +// host") as one, and a hostname that does not resolve - or a connection the host +// refuses - is a configuration fact no retry changes. Retrying it would spend six +// attempts and a shared cooldown on an answer that cannot get better, and delay the +// error an operator needs to see. +// +// Retried: a timeout, a DNS failure the resolver flagged as temporary, and a +// connection dropped after the request was accepted (reset / broken pipe / aborted). +// Everything else is returned to the caller as-is. +func isRetryableNetErr(err error, netErr net.Error) bool { + if netErr.Timeout() { + return true + } + var dnsErr *net.DNSError + if errors.As(err, &dnsErr) { + return dnsErr.IsTemporary || dnsErr.IsTimeout + } + var opErr *net.OpError + if errors.As(err, &opErr) { + return errors.Is(opErr.Err, syscall.ECONNRESET) || + errors.Is(opErr.Err, syscall.EPIPE) || + errors.Is(opErr.Err, syscall.ECONNABORTED) + } + return false } // newEmbedderResolver builds the production embedder resolver used by the diff --git a/internal/ingestion/task/embedder_limit_test.go b/internal/ingestion/task/embedder_limit_test.go new file mode 100644 index 0000000000..e26b572d37 --- /dev/null +++ b/internal/ingestion/task/embedder_limit_test.go @@ -0,0 +1,294 @@ +// +// 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 task + +import ( + "context" + "errors" + "strings" + "testing" + + "ragflow/internal/common" + "ragflow/internal/entity/models" + "ragflow/internal/tokenizer" +) + +// siliconflowOverLimitErr is the exact rejection observed on document 78785.md: +// a generic "parameter is invalid" message that actually means "this input is +// longer than the model accepts". +const siliconflowOverLimitErr = `SILICONFLOW API error: 400 Bad Request, body: {"code":20015,"message":"The parameter is invalid. Please check again.","data":null}` + +// tokenBoundDriver accepts a request only while every input stays under +// maxTokens tokens (counted with cl100k, which is what the ingest path counted +// with before this change). It mimics a provider whose window is maxTokens. +type tokenBoundDriver struct { + stubDriver + maxTokens int + batchOnly bool // reject only multi-input requests that are too long + rejects int + sawSingle bool + singleCalls int +} + +func (d *tokenBoundDriver) Embed(ctx context.Context, modelName *string, request models.EmbedRequest, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig, usage *common.ModelUsage) ([]models.EmbeddingData, error) { + counter := tokenizer.CountCL100K() + tooLong := false + for _, text := range request.Texts { + if counter.Count(text) > d.maxTokens { + tooLong = true + break + } + } + if tooLong && !d.batchOnly { + d.rejects++ + return nil, errors.New(siliconflowOverLimitErr) + } + if tooLong && d.batchOnly && len(request.Texts) > 1 { + d.rejects++ + return nil, errors.New(siliconflowOverLimitErr) + } + if len(request.Texts) == 1 { + d.sawSingle = true + } + d.singleCalls++ + // Report a request-level usage, exactly like a real provider. + if usage != nil { + usage.InputTokens = 7 * len(request.Texts) + } + return d.stubDriver.Embed(ctx, modelName, request, apiConfig, embeddingConfig, usage) +} + +func newLimitEmbedder(driver models.ModelDriver, declaredMaxTokens int) *embedder { + return &embedder{model: models.NewEmbeddingModel(driver, strPtr("stub"), &models.APIConfig{}, declaredMaxTokens)} +} + +func TestIsOverLimitErrClassification(t *testing.T) { + cases := []struct { + err string + want bool + }{ + {siliconflowOverLimitErr, true}, + {`OpenAI embeddings API error: 400 Bad Request, body: {"error":{"message":"This model's maximum context length is 8192 tokens"}}`, true}, + {`API error: 413 Request Entity Too Large, body: input is too long`, true}, + {`API error: 422 Unprocessable Entity, body: too many tokens`, true}, + // Rate limits and provider failures must not be mistaken for size. + {tpmLimitErr, false}, + {`SILICONFLOW API error: 500 Internal Server Error, body: too long`, false}, + {`OpenAI embeddings API error: 401 Unauthorized, body: invalid api key`, false}, + {`failed to send request: dial tcp: connection refused`, false}, + } + for _, c := range cases { + if got := isOverLimitErr(errors.New(c.err)); got != c.want { + t.Errorf("isOverLimitErr(%q) = %t, want %t", c.err, got, c.want) + } + } + if isOverLimitErr(nil) { + t.Error("isOverLimitErr(nil) = true, want false") + } +} + +func TestOverLimitLadderShrinksAndFloors(t *testing.T) { + limits := overLimitLadder(8192) + want := []int{8192, 6144, 4096, 2048, 1024} + if len(limits) != len(want) { + t.Fatalf("overLimitLadder(8192) = %v, want %v", limits, want) + } + for i := range want { + if limits[i] != want[i] { + t.Fatalf("overLimitLadder(8192) = %v, want %v", limits, want) + } + } + // A small window must not shrink past the floor, and must stay monotone. + small := overLimitLadder(100) + for i, v := range small { + if v < overLimitFloorTokens { + t.Fatalf("ladder step %d = %d, below the floor %d", i, v, overLimitFloorTokens) + } + if i > 0 && v >= small[i-1] { + t.Fatalf("ladder is not strictly shrinking: %v", small) + } + } + if got := overLimitLadder(0); len(got) != 1 || got[0] != 0 { + t.Fatalf("overLimitLadder(0) = %v, want [0]", got) + } + + // The per-input isolation path appends the documented floor: an input that no + // proportional step fits may still fit at 64 tokens, and without that step the + // document is declared unembeddable. + floored := overLimitLadderToFloor(8192) + if len(floored) != len(want)+1 || floored[len(floored)-1] != overLimitFloorTokens { + t.Fatalf("overLimitLadderToFloor(8192) = %v, want %v then the floor %d", floored, want, overLimitFloorTokens) + } + for i := range want { + if floored[i] != want[i] { + t.Fatalf("overLimitLadderToFloor(8192) = %v, want %v then the floor", floored, want) + } + } + if got := overLimitLadderToFloor(overLimitFloorTokens); len(got) != 1 || got[0] != overLimitFloorTokens { + t.Fatalf("overLimitLadderToFloor(%d) = %v, want no duplicate floor step", overLimitFloorTokens, got) + } +} + +func TestDistributeTokenCountSumsToProviderUsage(t *testing.T) { + limiter := tokenizer.NewExactLimiter(tokenizer.CountCL100K()) + texts := []string{strings.Repeat("word ", 10), strings.Repeat("word ", 40), strings.Repeat("word ", 20)} + embeds := make([]models.EmbeddingData, len(texts)) + + got := distributeTokenCount(embeds, 100, texts, limiter) + sum := 0 + largest, largestIdx := -1, -1 + for i, v := range got { + sum += v.TokenCount + if v.TokenCount > largest { + largest, largestIdx = v.TokenCount, i + } + } + if sum != 100 { + t.Fatalf("distributed token counts sum to %d, want the provider total 100", sum) + } + if largestIdx != 1 { + t.Fatalf("the remainder went to input %d, want the longest input (1)", largestIdx) + } +} + +func TestDistributeTokenCountFallsBackToOwnCountWithoutUsage(t *testing.T) { + limiter := tokenizer.NewExactLimiter(tokenizer.CountCL100K()) + texts := []string{"alpha beta", "gamma delta epsilon zeta"} + embeds := make([]models.EmbeddingData, len(texts)) + got := distributeTokenCount(embeds, 0, texts, limiter) + for i := range got { + want := limiter.Counter().Count(texts[i]) + if got[i].TokenCount != want { + t.Errorf("input %d token count = %d, want our own count %d", i, got[i].TokenCount, want) + } + } +} + +// TestEmbedderEncodeShrinksUntilTheProviderAccepts is the 78785.md scenario at +// unit scale: the provider's window is smaller than what our counter allowed, so +// the first attempt is rejected. The document must still be embedded — by +// shrinking — and the rejection must teach the calibration. +func TestEmbedderEncodeShrinksUntilTheProviderAccepts(t *testing.T) { + freezeBackoff(t) + driver := &tokenBoundDriver{maxTokens: 2000} + emb := newLimitEmbedder(driver, 8192) + + long := strings.Repeat("the quick brown fox jumps over the lazy dog ", 500) // ~5500 cl100k tokens + vecs, err := emb.Encode(context.Background(), []string{long, long, long}) + if err != nil { + t.Fatalf("Encode returned %v; an over-limit rejection must not fail the document", err) + } + if len(vecs) != 3 { + t.Fatalf("got %d vectors, want 3", len(vecs)) + } + if driver.rejects == 0 { + t.Fatal("the driver never rejected anything, so the shrink path was not exercised") + } + for i, text := range driver.capturedTexts { + if got := tokenizer.CountCL100K().Count(text); got > driver.maxTokens { + t.Errorf("attempt sent input %d with %d tokens, still above the provider window %d", i, got, driver.maxTokens) + } + } + // The rejection is an observation: the next attempt must start tighter. + if ratio := tokenizer.DefaultCalibration().RatioUpper(string(emb.quotaKey())); ratio <= 1 { + t.Fatalf("calibration ratio after an over-limit rejection = %v, want > 1", ratio) + } + // Every input still gets a real token count (previously always 0). + for i, v := range vecs { + if v.TokenCount <= 0 { + t.Errorf("vector %d has token count %d, want a positive count", i, v.TokenCount) + } + } +} + +// TestEmbedderEncodeIsolatesPathologicalInput covers a batch where the size +// problem only shows up when several inputs travel together: shrinking the +// shared budget cannot help, so the embedder has to isolate the inputs instead of +// failing the whole document. +func TestEmbedderEncodeIsolatesPathologicalInput(t *testing.T) { + freezeBackoff(t) + driver := &tokenBoundDriver{maxTokens: 200, batchOnly: true} + emb := newLimitEmbedder(driver, 8192) + + long := strings.Repeat("lorem ipsum dolor sit amet ", 200) + vecs, err := emb.Encode(context.Background(), []string{"short one", long, "short two"}) + if err != nil { + t.Fatalf("Encode returned %v; isolation should have saved the document", err) + } + if len(vecs) != 3 { + t.Fatalf("got %d vectors, want 3", len(vecs)) + } + if !driver.sawSingle { + t.Fatal("the embedder never fell back to per-input calls") + } +} + +func TestEmbedderEncodeEmptyInput(t *testing.T) { + emb := newLimitEmbedder(&stubDriver{}, 8192) + vecs, err := emb.Encode(context.Background(), nil) + if err != nil { + t.Fatalf("Encode(nil) returned %v", err) + } + if len(vecs) != 0 { + t.Fatalf("Encode(nil) returned %d vectors, want 0", len(vecs)) + } +} + +// TestEmbedderTrimPrefersModelTokenizer pins the contract the component relies +// on: the embedder's Trim honours the declared window with a proportional margin +// (not a flat 10 tokens), and never hands back more than the limit allows. +func TestEmbedderTrimPrefersModelTokenizer(t *testing.T) { + emb := newLimitEmbedder(&stubDriver{}, 8192) + if got := emb.ResolveMaxTokens(); got != 8192 { + t.Fatalf("ResolveMaxTokens() = %d, want the declared 8192", got) + } + text := strings.Repeat("token ", 20000) + trimmed, tokens := emb.Trim(text) + limit := tokenizer.EmbeddingTokenLimit(8192) + if tokens > limit { + t.Fatalf("Trim returned %d tokens, limit is %d", tokens, limit) + } + if got := tokenizer.CountCL100K().Count(trimmed); got > limit { + t.Fatalf("trimmed text counts %d tokens, limit is %d", got, limit) + } + // Defaults are also honoured when a model declares nothing. + unset := newLimitEmbedder(&stubDriver{}, 0) + if got := unset.ResolveMaxTokens(); got != tokenizer.EmbeddingTokenLimitDefault { + t.Fatalf("ResolveMaxTokens() with no declaration = %d, want %d", got, tokenizer.EmbeddingTokenLimitDefault) + } +} + +// TestEmbedderEncodeRefusesToSubstituteTheCalibratedEstimate pins that a model declaring a +// tokenizer it cannot load FAILS instead of counting with the calibrated cl100k estimate: +// that estimate belongs to a different tokenizer (cl100k under-counts XLM-R on some +// content), and an under-count is the direction that makes a provider answer 400. The id +// below is deliberately one nothing registers - the same path as an asset that never +// loaded - and it reaches the resolution through the same env override an operator uses to +// pin a family. +func TestEmbedderEncodeRefusesToSubstituteTheCalibratedEstimate(t *testing.T) { + t.Setenv("TOKENIZER_EMBEDDING_TOKENIZER", "no-such-family") + emb := newLimitEmbedder(&stubDriver{}, 8192) + _, err := emb.Encode(context.Background(), []string{"hello"}) + if err == nil { + t.Fatal("Encode succeeded with an unavailable declared tokenizer; it must refuse instead") + } + for _, want := range []string{`"no-such-family"`, "but its asset is unavailable", "refusing to count with the calibrated estimate"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error %q does not mention %q", err.Error(), want) + } + } +} diff --git a/internal/ingestion/task/embedder_ratelimit_test.go b/internal/ingestion/task/embedder_ratelimit_test.go new file mode 100644 index 0000000000..116bc53a70 --- /dev/null +++ b/internal/ingestion/task/embedder_ratelimit_test.go @@ -0,0 +1,333 @@ +// +// 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 task + +import ( + "context" + "errors" + "fmt" + "net" + "strings" + "syscall" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/entity/models" +) + +const tpmLimitErr = `SILICONFLOW API error: 429 Too Many Requests, body: {"message":"Request was rejected due to rate limiting. Details: TPM limit reached.","data":null}` + +// flakyDriver fails its first `failures` Embed calls with `err`. +type flakyDriver struct { + stubDriver + failures int + err error + calls int +} + +func (d *flakyDriver) Embed(ctx context.Context, modelName *string, request models.EmbedRequest, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig, usage *common.ModelUsage) ([]models.EmbeddingData, error) { + d.calls++ + if d.calls <= d.failures { + return nil, d.err + } + return d.stubDriver.Embed(ctx, modelName, request, apiConfig, embeddingConfig, usage) +} + +func newFlakyEmbedder(driver *flakyDriver) *embedder { + return newFlakyEmbedderWithConfig(driver, &models.APIConfig{}) +} + +func newFlakyEmbedderWithConfig(driver *flakyDriver, cfg *models.APIConfig) *embedder { + return &embedder{model: models.NewEmbeddingModel(driver, strPtr("embed"), cfg, 128)} +} + +// freezeBackoff shrinks the retry schedule to milliseconds and clears any +// cooldown left behind by an earlier test, restoring both on cleanup. +func freezeBackoff(t *testing.T) { + t.Helper() + prevBase, prevMax, prevJitter, prevAttempts := encodeBackoffBase, encodeBackoffMax, encodeBackoffJitter, maxEncodeAttempts + encodeBackoffBase = time.Millisecond + encodeBackoffMax = 5 * time.Millisecond + encodeBackoffJitter = 0 + resetCooldown() + t.Cleanup(func() { + encodeBackoffBase, encodeBackoffMax, encodeBackoffJitter, maxEncodeAttempts = prevBase, prevMax, prevJitter, prevAttempts + resetCooldown() + }) +} + +func resetCooldown() { + rateLimitCooldowns.Lock() + rateLimitCooldowns.until = make(map[providerQuotaKey]time.Time) + rateLimitCooldowns.Unlock() +} + +func TestEmbedderEncode_RetriesRateLimitThenSucceeds(t *testing.T) { + freezeBackoff(t) + driver := &flakyDriver{failures: 2, err: errors.New(tpmLimitErr)} + emb := newFlakyEmbedder(driver) + + vecs, err := emb.Encode(t.Context(), []string{"a", "b"}) + if err != nil { + t.Fatalf("Encode: %v", err) + } + if driver.calls != 3 { + t.Fatalf("driver calls = %d, want 3 (two 429s then success)", driver.calls) + } + if len(vecs) != 2 { + t.Fatalf("vectors = %d, want 2", len(vecs)) + } +} + +func TestEmbedderEncode_GivesUpAfterMaxAttempts(t *testing.T) { + freezeBackoff(t) + maxEncodeAttempts = 3 + driver := &flakyDriver{failures: 99, err: errors.New(tpmLimitErr)} + emb := newFlakyEmbedder(driver) + + if _, err := emb.Encode(t.Context(), []string{"a"}); err == nil { + t.Fatal("expected the rate-limit error to surface once attempts are exhausted") + } + if driver.calls != 3 { + t.Fatalf("driver calls = %d, want 3", driver.calls) + } +} + +func TestEmbedderEncode_NonRateLimitErrorFailsFast(t *testing.T) { + freezeBackoff(t) + driver := &flakyDriver{failures: 99, err: errors.New("SILICONFLOW API error: 400 Bad Request, body: bad input")} + emb := newFlakyEmbedder(driver) + + if _, err := emb.Encode(t.Context(), []string{"a"}); err == nil { + t.Fatal("expected an error") + } + if driver.calls != 1 { + t.Fatalf("driver calls = %d, want 1: only rate limits are retried", driver.calls) + } +} + +// A 429 parks every worker on the same provider instance, not just the caller +// that saw it: the quota they compete for is the same one. +func TestEmbedderEncode_WaitsOutCooldownStartedByAnotherWorker(t *testing.T) { + freezeBackoff(t) + driver := &flakyDriver{} + emb := newFlakyEmbedder(driver) + startCooldown(emb.quotaKey(), 120*time.Millisecond) + defer resetCooldown() + + start := time.Now() + if _, err := emb.Encode(t.Context(), []string{"a"}); err != nil { + t.Fatalf("Encode: %v", err) + } + if elapsed := time.Since(start); elapsed < 100*time.Millisecond { + t.Fatalf("Encode returned after %v; expected it to wait out the active cooldown", elapsed) + } + if driver.calls != 1 { + t.Fatalf("driver calls = %d, want 1", driver.calls) + } +} + +func TestEmbedderEncode_CooldownHonorsContextCancellation(t *testing.T) { + freezeBackoff(t) + driver := &flakyDriver{} + emb := newFlakyEmbedder(driver) + startCooldown(emb.quotaKey(), 5*time.Second) + defer resetCooldown() + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + if _, err := emb.Encode(ctx, []string{"a"}); err == nil { + t.Fatal("expected the cancelled context to surface") + } + if driver.calls != 0 { + t.Fatalf("driver calls = %d, want 0: the call must not be made during a cooldown", driver.calls) + } +} + +// The cooldown is scoped to one provider instance. A throttled credential (or +// endpoint) must not stall ingestion that embeds somewhere else - otherwise a +// single throttled provider would freeze the whole process. +func TestEmbedderEncode_CooldownDoesNotBlockOtherProviderInstances(t *testing.T) { + freezeBackoff(t) + throttled := newFlakyEmbedderWithConfig( + &flakyDriver{}, + &models.APIConfig{ApiKey: strPtr("key-a"), BaseURL: strPtr("https://a.example")}, + ) + startCooldown(throttled.quotaKey(), 5*time.Second) + defer resetCooldown() + + other := newFlakyEmbedderWithConfig( + &flakyDriver{}, + &models.APIConfig{ApiKey: strPtr("key-b"), BaseURL: strPtr("https://b.example")}, + ) + start := time.Now() + if _, err := other.Encode(t.Context(), []string{"a"}); err != nil { + t.Fatalf("Encode on the other provider instance: %v", err) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("a cooldown on provider A delayed provider B by %v", elapsed) + } +} + +// Conversely, two workers embedding through the same provider instance must +// land on one quota key: the resolver builds a fresh embedder per tokenizer +// invocation, so per-value state would never be shared between them. +func TestEmbedderQuotaKey_ScopesByProviderInstance(t *testing.T) { + cfg := &models.APIConfig{ApiKey: strPtr("key-a"), BaseURL: strPtr("https://a.example")} + first := newFlakyEmbedderWithConfig(&flakyDriver{}, cfg) + second := newFlakyEmbedderWithConfig(&flakyDriver{}, cfg) + if first.quotaKey() != second.quotaKey() { + t.Fatal("two embedders on the same provider instance must share one quota key") + } + + otherKey := newFlakyEmbedderWithConfig( + &flakyDriver{}, + &models.APIConfig{ApiKey: strPtr("key-b"), BaseURL: strPtr("https://a.example")}, + ) + if first.quotaKey() == otherKey.quotaKey() { + t.Fatal("a different credential owns a different quota and must not share the key") + } + + otherEndpoint := newFlakyEmbedderWithConfig( + &flakyDriver{}, + &models.APIConfig{ApiKey: strPtr("key-a"), BaseURL: strPtr("https://b.example")}, + ) + if first.quotaKey() == otherEndpoint.quotaKey() { + t.Fatal("a different endpoint must not share the key") + } + + // The key must not carry the secret itself. + if key := string(first.quotaKey()); strings.Contains(key, "key-a") { + t.Fatalf("quota key leaks the API key: %q", key) + } +} + +func TestEncodeBackoff_GrowsAndIsCapped(t *testing.T) { + prevBase, prevMax, prevJitter := encodeBackoffBase, encodeBackoffMax, encodeBackoffJitter + defer func() { + encodeBackoffBase, encodeBackoffMax, encodeBackoffJitter = prevBase, prevMax, prevJitter + }() + encodeBackoffBase = 5 * time.Second + encodeBackoffMax = 120 * time.Second + encodeBackoffJitter = 0 + + want := []time.Duration{5 * time.Second, 10 * time.Second, 20 * time.Second, 40 * time.Second, 80 * time.Second, 120 * time.Second} + for i, w := range want { + if got := encodeBackoff(i+1, errors.New(tpmLimitErr)); got != w { + t.Errorf("attempt %d backoff = %v, want %v", i+1, got, w) + } + } + // The whole schedule has to outlast a one-minute TPM window, otherwise every + // attempt lands inside the same exhausted quota. + total := time.Duration(0) + for i := range want { + total += encodeBackoff(i+1, errors.New(tpmLimitErr)) + } + if total < 2*time.Minute { + t.Errorf("total backoff = %v, want >= 2m to span more than one TPM window", total) + } +} + +func TestEncodeBackoff_HonorsRetryAfter(t *testing.T) { + prevMax := encodeBackoffMax + defer func() { encodeBackoffMax = prevMax }() + encodeBackoffMax = 120 * time.Second + + cases := []struct { + msg string + want time.Duration + }{ + {"SILICONFLOW API error: 429 Too Many Requests, retry-after: 30, body: {}", 30 * time.Second}, + {"429 retry_after=7", 7 * time.Second}, + {"429 too many requests, retry after 12 seconds", 12 * time.Second}, + {`429 {"retry-after": 45}`, 45 * time.Second}, + } + for _, tc := range cases { + if got := encodeBackoff(1, errors.New(tc.msg)); got != tc.want { + t.Errorf("encodeBackoff(%q) = %v, want %v", tc.msg, got, tc.want) + } + } + // A provider hint larger than the cap is clamped. + if got := encodeBackoff(1, errors.New("429 retry-after: 3600")); got != encodeBackoffMax { + t.Errorf("huge Retry-After = %v, want the %v cap", got, encodeBackoffMax) + } +} + +func TestIsRetryableErr(t *testing.T) { + retryable := []string{ + tpmLimitErr, + "SILICONFLOW API error: 429 Too Many Requests, body: {}", + "provider says: too many requests", + "rate limit exceeded", + "TPM limit reached", + // Provider-side failures: a 5xx and Anthropic's overloaded 529 mean "not + // now", so they are retried like a 429. + "API request failed with status 503: unavailable", + "error: 529 overloaded", + "500 Internal Server Error", + } + for _, msg := range retryable { + if !isRetryableErr(errors.New(msg)) { + t.Errorf("isRetryableErr(%q) = false, want true", msg) + } + } + notRetryable := []string{ + "", + "402 Payment Required", + // An explicit status decides in both directions: a permanent 4xx forbids + // the retry even when the body also mentions rate limits. + "SILICONFLOW API error: 400 Bad Request, body: rate limit header malformed", + // A bare three-digit number in a body is data, not a status. + `{"usage":{"total_tokens":429}}`, + } + for _, msg := range notRetryable { + if isRetryableErr(errors.New(msg)) { + t.Errorf("isRetryableErr(%q) = true, want false", msg) + } + } + // Structured status from the shared HTTP helpers, no message parsing at all. + if !isRetryableErr(&models.APIStatusError{Status: 429, Body: "{}"}) { + t.Error("a typed 429 must be retryable") + } + if isRetryableErr(&models.APIStatusError{Status: 400, Body: "bad input"}) { + t.Error("a typed 400 must not be retryable") + } + if isRetryableErr(nil) { + t.Error("isRetryableErr(nil) = true, want false") + } + // A net.Error is not automatically transient. An unresolvable host or a refused + // connection is configuration the operator has to fix; retrying it spends six + // attempts and the shared cooldown on an answer that cannot get better. + netCases := []struct { + name string + err error + want bool + }{ + {"dns NXDOMAIN", &net.DNSError{Err: "no such host", Name: "provider.invalid", IsNotFound: true}, false}, + {"dns temporary", &net.DNSError{Err: "server misbehaving", Name: "provider.example", IsTemporary: true}, true}, + {"timeout", &net.DNSError{Err: "i/o timeout", Name: "provider.example", IsTimeout: true}, true}, + {"connection refused", &net.OpError{Op: "dial", Err: syscall.ECONNREFUSED}, false}, + {"connection reset after accept", &net.OpError{Op: "read", Err: syscall.ECONNRESET}, true}, + {"wrapped dns NXDOMAIN", fmt.Errorf("embed: %w", &net.DNSError{Err: "no such host", IsNotFound: true}), false}, + } + for _, tc := range netCases { + if got := isRetryableErr(tc.err); got != tc.want { + t.Errorf("isRetryableErr(%s) = %v, want %v", tc.name, got, tc.want) + } + } +} diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go index b45a7ac69a..bac2bd0ad4 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -97,6 +97,32 @@ func validateTaskContext(taskCtx *TaskContext) error { return nil } +// noRefreshChunkInserter is the narrower inserter the ingestion path prefers: +// an engine that implements it can write chunks without waiting for an index +// refresh. +type noRefreshChunkInserter interface { + InsertChunksNoRefresh(ctx context.Context, chunks []map[string]interface{}, baseName string, datasetID string) ([]string, error) +} + +// insertChunksForIngestion writes chunks through the engine's no-refresh path +// when it offers one, falling back to the plain inserter otherwise. +// +// Waiting for a refresh costs ~0.5s per write (median; p90 1.0s, max 5.0s) with +// the KB index's refresh_interval=1000ms, and it buys nothing during ingestion: +// the index publishes the chunks on its own refresh cycle a moment later, so the +// only difference is up to a second before they are searchable. It is also what +// Python's ingestion does - it inserts with refresh=False +// (rag/svr/task_executor_refactor/chunk_service.py:386 and :423) while its API +// default stays "wait_for" - so this keeps the two ingestion paths in step. +func insertChunksForIngestion(eng engine.DocEngine) InsertFunc { + return func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + if bulk, ok := eng.(noRefreshChunkInserter); ok { + return bulk.InsertChunksNoRefresh(ctx, chunks, baseName, datasetID) + } + return eng.InsertChunks(ctx, chunks, baseName, datasetID) + } +} + func NewPipelineExecutor( taskCtx *TaskContext, canvasID string, @@ -112,14 +138,7 @@ func NewPipelineExecutor( taskCtx: taskCtx, canvasID: canvasID, docBulkSize: docBulkSize, - indexWriter: newChunkIndexWriter( - func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { - return engine.Get().InsertChunks(ctx, chunks, baseName, datasetID) - }, - fmt.Sprintf("ragflow_%s", taskCtx.Tenant.ID), - taskCtx.Doc.KbID, - docBulkSize, - ), + indexWriter: newChunkIndexWriter(insertChunksForIngestion(engine.Get()), fmt.Sprintf("ragflow_%s", taskCtx.Tenant.ID), taskCtx.Doc.KbID, docBulkSize), } svc.loadDSLFunc = svc.loadDSLFromCanvas svc.runPipelineFunc = svc.runPipelineWithDSL diff --git a/internal/parser/parser/markdown_parser.go b/internal/parser/parser/markdown_parser.go index 4986bb357b..5e35e8e173 100644 --- a/internal/parser/parser/markdown_parser.go +++ b/internal/parser/parser/markdown_parser.go @@ -128,7 +128,25 @@ func (p *MarkdownParser) String() string { // markdownNew is a thin constructor so the extension set is owned // in one place (both Parse and ParseWithResult consume it). func markdownNew() *mdparser.Parser { - extensions := mdparser.CommonExtensions | mdparser.AutoHeadingIDs | mdparser.NoEmptyLineBeforeBlock + // MathJax is masked out on purpose. CommonExtensions enables `$...$` + // inline math, whose payload lives on ast.Math rather than on a Text + // child; a real-world document that uses `$` as a currency symbol (a + // billing report, a price list) therefore had every span between two + // `$` silently deleted from the parsed text - with two far-apart `$` + // the whole middle of the document disappears. Python's markdown + // library has no math extension, so masking it also keeps the Go and + // Python parsers in parity. walkLeaf still handles ast.Math / + // ast.MathBlock defensively so that turning the extension back on + // cannot lose text again. + // OrderedListStart is not part of CommonExtensions, and without it the parser leaves + // ast.List.Start at zero for every ordered list: listEntries then falls back to 1 and + // renumbers the entries 1..N, so a list that the source spells "996. graphql: ..." is + // indexed as "1. graphql: ...". The numbers are then gone from the index (a query for + // "996" cannot match anything, and the gap counts against coverage), while Python's + // _extract_list_block keeps the raw lines - markers included - verbatim. Enabling it + // makes Start carry the number the source wrote (gomarkdown normalises a literal 1 to 0, + // which the fallback in listEntries turns back into 1). + extensions := mdparser.CommonExtensions&^mdparser.MathJax | mdparser.AutoHeadingIDs | mdparser.NoEmptyLineBeforeBlock | mdparser.OrderedListStart return mdparser.NewWithExtensions(extensions) } @@ -163,7 +181,7 @@ func renderMarkdownTablesInline(text string) (string, bool) { for i < len(lines) && strings.TrimSpace(lines[i]) == "" { i++ } - tableHTML := markdownlib.ToHTML([]byte(strings.Join(lines[start:i], "")), markdownNew(), nil) + tableHTML := markdownlib.ToHTML([]byte(strings.Join(balanceTableColumns(lines[start:i]), "")), markdownNew(), nil) // Wrap the inlined HTML in blank lines so gomarkdown // keeps it as a single HTML block (one item) instead of // re-parsing it into scattered cell text. See @@ -180,6 +198,84 @@ func renderMarkdownTablesInline(text string) (string, bool) { return buf.String(), changed } +// balanceTableColumns pads the header row and its separator up to the widest row's cell +// count, and returns the lines unchanged when no row is wider. +// +// This is a deliberate deviation from GFM, which ignores every cell beyond the header's +// column count. The corpus is Wikipedia-derived and its rows carry unescaped pipes inside a +// cell - an image spec such as "150x150px|alt=..." - so a row legitimately has more cells +// than its header, and GFM therefore drops the tail of the row without a word. On a list of +// World Heritage Sites that tail is the Year and Description columns: the words are in the +// file and nowhere in the index. Measured on the production index, tables of this shape cost +// 514 documents and 18,185 words that no query can reach. +// +// Padding the header keeps every cell, at the price of a few empty "), which the + // order-sensitive metric reports as missing text on every table-heavy document. + WordPct float64 `json:"word_pct"` + FileWords int `json:"file_words"` + MissingOccurrences int `json:"missing_occurrences"` + MissingWords int `json:"missing_words"` + // AbsentWords counts file words the chunks do not contain at all - not even inside a longer + // token. This is the "content never arrived" breakdown of MissingStrictWords. + AbsentWords int `json:"absent_words"` + // FusedWords counts file words reachable ONLY inside a longer chunk token, i.e. text the ingest + // glued together by dropping the separator between entries. Gluing is a defect, not a cosmetic + // difference: the merged token cannot be matched (a query for "China" does not match + // "chinachina"), so these words are counted as not served. They are broken out here only so a + // glued document can be told apart from one whose text is missing. + FusedWords int `json:"fused_words"` + // MissingStrictWords is the headline defect count: file words that are not available as their own + // token in the chunk text - absent plus fused. It does not tolerate gluing. + MissingStrictWords int `json:"missing_strict_words"` + StrictSamples []string `json:"strict_samples,omitempty"` + AbsentSamples []string `json:"absent_samples,omitempty"` + FusedSamples []string `json:"fused_samples,omitempty"` + ComputeMs float64 `json:"compute_ms"` + Gaps []gap `json:"gaps,omitempty"` + Error string `json:"error,omitempty"` +} + +// stats is the phase breakdown. Every counter is atomic because the phases run in different +// goroutines; wall time is what matters for the fetch (it is one socket per shard). +type stats struct { + fetchNanos int64 + fetchBytes int64 + pages int64 + decodeNanos int64 + computeNano int64 + docs int64 + emptyDocs int64 // documents emitted with no chunks at all though MySQL expected some + errs int64 + slowest sync.Map // name -> float64 ms +} + +func (s *stats) addEmpty() { atomic.AddInt64(&s.emptyDocs, 1) } +func (s *stats) addError() { atomic.AddInt64(&s.errs, 1) } + +func (s *stats) addFetch(d time.Duration, bytes int) { + atomic.AddInt64(&s.fetchNanos, int64(d)) + atomic.AddInt64(&s.fetchBytes, int64(bytes)) + atomic.AddInt64(&s.pages, 1) +} + +func (s *stats) addDecode(d time.Duration) { atomic.AddInt64(&s.decodeNanos, int64(d)) } +func (s *stats) addCompute(d time.Duration) { + atomic.AddInt64(&s.computeNano, int64(d)) + atomic.AddInt64(&s.docs, 1) +} + +func (s *stats) noteSlow(name string, ms float64) { s.slowest.Store(name, ms) } + +func (s *stats) report(wall time.Duration, cpu time.Duration) { + secs := wall.Seconds() + fmt.Printf("\n=== phase breakdown (wall %.1fs, cpu %.1fs, %.2f cores busy) ===\n", + secs, cpu.Seconds(), cpu.Seconds()/max(secs, 1e-9)) + fmt.Printf(" fetch : %6.1fs wall over %d pages, %.1f MB from ES\n", + float64(atomic.LoadInt64(&s.fetchNanos))/1e9, s.pages, float64(s.fetchBytes)/1e6) + fmt.Printf(" decode : %6.1fs cpu (JSON parse of those pages)\n", float64(atomic.LoadInt64(&s.decodeNanos))/1e9) + fmt.Printf(" compute: %6.1fs cpu over %d docs (tokenise+hash+compare)\n", + float64(atomic.LoadInt64(&s.computeNano))/1e9, s.docs) + if empty := atomic.LoadInt64(&s.emptyDocs); empty > 0 { + fmt.Printf(" WARNING: %d documents were emitted with no chunks at all although MySQL expects some\n"+ + " (paging lost them - check the sort key before trusting this run)\n", empty) + } + if errs := atomic.LoadInt64(&s.errs); errs > 0 { + fmt.Printf(" WARNING: %d shards aborted on an error - the numbers above are incomplete\n", errs) + } + type kv struct { + name string + ms float64 + } + var all []kv + s.slowest.Range(func(k, v any) bool { + all = append(all, kv{k.(string), v.(float64)}) + return true + }) + sort.Slice(all, func(i, j int) bool { return all[i].ms > all[j].ms }) + if len(all) > 0 { + fmt.Printf(" slowest documents: ") + for i := 0; i < len(all) && i < 5; i++ { + fmt.Printf("%s(%.0fms) ", all[i].name, all[i].ms) + } + fmt.Println() + } +} + +var whitespaceRun = regexp.MustCompile(`\s+`) + +func normalise(s string) string { return strings.TrimSpace(whitespaceRun.ReplaceAllString(s, " ")) } + +// --------------------------------------------------------------------------- +// Elasticsearch + +type esClient struct { + base string + auth string + http *http.Client +} + +type searchHits struct { + Hits struct { + Hits []struct { + Sort []any `json:"sort"` + Source struct { + Name string `json:"docnm_kwd"` + Content string `json:"content_with_weight"` + } `json:"_source"` + } `json:"hits"` + } `json:"hits"` +} + +func newESClient(base, user, pass string, conns int) *esClient { + transport := &http.Transport{ + // The default transport keeps only 2 idle connections per host, which is the wrong shape + // for a sharded reader: every shard past the second would re-dial instead of reusing. + MaxIdleConns: conns * 4, + MaxIdleConnsPerHost: conns * 2, + MaxConnsPerHost: 0, + IdleConnTimeout: 90 * time.Second, + DialContext: (&net.Dialer{Timeout: 30 * time.Second}).DialContext, + ForceAttemptHTTP2: false, + } + return &esClient{ + base: strings.TrimRight(base, "/"), + auth: base64.StdEncoding.EncodeToString([]byte(user + ":" + pass)), + http: &http.Client{Transport: transport, Timeout: 10 * time.Minute}, + } +} + +// searchRaw posts a query and returns the response body untouched, so the caller can decode it on +// its own goroutine while the next page is already being requested. +func (c *esClient) searchRaw(index string, body map[string]any) ([]byte, error) { + payload, err := json.Marshal(body) + if err != nil { + return nil, err + } + req, err := http.NewRequest(http.MethodPost, c.base+"/"+index+"/_search", bytes.NewReader(payload)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Basic "+c.auth) + resp, err := c.http.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("es %s: %s", resp.Status, truncate(string(raw), 300)) + } + return raw, nil +} + +func decodeHits(raw []byte) (*searchHits, error) { + var out searchHits + if err := json.Unmarshal(raw, &out); err != nil { + return nil, err + } + return &out, nil +} + +// produceShard walks one shard of the document set and emits each document as soon as every chunk +// MySQL expects has arrived. The shard runs two goroutines: one keeps ES busy, one parses the page +// that just arrived. +func produceShard(es *esClient, index, kb string, shard []doc, pageSize int, jobs chan<- job, st *stats, wg *sync.WaitGroup) { + defer wg.Done() + pending := make(map[string]*job, len(shard)) + names := make([]string, 0, len(shard)) + for _, d := range shard { + pending[d.Name] = &job{Name: d.Name, ID: d.ID, Chunks: d.Chunks} + names = append(names, d.Name) + } + body := map[string]any{ + "size": pageSize, + "_source": []string{"docnm_kwd", "content_with_weight"}, + "query": map[string]any{"bool": map[string]any{"filter": []any{ + map[string]any{"term": map[string]any{"kb_id": kb}}, + // A name list keeps the response to the target documents only; a name RANGE would + // drag in every other document of the dataset that sorts between them. + map[string]any{"terms": map[string]any{"docnm_kwd": names}}, + }}}, + // A stable sort is not optional here: search_after over sort:_doc is documented as + // unreliable and skipped documents once the pages got small (chunks=0 for documents whose + // chunks were never returned). Sorting by name then chunk order keeps a document's chunks + // contiguous across pages, which is also why paging can stop as soon as a name changes. + "sort": []any{ + map[string]any{"docnm_kwd": map[string]any{"order": "asc", "missing": "_last"}}, + map[string]any{"chunk_order_int": map[string]any{"order": "asc", "missing": "_last"}}, + }, + } + + type page struct { + raw []byte + err error + } + pages := make(chan page, 2) + go func() { + defer close(pages) + var after []any + for { + if after != nil { + body["search_after"] = after + } + start := time.Now() + raw, err := es.searchRaw(index, body) + st.addFetch(time.Since(start), len(raw)) + if err != nil { + pages <- page{err: err} + return + } + // Peek at the sort of the last hit to page on, then hand the raw body over: the + // decoding itself belongs to the other goroutine. + last, count, perr := lastSortAndCount(raw) + if perr != nil { + pages <- page{err: perr} + return + } + pages <- page{raw: raw} + if count == 0 || count < pageSize { + return + } + after = last + } + }() + + for p := range pages { + if p.err != nil { + fmt.Fprintf(os.Stderr, "\nFATAL shard fetch failed: %v\n", p.err) + st.addError() + break + } + start := time.Now() + resp, err := decodeHits(p.raw) + st.addDecode(time.Since(start)) + if err != nil { + fmt.Fprintf(os.Stderr, "\nFATAL shard decode failed (response shape changed?): %v\n", err) + st.addError() + break + } + for _, hit := range resp.Hits.Hits { + if j, ok := pending[hit.Source.Name]; ok { + j.Texts = append(j.Texts, hit.Source.Content) + } + } + for name, j := range pending { + if len(j.Texts) >= j.Chunks { + if len(j.Texts) == 0 && j.Chunks > 0 { + st.addEmpty() + } + jobs <- *j + delete(pending, name) + } + } + } + for _, j := range pending { + if len(j.Texts) == 0 && j.Chunks > 0 { + st.addEmpty() + } + jobs <- *j + } +} + +// lastSortAndCount reads just enough of a search response to continue paging: the sort value of +// the last hit and how many hits came back. +func lastSortAndCount(raw []byte) ([]any, int, error) { + var envelope struct { + Hits struct { + Hits []struct { + Sort []any `json:"sort"` + } `json:"hits"` + } `json:"hits"` + } + if err := json.Unmarshal(raw, &envelope); err != nil { + return nil, 0, err + } + n := len(envelope.Hits.Hits) + if n == 0 { + return nil, 0, nil + } + return envelope.Hits.Hits[n-1].Sort, n, nil +} + +// --------------------------------------------------------------------------- +// Coverage: character k-grams over an alphanumeric-only projection +// +// Comparing raw bytes, or even word tokens, over-reports missing text on this corpus. The ingest +// removes punctuation and closes the gap it leaves behind, so the source line +// +// * China at the Asian Games * China at the Paralympics +// +// is served as +// +// China at the Asian GamesChina at the Paralympics +// +// Every word is present, but "Games" and "China" are one token in the index and two in the file, so +// any token-boundary comparison calls that passage missing - which is how the first version of this +// tool managed to report a uniform ~78% coverage for 8,398 documents with no document above 90%. +// +// The fix is to compare the two sides on a projection that keeps only letters and digits, which is +// invariant under punctuation removal and word concatenation, and to match 16-character k-grams on +// it with a rolling hash. Each file character remembers where it came from, so matched k-grams still +// translate back into byte ranges of the original file - the report stays in file coordinates. + +const kGram = 16 + +// hashBase is the FNV-1a prime; any odd multiplier works for a rolling hash. +const hashBase = 1099511628211 + +type projection struct { + text []byte // alphanumeric-only, lowercased + offsets []int32 // offsets[i] = byte offset in the original string of text[i] +} + +// project folds a string onto its alphanumeric characters. Letters are lowercased so ASCII case +// differences in the ingest cannot register as missing content either. +func project(s string) projection { + p := projection{text: make([]byte, 0, len(s)), offsets: make([]int32, 0, len(s))} + for i, r := range s { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9': + p.text = append(p.text, byte(r)) + case r >= 'A' && r <= 'Z': + p.text = append(p.text, byte(r-'A'+'a')) + case r > unicode.MaxASCII && (unicode.IsLetter(r) || unicode.IsDigit(r)): + // Keep non-ASCII letters/digits as their rune bytes; both sides are folded the same way, + // so equality still holds where the ingest preserved the text. + start := len(p.text) + for _, b := range []byte(string(r)) { + p.text = append(p.text, b) + p.offsets = append(p.offsets, int32(i)) + _ = start + } + continue + default: + continue + } + p.offsets = append(p.offsets, int32(i)) + } + return p +} + +// rollingKGramHashes hashes every k-length window of a byte slice. +func rollingKGramHashes(b []byte, k int) []uint64 { + if len(b) < k { + return nil + } + out := make([]uint64, len(b)-k+1) + var pow uint64 = 1 + for i := 0; i < k-1; i++ { + pow *= hashBase + } + var h uint64 + for i := 0; i < k; i++ { + h = h*hashBase + uint64(b[i]) + } + out[0] = h + for i := k; i < len(b); i++ { + h = (h-uint64(b[i-k])*pow)*hashBase + uint64(b[i]) + out[i-k+1] = h + } + return out +} + +// coverageAlnum covers every file byte that takes part in a k-gram the chunk text also contains. +// Whatever is left over in runs longer than minGap bytes is a passage the index does not serve. +func coverageAlnum(fileNorm, joined string, minGap int) (covered []byte, fileAlnum int, gaps []gap) { + fileProj := project(fileNorm) + chunkProj := project(joined) + fileAlnum = len(fileProj.text) + covered = make([]byte, len(fileNorm)) + if len(fileProj.text) < kGram || len(chunkProj.text) < kGram { + return covered, fileAlnum, nil + } + chunkHashes := rollingKGramHashes(chunkProj.text, kGram) + set := make(map[uint64]struct{}, len(chunkHashes)) + for _, h := range chunkHashes { + set[h] = struct{}{} + } + for i, h := range rollingKGramHashes(fileProj.text, kGram) { + if _, ok := set[h]; !ok { + continue + } + from := int(fileProj.offsets[i]) + to := int(fileProj.offsets[i+kGram-1]) + 1 + for b := from; b < to; b++ { + covered[b] = 1 + } + } + // Uncovered runs, reported in file coordinates. Text the ingest merely re-punctuated shows up as + // scattered single bytes and never reaches minGap; a dropped passage does. + start := -1 + for i := 0; i <= len(covered); i++ { + flag := i < len(covered) && covered[i] == 1 + switch { + case !flag && start < 0: + start = i + case start >= 0 && (flag || i == len(covered)): + // The run that began at `start` ends here: either coverage resumed, or + // the sentinel iteration (i == len(covered), flag == false) was reached. + // A tail the ingest never covered is a gap like any other — without this + // arm it is discarded while its bytes still count toward GapChars. + if i-start >= minGap { + gaps = append(gaps, gap{Start: start, Length: i - start, + Context: fileNorm[max(0, start-40):min(len(fileNorm), i+40)]}) + } + start = -1 + } + } + return covered, fileAlnum, gaps +} + +// coverageWindow is the byte-window comparison kept for the record (see the package comment): it +// reports markup rewrites as missing text. +func coverageWindow(fileNorm, joined string, window, stride, minGap int) ([]byte, []gap) { + covered := make([]byte, len(fileNorm)) + if joined == "" || len(fileNorm) == 0 { + return covered, nil + } + for pos := 0; pos < len(fileNorm); { + end := min(pos+window, len(fileNorm)) + if strings.Contains(joined, fileNorm[pos:end]) { + for i := pos; i < end; i++ { + covered[i] = 1 + } + pos += max(end-pos-stride, 1) + continue + } + pos += stride + } + var gaps []gap + start := -1 + for i := 0; i <= len(covered); i++ { + flag := i < len(covered) && covered[i] == 1 + switch { + case !flag && start < 0: + start = i + case start >= 0 && (flag || i == len(covered)): + // The run that began at `start` ends here: either coverage resumed, or + // the sentinel iteration (i == len(covered), flag == false) was reached. + // A tail the ingest never covered is a gap like any other — without this + // arm it is discarded while its bytes still count toward GapChars. + if i-start >= minGap { + gaps = append(gaps, gap{Start: start, Length: i - start, + Context: fileNorm[max(0, start-40):min(len(fileNorm), i+40)]}) + } + start = -1 + } + } + return covered, gaps +} + +// tokeniseWords splits lowercased alphanumeric runs out of the text. +func tokeniseWords(s string) []string { + out := make([]string, 0, len(s)/6+8) + start := -1 + for i, r := range s { + if unicode.IsLetter(r) || unicode.IsDigit(r) { + if start < 0 { + start = i + } + continue + } + if start >= 0 { + out = append(out, strings.ToLower(s[start:i])) + start = -1 + } + } + if start >= 0 { + out = append(out, strings.ToLower(s[start:])) + } + return out +} + +type wordStats struct { + Total int // word occurrences in the file + Short int // distinct words whose occurrence count falls short in the chunks + OccShort int // occurrences that fall short (multiset attrition) + Absent []string // file words the chunks do not contain at all, not even inside a longer token + Fused []string // file words reachable only inside a longer chunk token + Strict []string // Absent + Fused: words with no token of their own (sorted) + MultisetPct float64 // share of the file's occurrences the chunks carry +} + +// computeWordStats is the order-insensitive answer to two questions at once, and it deliberately +// gives gluing no quarter. +// +// A file word is SERVED only if the chunk text carries it as its own token. Two ways to fail: +// +// Absent - the chunk text does not contain the word at all (content never arrived); +// Fused - the word exists only inside a longer token, because the ingest dropped a separator +// between entries ("GamesChina", "HathiTrustThe"). The characters are there, but the +// token is not: a query for "China" cannot match "chinachina", so retrieval loses the +// word. Counting it as served would hide a real defect behind a substring test. +// +// Strict is the union of the two - the words this document does not actually serve - and it is what +// the sweep flags on. The split exists only to tell a glued document from one whose text is gone. +// +// MultisetPct reads low (~50%) for documents whose source file repeats the same blocks twice: the +// ingest de-duplicates the repeats, so the ratio of file words to chunk words lands at ~1.98. +func computeWordStats(fileNorm, joined string) wordStats { + stats := wordStats{} + fileWords := map[string]int{} + for _, w := range tokeniseWords(fileNorm) { + fileWords[w]++ + } + chunkWords := map[string]int{} + for _, w := range tokeniseWords(joined) { + chunkWords[w]++ + } + stream := project(joined).text + for w, n := range fileWords { + stats.Total += n + have := chunkWords[w] + if have < n { + stats.Short++ + stats.OccShort += n - have + } + if have > 0 { + continue // served as its own token + } + if bytes.Contains(stream, []byte(w)) { + stats.Fused = append(stats.Fused, w) + } else { + stats.Absent = append(stats.Absent, w) + } + } + if stats.Total > 0 { + stats.MultisetPct = 100 * float64(stats.Total-stats.OccShort) / float64(stats.Total) + } else { + stats.MultisetPct = 100 + } + sort.Strings(stats.Absent) + sort.Strings(stats.Fused) + stats.Strict = append(append(make([]string, 0, len(stats.Absent)+len(stats.Fused)), stats.Absent...), stats.Fused...) + sort.Strings(stats.Strict) + return stats +} + +func measure(corpusDir string, j job, opts options, st *stats) result { + start := time.Now() + defer func() { st.noteSlow(j.Name, float64(time.Since(start).Microseconds())/1000) }() + res := result{Name: j.Name, ID: j.ID, MySQLChunks: j.Chunks, ChunkCount: len(j.Texts)} + raw, err := os.ReadFile(filepath.Join(corpusDir, j.Name)) + if err != nil { + res.Error = "FILE MISSING" + return res + } + fileNorm := normalise(string(raw)) + res.FileChars = len(fileNorm) + if fileNorm == "" { + res.Error = "EMPTY FILE" + return res + } + joined := normalise(strings.Join(j.Texts, "\n")) + + computeStart := time.Now() + var covered []byte + var gaps []gap + if opts.method == "window" { + covered, gaps = coverageWindow(fileNorm, joined, opts.window, opts.stride, opts.minGap) + } else { + covered, res.FileAlnum, gaps = coverageAlnum(fileNorm, joined, opts.minGap) + } + ws := computeWordStats(fileNorm, joined) + res.WordPct, res.FileWords = ws.MultisetPct, ws.Total + res.MissingWords, res.MissingOccurrences = ws.Short, ws.OccShort + res.AbsentWords, res.FusedWords = len(ws.Absent), len(ws.Fused) + // The strict count has to be copied here explicitly: it is the verdict the summary keys off + // ("a word is served only when it is its own token"), and when it was left at its zero value the + // report said "0 documents with unserved words" while listing thousands of glued ones - a gate + // that passes because it was never wired, not because the index is clean. + res.MissingStrictWords = len(ws.Strict) + res.AbsentSamples, res.FusedSamples = firstN(ws.Absent, 8), firstN(ws.Fused, 8) + res.StrictSamples = firstN(ws.Strict, 8) + st.addCompute(time.Since(computeStart)) + + count := 0 + for _, b := range covered { + count += int(b) + } + res.CoveredChar = count + res.GapChars = len(fileNorm) - count + res.CoveragePct = float64(count) / float64(len(fileNorm)) * 100 + res.ComputeMs = float64(time.Since(start).Microseconds()) / 1000 + if len(gaps) > 5 { + gaps = gaps[:5] + } + res.Gaps = gaps + return res +} + +// --------------------------------------------------------------------------- +// helpers + +type options struct { + method string + window int + stride int + minGap int +} + +func firstN(in []string, n int) []string { + if len(in) <= n { + return in + } + return in[:n] +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "…" +} + +func envPassword(path, key string) string { + fh, err := os.Open(path) + if err != nil { + return "" + } + defer fh.Close() + scanner := bufio.NewScanner(fh) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if strings.HasPrefix(line, key+"=") { + return strings.TrimPrefix(line, key+"=") + } + } + return "" +} + +func loadDocs(path string, limit int) ([]doc, error) { + fh, err := os.Open(path) + if err != nil { + return nil, err + } + defer fh.Close() + var out []doc + scanner := bufio.NewScanner(fh) + scanner.Buffer(make([]byte, 1<<20), 1<<20) + for scanner.Scan() { + fields := strings.Split(scanner.Text(), "\t") + if len(fields) < 3 || fields[0] == "" { + continue + } + var n int + fmt.Sscanf(fields[2], "%d", &n) + out = append(out, doc{ID: fields[0], Name: fields[1], Chunks: n}) + if limit > 0 && len(out) >= limit { + break + } + } + return out, scanner.Err() +} + +func loadDone(path string) map[string]bool { + done := map[string]bool{} + fh, err := os.Open(path) + if err != nil { + return done + } + defer fh.Close() + scanner := bufio.NewScanner(fh) + scanner.Buffer(make([]byte, 1<<20), 1<<20) + for scanner.Scan() { + var r result + if err := json.Unmarshal(scanner.Bytes(), &r); err == nil && r.Name != "" { + done[r.Name] = true + } + } + return done +} + +func loadNames(path string) map[string]bool { + if path == "" { + return nil + } + names := map[string]bool{} + fh, err := os.Open(path) + if err != nil { + return names + } + defer fh.Close() + scanner := bufio.NewScanner(fh) + for scanner.Scan() { + if n := strings.TrimSpace(scanner.Text()); n != "" { + names[n] = true + } + } + return names +} + +func processCPU() time.Duration { + var ru syscall.Rusage + if err := syscall.Getrusage(syscall.RUSAGE_SELF, &ru); err != nil { + return 0 + } + user := time.Duration(ru.Utime.Sec)*time.Second + time.Duration(ru.Utime.Usec)*time.Microsecond + sys := time.Duration(ru.Stime.Sec)*time.Second + time.Duration(ru.Stime.Usec)*time.Microsecond + return user + sys +} + +func main() { + var ( + docsPath = flag.String("docs", "", "TSV of doc id / name / chunk_num rows (required)") + corpusDir = flag.String("corpus", "", "directory holding the source file named by each row (required)") + outDir = flag.String("out", "outputs/coverage", "output directory") + esURL = flag.String("es", "http://localhost:1200", "Elasticsearch base URL") + esUser = flag.String("es-user", "elastic", "Elasticsearch user") + esPass = flag.String("es-pass", "", "Elasticsearch password (default: ELASTIC_PASSWORD from --env)") + envFile = flag.String("env", "docker/.env", "file to read ELASTIC_PASSWORD from") + index = flag.String("index", "", "ES index to read (required)") + kb = flag.String("kb", "", "dataset id every checked document belongs to (required)") + shards = flag.Int("shards", 8, "parallel ES producer goroutines") + workers = flag.Int("workers", 0, "coverage workers (0 = number of CPUs)") + pageSize = flag.Int("page", 1000, "chunk documents per ES page") + method = flag.String("method", "alnum", "alnum (k-grams over letters/digits) or window (raw byte windows)") + window = flag.Int("window", 200, "byte window length for --method window") + stride = flag.Int("stride", 50, "byte window step for --method window") + minGap = flag.Int("min-gap", 200, "report uncovered runs at least this many bytes") + limit = flag.Int("limit", 0, "only the first N documents (calibration)") + onlyFile = flag.String("only-file", "", "only documents named in this file") + every = flag.Int("every", 200, "progress line every N documents") + ) + flag.Parse() + + // Nothing about the document set is baked in. A default corpus or index would mean a run that + // forgot a flag still produced numbers - about whatever was compiled in. + var missing []string + for _, required := range []struct { + name string + value string + }{ + {"docs", *docsPath}, + {"corpus", *corpusDir}, + {"index", *index}, + {"kb", *kb}, + } { + if strings.TrimSpace(required.value) == "" { + missing = append(missing, "-"+required.name) + } + } + if len(missing) > 0 { + fmt.Fprintf(os.Stderr, "required flag(s) missing: %s\n", strings.Join(missing, ", ")) + flag.Usage() + os.Exit(2) + } + + if *workers <= 0 { + *workers = runtime.NumCPU() + } + if *esPass == "" { + *esPass = envPassword(*envFile, "ELASTIC_PASSWORD") + } + if *esPass == "" { + fmt.Fprintln(os.Stderr, "no ES password: pass --es-pass or set ELASTIC_PASSWORD in --env") + os.Exit(2) + } + if err := os.MkdirAll(*outDir, 0o755); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + opts := options{method: *method, window: *window, stride: *stride, minGap: *minGap} + + all, err := loadDocs(*docsPath, *limit) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + outPath := filepath.Join(*outDir, "coverage_go.jsonl") + done := loadDone(outPath) + only := loadNames(*onlyFile) + todo := make([]doc, 0, len(all)) + for _, d := range all { + if done[d.Name] { + continue + } + if len(only) > 0 && !only[d.Name] { + continue + } + todo = append(todo, d) + } + fmt.Printf("docs=%d already=%d todo=%d shards=%d workers=%d method=%s min-gap=%d bytes\n", + len(all), len(done), len(todo), *shards, *workers, opts.method, opts.minGap) + + // Shard the work into contiguous blocks so each shard's ES pages are a narrow slice of the + // index rather than a random sprinkle of documents. + shardSets := make([][]doc, 0, *shards) + per := (len(todo) + *shards - 1) / max(*shards, 1) + for start := 0; start < len(todo); start += per { + shardSets = append(shardSets, todo[start:min(start+per, len(todo))]) + } + + es := newESClient(*esURL, *esUser, *esPass, *shards) + st := &stats{} + jobs := make(chan job, *workers*2) + results := make(chan result, *workers*2) + var producerWG sync.WaitGroup + for _, shard := range shardSets { + if len(shard) == 0 { + continue + } + producerWG.Add(1) + go produceShard(es, *index, *kb, shard, *pageSize, jobs, st, &producerWG) + } + go func() { producerWG.Wait(); close(jobs) }() + + var workerWG sync.WaitGroup + for i := 0; i < *workers; i++ { + workerWG.Add(1) + go func() { + defer workerWG.Done() + for j := range jobs { + results <- measure(*corpusDir, j, opts, st) + } + }() + } + go func() { workerWG.Wait(); close(results) }() + + sink, err := os.OpenFile(outPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + buffered := bufio.NewWriter(sink) + started := time.Now() + var measured int64 + lossy := make([]result, 0, 64) + glued := make([]result, 0, 64) + for res := range results { + payload, _ := json.Marshal(res) + buffered.Write(payload) + buffered.WriteByte('\n') + if atomic.AddInt64(&measured, 1)%10 == 0 { + buffered.Flush() + } + // The verdict is strict: a word that exists only glued inside a longer token is not served + // either, so gluing is a defect rather than a lesser category. The absent/fused split is + // kept only to tell the two failure modes apart. + if res.Error != "" || res.MissingStrictWords > 0 { + lossy = append(lossy, res) + } + if res.FusedWords > 0 { + glued = append(glued, res) + } + n := atomic.LoadInt64(&measured) + if int(n)%*every == 0 || int(n) == len(todo) { + elapsed := time.Since(started).Seconds() + rate := float64(n) / elapsed + fmt.Printf(" %d/%d %.1f docs/s ETA %.1f min\n", n, len(todo), rate, float64(len(todo)-int(n))/rate/60) + } + } + buffered.Flush() + sink.Close() + + sort.Slice(lossy, func(i, j int) bool { return lossy[i].MissingStrictWords > lossy[j].MissingStrictWords }) + sort.Slice(glued, func(i, j int) bool { return glued[i].FusedWords > glued[j].FusedWords }) + totalStrict, totalAbsent, totalFused := 0, 0, 0 + for _, r := range lossy { + totalStrict += r.MissingStrictWords + totalAbsent += r.AbsentWords + } + for _, r := range glued { + totalFused += r.FusedWords + } + fmt.Printf("\n=== summary: %d documents in %.1f min ===\n", measured, time.Since(started).Minutes()) + fmt.Printf(" documents with unserved words / errors : %d (sum %d words)\n", len(lossy), totalStrict) + // Consistency check, not a warning to ignore: every glued document is by definition a document + // with unserved words (gluing is not tolerated), so an empty strict verdict while glued + // documents exist means the strict count was not populated - the exact way this report managed + // to look clean while flagging thousands of glued words in its own next line. + if len(lossy) == 0 && len(glued) > 0 { + fmt.Printf(" INCONSISTENT: %d documents have glued words but none counted as unserved; the strict verdict is not wired\n", len(glued)) + os.Exit(1) + } + fmt.Printf(" of which content never arrived : %d words\n", totalAbsent) + fmt.Printf(" of which glued into a longer token : %d words in %d documents (NOT tolerated)\n", + totalFused, len(glued)) + for i, r := range lossy { + if i < 20 { + fmt.Printf(" STRICT %-16s unserved=%-6d (absent=%-5d glued=%-5d) words=%.2f%% chunks=%d %v\n", + r.Name, r.MissingStrictWords, r.AbsentWords, r.FusedWords, r.WordPct, r.ChunkCount, r.StrictSamples) + } + } + lossyNames := make([]string, 0, len(lossy)) + for _, r := range lossy { + lossyNames = append(lossyNames, r.Name) + } + gluedNames := make([]string, 0, len(glued)) + for _, r := range glued { + gluedNames = append(gluedNames, r.Name) + } + flagPath := filepath.Join(*outDir, "below99.txt") + _ = os.WriteFile(flagPath, []byte(strings.Join(lossyNames, "\n")+"\n"), 0o644) + fusedPath := filepath.Join(*outDir, "fused.txt") + _ = os.WriteFile(fusedPath, []byte(strings.Join(gluedNames, "\n")+"\n"), 0o644) + st.report(time.Since(started), processCPU()) + fmt.Printf("wrote %s, %s and %s\n", outPath, flagPath, fusedPath) +} diff --git a/tools/mdprobe/main.go b/tools/mdprobe/main.go new file mode 100644 index 0000000000..85d519411f --- /dev/null +++ b/tools/mdprobe/main.go @@ -0,0 +1,83 @@ +// Command mdprobe runs the ingest Markdown parser over a corpus file and reports, item by item, +// what the parser produced - the fastest way to tell a parser-side loss from a downstream one. +// +// The coverage sweep flagged documents whose indexed chunks do not contain passages of the source +// file; on this corpus the missing passages are tables. The parser is the first suspect, so this +// probe answers one question: does the table text survive the parser? +// +// Usage: +// +// go run ./tools/mdprobe [needle ...] +// +// With needles it also prints, per needle, which item contains it (or that none does). +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "ragflow/internal/parser/parser" +) + +func main() { + if len(os.Args) < 2 { + fmt.Fprintln(os.Stderr, "usage: mdprobe [needle ...]") + os.Exit(2) + } + path := os.Args[1] + needles := os.Args[2:] + data, err := os.ReadFile(path) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + p, err := parser.NewMarkdownParser(parser.GoMarkdown) + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + // The pipeline's setups for the markdown family: output_format=json. + p.ConfigureFromSetup(map[string]any{"output_format": "json"}) + res := p.ParseWithResult(context.Background(), filepath.Base(path), data) + + fmt.Printf("file=%s bytes=%d output_format=%s items=%d\n", filepath.Base(path), len(data), res.OutputFormat, len(res.JSON)) + var total int + for i, item := range res.JSON { + text, _ := item["text"].(string) + total += len(text) + ck := fmt.Sprint(item["ck_type"]) + docType := fmt.Sprint(item["doc_type_kwd"]) + fmt.Printf(" [%3d] %-7s %-6s %6d chars %q\n", i, ck, docType, len(text), truncate(text, 90)) + } + fmt.Printf("parser emitted %d chars total (file %d)\n", total, len(data)) + + for _, needle := range needles { + lower := strings.ToLower(needle) + found := -1 + for i, item := range res.JSON { + if text, _ := item["text"].(string); strings.Contains(strings.ToLower(text), lower) { + found = i + break + } + } + inFile := strings.Contains(strings.ToLower(string(data)), lower) + fmt.Printf("needle %-28q in file=%v in parser output=%v (item %d)\n", needle, inFile, found >= 0, found) + } + + if os.Getenv("MDPROBE_JSON") != "" { + out, _ := json.MarshalIndent(res.JSON, "", " ") + fmt.Println(string(out)) + } +} + +func truncate(s string, n int) string { + s = strings.ReplaceAll(s, "\n", "\\n") + if len(s) <= n { + return s + } + return s[:n] + "…" +} diff --git a/web/package.json b/web/package.json index de05c02d72..6afa005e59 100644 --- a/web/package.json +++ b/web/package.json @@ -203,5 +203,11 @@ }, "engines": { "node": ">=18.20.4" + }, + "allowScripts": { + "esbuild@0.25.12": true, + "esbuild@0.27.2": true, + "@swc/core@1.15.7": true, + "less@4.5.1": true } } diff --git a/web/src/assets/exa.png b/web/src/assets/exa.png new file mode 100644 index 0000000000..82bdbacc5d Binary files /dev/null and b/web/src/assets/exa.png differ diff --git a/web/src/assets/firecrawl.png b/web/src/assets/firecrawl.png new file mode 100644 index 0000000000..a30a034a3a Binary files /dev/null and b/web/src/assets/firecrawl.png differ diff --git a/web/src/assets/linkup.png b/web/src/assets/linkup.png new file mode 100644 index 0000000000..f0b72ff9c7 Binary files /dev/null and b/web/src/assets/linkup.png differ diff --git a/web/src/assets/svg/brave.svg b/web/src/assets/svg/brave.svg new file mode 100644 index 0000000000..8c56fec961 --- /dev/null +++ b/web/src/assets/svg/brave.svg @@ -0,0 +1 @@ +Brave diff --git a/web/src/assets/svg/parallel.svg b/web/src/assets/svg/parallel.svg new file mode 100644 index 0000000000..ccd0c3198a --- /dev/null +++ b/web/src/assets/svg/parallel.svg @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/web/src/components/message-input/next.tsx b/web/src/components/message-input/next.tsx index 952e462e4e..ff32a3c211 100644 --- a/web/src/components/message-input/next.tsx +++ b/web/src/components/message-input/next.tsx @@ -37,7 +37,7 @@ import { } from '@/components/ui/tooltip'; import { CircleStop, Globe, Paperclip, Send, Upload, X } from 'lucide-react'; import * as React from 'react'; -import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { toast } from 'sonner'; import storage from '@/utils/authorization-util'; @@ -51,6 +51,13 @@ export type NextMessageInputOnPressEnterParameter = { omitSessionId?: boolean; }; +// The selector only offers numeric thinking levels; a value persisted by a +// since-removed mode is not one, so it falls back to the default rather than +// reaching the request as NaN. +function normalizeThinkingLevel(level: string) { + return /^[0-4]$/.test(level) ? level : '1'; +} + interface NextMessageInputProps { disabled: boolean; value: string; @@ -98,7 +105,7 @@ export function NextMessageInput({ ); const [enableThinking, setEnableThinking] = useState(() => - storage.getThinkingLevel(), + normalizeThinkingLevel(storage.getThinkingLevel()), ); const [enableInternet, setEnableInternet] = useState(false); diff --git a/web/src/components/web-search-form-field.tsx b/web/src/components/web-search-form-field.tsx index e861baa270..943b432cbf 100644 --- a/web/src/components/web-search-form-field.tsx +++ b/web/src/components/web-search-form-field.tsx @@ -14,6 +14,11 @@ * limitations under the License. */ +import braveLogo from '@/assets/svg/brave.svg'; +import exaLogo from '@/assets/exa.png'; +import firecrawlLogo from '@/assets/firecrawl.png'; +import linkupLogo from '@/assets/linkup.png'; +import parallelLogo from '@/assets/svg/parallel.svg'; import queritLogo from '@/assets/querit.png'; import serplyLogo from '@/assets/serply.png'; import tavilyLogo from '@/assets/svg/tavily.svg'; @@ -21,6 +26,7 @@ import youcomLogo from '@/assets/svg/youcom.svg'; import { RAGFlowFormItem } from '@/components/ragflow-form'; import { WebSearchProvider } from '@/constants/chat'; import { useTranslate } from '@/hooks/common-hooks'; +import { isWebSearchApiKeyRequired } from '@/pages/next-chats/chat/web-search-api-key'; import { prefixName } from '@/utils/form'; import { useFormContext, useWatch } from 'react-hook-form'; import PasswordInput from './originui/password-input'; @@ -30,11 +36,33 @@ interface IProps { prefix?: string; } -const providerOptions = [ +// One entry per provider. `name` is the BRAND, so it stays untranslated; it is what +// the key field's label is interpolated into (chat.webSearchApiKeyLabel). +const webSearchProviderCatalog = [ { - name: 'Tavily', - logo: tavilyLogo, - value: WebSearchProvider.Tavily, + name: 'Brave Search', + logo: braveLogo, + value: WebSearchProvider.Brave, + }, + { + name: 'Exa', + logo: exaLogo, + value: WebSearchProvider.Exa, + }, + { + name: 'Firecrawl', + logo: firecrawlLogo, + value: WebSearchProvider.Firecrawl, + }, + { + name: 'Linkup', + logo: linkupLogo, + value: WebSearchProvider.Linkup, + }, + { + name: 'Parallel', + logo: parallelLogo, + value: WebSearchProvider.Parallel, }, { name: 'Querit', @@ -46,12 +74,20 @@ const providerOptions = [ logo: serplyLogo, value: WebSearchProvider.Serply, }, + { + name: 'Tavily', + logo: tavilyLogo, + value: WebSearchProvider.Tavily, + }, { name: 'You.com', logo: youcomLogo, value: WebSearchProvider.YouCom, }, -] +]; + +const providerOptions = webSearchProviderCatalog + .slice() .sort((left, right) => left.name.localeCompare(right.name)) .map(({ name, logo, value }) => ({ label: ( @@ -65,34 +101,67 @@ const providerOptions = [ {name} ), + // The dropdown filters on value + keywords, so the shown name has to be + // searchable too: typing "You.com" must find the "youcom" entry. + keywords: [name], value, })); +const providerDisplayName = (provider?: WebSearchProvider) => + webSearchProviderCatalog.find((entry) => entry.value === provider)?.name ?? + ''; + const providerKeyConfig = { - [WebSearchProvider.Tavily]: { - name: 'prompt_config.tavily_api_key', - label: 'Tavily API Key', - tip: 'tavilyApiKeyTip', - placeholder: 'tavilyApiKeyMessage', - helpUrl: 'https://app.tavily.com/home', + [WebSearchProvider.Brave]: { + name: 'prompt_config.brave_api_key', + tip: 'braveApiKeyTip', + placeholder: 'braveApiKeyMessage', + helpUrl: 'https://brave.com/search/api/', + }, + [WebSearchProvider.Exa]: { + name: 'prompt_config.exa_api_key', + tip: 'exaApiKeyTip', + placeholder: 'exaApiKeyMessage', + helpUrl: 'https://dashboard.exa.ai/api-keys', + }, + [WebSearchProvider.Firecrawl]: { + name: 'prompt_config.firecrawl_api_key', + tip: 'firecrawlApiKeyTip', + placeholder: 'firecrawlApiKeyMessage', + helpUrl: 'https://www.firecrawl.dev/app/api-keys', + }, + [WebSearchProvider.Linkup]: { + name: 'prompt_config.linkup_api_key', + tip: 'linkupApiKeyTip', + placeholder: 'linkupApiKeyMessage', + helpUrl: 'https://app.linkup.so', + }, + [WebSearchProvider.Parallel]: { + name: 'prompt_config.parallel_api_key', + tip: 'parallelApiKeyTip', + placeholder: 'parallelApiKeyMessage', + helpUrl: 'https://platform.parallel.ai', }, [WebSearchProvider.Querit]: { name: 'prompt_config.querit_api_key', - label: 'Querit API Key', tip: 'queritApiKeyTip', placeholder: 'queritApiKeyMessage', helpUrl: 'https://querit.ai', }, [WebSearchProvider.Serply]: { name: 'prompt_config.serply_api_key', - label: 'Serply API Key', tip: 'serplyApiKeyTip', placeholder: 'serplyApiKeyMessage', helpUrl: 'https://serply.io', }, + [WebSearchProvider.Tavily]: { + name: 'prompt_config.tavily_api_key', + tip: 'tavilyApiKeyTip', + placeholder: 'tavilyApiKeyMessage', + helpUrl: 'https://app.tavily.com/home', + }, [WebSearchProvider.YouCom]: { name: 'prompt_config.youcom_api_key', - label: 'You.com API Key', tip: 'youcomApiKeyTip', placeholder: 'youcomApiKeyMessage', helpUrl: @@ -109,6 +178,12 @@ export function WebSearchFormField({ prefix = '' }: IProps) { name: providerName, }); const keyConfig = providerKeyConfig[selectedProvider as WebSearchProvider]; + // Leaving the key blank on a provider that needs one is a SILENT failure: the + // chat box simply never offers the Internet switch. Mark the field required + // so the asterisk says so before the user has to discover it. + const keyRequired = isWebSearchApiKeyRequired( + selectedProvider as WebSearchProvider, + ); return ( <> @@ -133,11 +208,22 @@ export function WebSearchFormField({ prefix = '' }: IProps) { + !keyRequired || + Boolean(String(value ?? '').trim()) || + t('webSearchApiKeyRequired'), + }} description={ - {t('tavilyApiKeyHelp')} + {t('webSearchApiKeyHelp')} } > diff --git a/web/src/constants/chat.ts b/web/src/constants/chat.ts index a898b983fc..f242998474 100644 --- a/web/src/constants/chat.ts +++ b/web/src/constants/chat.ts @@ -57,15 +57,28 @@ export enum DatasetMetadata { } export enum WebSearchProvider { - Tavily = 'tavily', + Brave = 'brave', + Exa = 'exa', + Firecrawl = 'firecrawl', + Linkup = 'linkup', + Parallel = 'parallel', Querit = 'querit', Serply = 'serply', + Tavily = 'tavily', YouCom = 'youcom', } /** - * Providers usable with no credentials at all. You.com serves a rate-limited - * keyless endpoint; every other provider requires a key before it can be used. + * Providers usable with no credentials at all — the key field is optional for + * these, and the Internet switch stays available when it is left blank. + * + * You.com serves a dedicated keyless ENDPOINT: a different path from the keyed + * one, the same response shape, rate-limited per source IP. Adding a key + * switches to the keyed endpoint and lifts that limit. + * + * "Free tier" is not the same as "keyless". Exa's free tier (1,000 + * requests/month, no credit card) still requires a key on every request, so it + * is deliberately absent from this list. */ export const KEYLESS_WEB_SEARCH_PROVIDERS: readonly WebSearchProvider[] = [ WebSearchProvider.YouCom, diff --git a/web/src/interfaces/database/chat.ts b/web/src/interfaces/database/chat.ts index c636fdcc02..ae7a8df662 100644 --- a/web/src/interfaces/database/chat.ts +++ b/web/src/interfaces/database/chat.ts @@ -20,9 +20,14 @@ export interface PromptConfig { use_kg: boolean; reasoning?: boolean; cross_languages?: Array; - tavily_api_key?: string; + brave_api_key?: string; + exa_api_key?: string; + firecrawl_api_key?: string; + linkup_api_key?: string; + parallel_api_key?: string; querit_api_key?: string; serply_api_key?: string; + tavily_api_key?: string; youcom_api_key?: string; web_search_provider?: WebSearchProvider; toc_enhance?: boolean; diff --git a/web/src/locales/ar.ts b/web/src/locales/ar.ts index d76959edf7..e63ba6a59f 100644 --- a/web/src/locales/ar.ts +++ b/web/src/locales/ar.ts @@ -806,7 +806,7 @@ export default { tavilyApiKeyTip: 'إذا تم تعيين مفتاح API بشكل صحيح هنا، فسيتم استخدام عمليات البحث على الويب المستندة إلى Tavily لتكملة استرجاع قاعدة المعرفة.', tavilyApiKeyMessage: 'الرجاء إدخال مفتاح Tavily API الخاص بك', - tavilyApiKeyHelp: 'كيفية الحصول عليه؟', + webSearchApiKeyHelp: 'كيفية الحصول عليه؟', crossLanguage: 'البحث عبر اللغات', crossLanguageTip: 'حدد لغة واحدة أو أكثر للبحث بين اللغات. إذا لم يتم تحديد أي لغة، فسيقوم النظام بالبحث باستخدام الاستعلام الأصلي.', diff --git a/web/src/locales/bg.ts b/web/src/locales/bg.ts index 32a401cbbe..786f8daa90 100644 --- a/web/src/locales/bg.ts +++ b/web/src/locales/bg.ts @@ -884,7 +884,7 @@ The above is the content you need to summarize.`, tavilyApiKeyTip: 'Ако API ключът е зададен правилно тук, ще бъдат използвани уеб търсения базирани на Tavily за допълване на извличането от базата от знания.', tavilyApiKeyMessage: 'Моля, въведете вашия Tavily API ключ', - tavilyApiKeyHelp: 'Как да го получите?', + webSearchApiKeyHelp: 'Как да го получите?', crossLanguage: 'Крос-езиково търсене', crossLanguageTip: `Изберете един или повече езици за крос-езиково търсене. Ако не е избран език, системата търси с оригиналната заявка.`, createChat: 'Създай чат', diff --git a/web/src/locales/de.ts b/web/src/locales/de.ts index 0f6aba99c0..41bdcc0429 100644 --- a/web/src/locales/de.ts +++ b/web/src/locales/de.ts @@ -900,7 +900,7 @@ Diese Auto-Tag-Funktion verbessert den Abruf, indem sie eine weitere Schicht dom tavilyApiKeyTip: 'Wenn hier ein API-SCHLÜSSEL korrekt eingestellt ist, werden Tavily-basierte Websuchen verwendet, um den Abruf aus der Wissensdatenbank zu ergänzen.', tavilyApiKeyMessage: 'Bitte geben Sie Ihren Tavily-API-SCHLÜSSEL ein', - tavilyApiKeyHelp: 'Wie bekomme ich ihn?', + webSearchApiKeyHelp: 'Wie bekomme ich ihn?', crossLanguage: 'Sprachübergreifende Suche', crossLanguageTip: 'Wählen Sie eine oder mehrere Sprachen für die sprachübergreifende Suche aus. Wenn keine Sprache ausgewählt ist, sucht das System mit der ursprünglichen Abfrage.', diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 8abd489686..3570c7ff8a 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -1236,24 +1236,47 @@ This auto-tagging feature enhances retrieval by adding another layer of domain-s selectLanguage: 'Select a language', reasoning: 'Reasoning', reasoningTip: `Whether to enable a reasoning workflow during question answering, as seen in models like Deepseek-R1. When enabled, this allows the model to access external knowledge and tackle complex questions in a step-by-step manner, leveraging techniques like chain-of-thought reasoning. This approach enhances the model's ability to provide accurate responses by breaking down problems into manageable steps, improving performance on tasks that require logical reasoning and multi-step thinking.`, - tavilyApiKeyTip: - 'If an API Key is correctly set here, Tavily-based web searches will be used to supplement dataset retrieval.', - tavilyApiKeyMessage: 'Please enter your Tavily API Key', webSearch: 'Web search', webSearchProvider: 'Web search provider', webSearchProviderTip: 'Select the service used when Internet search is enabled.', webSearchProviderPlaceholder: 'Select a web search provider', + webSearchApiKeyRequired: + 'An API key is required for the selected provider — without one no Internet search happens and the switch never appears.', + // The key field's label. {{provider}} is the provider's BRAND name, which is + // deliberately not translated, so this single template covers all nine. + webSearchApiKeyLabel: '{{provider}} API Key', + // One Tip/Message pair per provider, alphabetical by provider id. + braveApiKeyTip: + 'When Brave Search is selected, its web results supplement dataset retrieval. Every Brave endpoint requires a key.', + braveApiKeyMessage: 'Please enter your Brave Search API Key', + exaApiKeyTip: + 'Required. When Exa is selected, its web results supplement dataset retrieval. A key is needed even on the free tier of 1,000 requests/month.', + exaApiKeyMessage: 'Please enter your Exa API Key', + firecrawlApiKeyTip: + 'When Firecrawl is selected, its search results supplement dataset retrieval. Only the search snippets are pulled, not a full page scrape.', + firecrawlApiKeyMessage: 'Please enter your Firecrawl API Key', + linkupApiKeyTip: + 'When Linkup is selected, its web results supplement dataset retrieval.', + linkupApiKeyMessage: 'Please enter your Linkup API Key', + parallelApiKeyTip: + 'When Parallel is selected, its search excerpts supplement dataset retrieval.', + parallelApiKeyMessage: 'Please enter your Parallel API Key', queritApiKeyTip: 'When Querit is selected, its web search results supplement dataset retrieval.', queritApiKeyMessage: 'Please enter your Querit API Key', serplyApiKeyTip: 'When Serply is selected, its web search results supplement dataset retrieval.', serplyApiKeyMessage: 'Please enter your Serply API Key', + tavilyApiKeyTip: + 'If an API Key is correctly set here, Tavily-based web searches will be used to supplement dataset retrieval.', + tavilyApiKeyMessage: 'Please enter your Tavily API Key', youcomApiKeyTip: 'Optional. You.com works without a key on its rate-limited endpoint; add a key to lift those limits.', youcomApiKeyMessage: 'Optional — leave blank to use the free tier', - tavilyApiKeyHelp: 'How to get it?', + // Shared anchor text for every provider's "get a key" link (only the href is + // provider-specific), so keep it provider-agnostic — do not name a provider. + webSearchApiKeyHelp: 'How to get it?', crossLanguage: 'Cross-language search', crossLanguagePlaceholder: 'Select value', crossLanguageTip: `Select one or more languages for cross‑language search. If no language is selected, the system searches with the original query.`, diff --git a/web/src/locales/fr.ts b/web/src/locales/fr.ts index 7ddaebbd4c..7a18be549b 100644 --- a/web/src/locales/fr.ts +++ b/web/src/locales/fr.ts @@ -1074,7 +1074,7 @@ Applicable lorsque vous avez besoin que le LLM résume le document entier. tavilyApiKeyTip: 'Si une clé API est correctement configurée ici, les recherches web basées sur Tavily seront utilisées pour compléter la récupération des bases de connaissances.', tavilyApiKeyMessage: 'Veuillez saisir votre clé API Tavily', - tavilyApiKeyHelp: "Comment l'obtenir ?", + webSearchApiKeyHelp: "Comment l'obtenir ?", crossLanguage: 'Recherche inter-langues', crossLanguagePlaceholder: 'Sélectionner une valeur', crossLanguageTip: `Sélectionnez une ou plusieurs langues pour la recherche inter-langues. Si aucune langue n'est sélectionnée, le système recherche avec la requête originale.`, diff --git a/web/src/locales/it.ts b/web/src/locales/it.ts index f2e14a040b..e3009fc3ed 100644 --- a/web/src/locales/it.ts +++ b/web/src/locales/it.ts @@ -1147,7 +1147,7 @@ Questa funzionalità di auto-tagging migliora il recupero aggiungendo un ulterio tavilyApiKeyTip: 'Se una chiave API è impostata correttamente qui, le ricerche web basate su Tavily saranno usate per integrare il recupero dal dataset.', tavilyApiKeyMessage: 'Inserisci la tua chiave API Tavily', - tavilyApiKeyHelp: 'Come ottenerla?', + webSearchApiKeyHelp: 'Come ottenerla?', crossLanguage: 'Ricerca cross-lingua', crossLanguagePlaceholder: 'Seleziona valore', crossLanguageTip: `Seleziona una o più lingue per la ricerca cross-lingua. Se non viene selezionata nessuna lingua, il sistema cerca con la query originale.`, diff --git a/web/src/locales/ja.ts b/web/src/locales/ja.ts index a7f499cf57..6653ea86a5 100644 --- a/web/src/locales/ja.ts +++ b/web/src/locales/ja.ts @@ -950,7 +950,7 @@ export default { 'ここにTavilyのAPIキーを設定すると、ナレッジベース検索に加えてウェブ検索も利用できます。', tavilyApiKeyMessage: 'Tavily APIキーを入力してください', - tavilyApiKeyHelp: '取得方法はこちら', + webSearchApiKeyHelp: '取得方法はこちら', crossLanguage: 'クロス言語検索', crossLanguageTip: `1つ以上の言語を選択すると、その言語でも検索します。未選択の場合は元の言語で検索します。`, createChat: 'チャットを作成', diff --git a/web/src/locales/ko.ts b/web/src/locales/ko.ts index adca40707f..5edf8633a9 100644 --- a/web/src/locales/ko.ts +++ b/web/src/locales/ko.ts @@ -1124,7 +1124,7 @@ export default { tavilyApiKeyTip: 'API 키가 올바르게 설정되면 Tavily 기반 웹 검색이 데이터셋 검색을 보완하는 데 사용됩니다.', tavilyApiKeyMessage: 'Tavily API 키를 입력해 주세요', - tavilyApiKeyHelp: '어떻게 얻나요?', + webSearchApiKeyHelp: '어떻게 얻나요?', crossLanguage: '교차 언어 검색', crossLanguagePlaceholder: '값 선택', crossLanguageTip: `교차 언어(cross-language) 검색을 위해 언어를 하나 이상 선택하세요. 언어를 선택하지 않으면 시스템은 원본 쿼리로 검색합니다.`, diff --git a/web/src/locales/ru.ts b/web/src/locales/ru.ts index 4685a6e78c..0dcc23db92 100644 --- a/web/src/locales/ru.ts +++ b/web/src/locales/ru.ts @@ -947,7 +947,7 @@ export default { tavilyApiKeyTip: 'Если здесь правильно установлен API ключ, для дополнения извлечения из базы знаний будут использоваться веб-поиски на основе Tavily.', tavilyApiKeyMessage: 'Введите Tavily API Key', - tavilyApiKeyHelp: 'Как его получить?', + webSearchApiKeyHelp: 'Как его получить?', crossLanguage: 'Кросс-языковой поиск', crossLanguagePlaceholder: 'Выберите значение', crossLanguageTip: `Выберите один или несколько языков для кросс-языкового поиска. Если не выбрано — поиск по исходному запросу.`, diff --git a/web/src/locales/tr.ts b/web/src/locales/tr.ts index cb5fc8b583..52aff1bda3 100644 --- a/web/src/locales/tr.ts +++ b/web/src/locales/tr.ts @@ -1205,25 +1205,46 @@ Bu otomatik etiketleme özelliği, mevcut datasete alanına özgü bilgi katman selectLanguage: 'Dil seçin', reasoning: 'Akıl yürütme', reasoningTip: `Deepseek-R1 gibi modellerde görüldüğü şekilde soru yanıtlama sırasında bir akıl yürütme iş akışı etkinleştirilsin mi. Etkinleştirildiğinde, model harici bilgiye erişebilir ve karmaşık soruları düşünce zinciri akıl yürütmesi gibi tekniklerden yararlanarak adım adım ele alabilir. Bu yaklaşım, problemleri yönetilebilir adımlara bölerek modelin doğru yanıtlar sağlama yeteneğini artırır.`, - tavilyApiKeyTip: - 'Burada bir API anahtarı doğru şekilde ayarlanırsa, Tavily tabanlı web aramaları dataset alımını desteklemek için kullanılacaktır.', - tavilyApiKeyMessage: 'Lütfen Tavily API anahtarınızı girin', webSearch: 'Web araması', webSearchProvider: 'Web arama sağlayıcısı', webSearchProviderTip: 'İnternet araması etkinleştirildiğinde kullanılacak hizmeti seçin.', webSearchProviderPlaceholder: 'Bir web arama sağlayıcısı seçin', + webSearchApiKeyRequired: + 'Seçilen sağlayıcı için bir API anahtarı gerekir — olmadan internet araması yapılmaz ve anahtar hiç görünmez.', + // Anahtar alanının etiketi. {{provider}} sağlayıcının MARKA adıdır ve + // bilinçli olarak çevrilmez; bu tek şablon dokuz sağlayıcıyı da kapsar. + webSearchApiKeyLabel: '{{provider}} API Anahtarı', + // Her sağlayıcı için bir Tip/Mesaj çifti, sağlayıcı id'sine göre alfabetik. + braveApiKeyTip: + 'Brave Search seçildiğinde, web arama sonuçları dataset alımını tamamlar. Tüm Brave uç noktaları anahtar gerektirir.', + braveApiKeyMessage: 'Lütfen Brave Search API anahtarınızı girin', + exaApiKeyTip: + 'Zorunlu. Exa seçildiğinde, web arama sonuçları dataset alımını tamamlar. Aylık 1.000 isteğe kadar olan ücretsiz katmanda bile anahtar gereklidir.', + exaApiKeyMessage: 'Lütfen Exa API anahtarınızı girin', + firecrawlApiKeyTip: + 'Firecrawl seçildiğinde, arama sonuçları dataset alımını tamamlar. Yalnızca arama özetleri alınır, tam sayfa kazınmaz.', + firecrawlApiKeyMessage: 'Lütfen Firecrawl API anahtarınızı girin', + linkupApiKeyTip: + 'Linkup seçildiğinde, web arama sonuçları dataset alımını tamamlar.', + linkupApiKeyMessage: 'Lütfen Linkup API anahtarınızı girin', + parallelApiKeyTip: + 'Parallel seçildiğinde, arama alıntıları dataset alımını tamamlar.', + parallelApiKeyMessage: 'Lütfen Parallel API anahtarınızı girin', queritApiKeyTip: 'Querit seçildiğinde, web arama sonuçları dataset alımını tamamlar.', queritApiKeyMessage: 'Lütfen Querit API Anahtarınızı girin', serplyApiKeyTip: 'Serply seçildiğinde, web arama sonuçları dataset alımını tamamlar.', serplyApiKeyMessage: 'Lütfen Serply API Anahtarınızı girin', + tavilyApiKeyTip: + 'Burada bir API anahtarı doğru şekilde ayarlanırsa, Tavily tabanlı web aramaları dataset alımını desteklemek için kullanılacaktır.', + tavilyApiKeyMessage: 'Lütfen Tavily API anahtarınızı girin', youcomApiKeyTip: 'İsteğe bağlı. You.com, hız sınırlı uç noktasında anahtarsız çalışır; bu sınırları kaldırmak için bir anahtar ekleyin.', youcomApiKeyMessage: 'İsteğe bağlı — ücretsiz katmanı kullanmak için boş bırakın', - tavilyApiKeyHelp: 'Nasıl alınır?', + webSearchApiKeyHelp: 'Nasıl alınır?', crossLanguage: 'Çapraz dil araması', crossLanguagePlaceholder: 'Değer seçin', crossLanguageTip: `Çapraz dil araması için bir veya daha fazla dil seçin. Hiçbir dil seçilmezse sistem orijinal sorguyla arama yapar.`, diff --git a/web/src/locales/zh-traditional.ts b/web/src/locales/zh-traditional.ts index 3b6a6d82fc..448ffe1b14 100644 --- a/web/src/locales/zh-traditional.ts +++ b/web/src/locales/zh-traditional.ts @@ -580,10 +580,41 @@ export default { reasoning: '推理', reasoningTip: '在問答過程中是否啟用推理工作流程，例如Deepseek-R1或OpenAI o1等模型所採用的方式。啟用後，該功能允許模型存取外部知識，並借助思維鏈推理等技術逐步解決複雜問題。通過將問題分解為可處理的步驟，這種方法增強了模型提供準確回答的能力，從而在需要邏輯推理和多步思考的任務上表現更優。', + webSearchProvider: '網路搜尋服務', + webSearchProviderTip: '選擇啟用聯網搜尋時使用的搜尋服務。', + webSearchProviderPlaceholder: '請選擇網路搜尋服務', + webSearchApiKeyRequired: + '所選服務必須填寫 API Key —— 否則不會發起聯網搜尋，聊天框中也不會出現聯網開關。', + // 密鑰輸入框的標籤。{{provider}} 是 provider 的品牌名，故意不翻譯， + // 因此一個範本即可涵蓋全部 9 個 provider。 + webSearchApiKeyLabel: '{{provider}} API Key', tavilyApiKeyTip: '如果 API 金鑰設定正確，它將利用 Tavily 進行網路搜尋作為知識庫的補充。', tavilyApiKeyMessage: '請輸入你的 Tavily API Key', - tavilyApiKeyHelp: '如何獲取？', + // 每個 provider 一組 Tip/Message，按 provider id 字典序排列。 + braveApiKeyTip: + '選擇 Brave Search 後，將使用其搜尋結果補充知識庫檢索。Brave 的所有端點都需要 Key。', + braveApiKeyMessage: '請輸入你的 Brave Search API Key', + exaApiKeyTip: + '必填。選擇 Exa 後，將使用其搜尋結果補充知識庫檢索。即使是每月 1,000 次的免費額度，也仍然需要 Key。', + exaApiKeyMessage: '請輸入你的 Exa API Key', + firecrawlApiKeyTip: + '選擇 Firecrawl 後，將使用其搜尋結果補充知識庫檢索。只取搜尋摘要，不抓取整頁。', + firecrawlApiKeyMessage: '請輸入你的 Firecrawl API Key', + linkupApiKeyTip: '選擇 Linkup 後，將使用其搜尋結果補充知識庫檢索。', + linkupApiKeyMessage: '請輸入你的 Linkup API Key', + parallelApiKeyTip: '選擇 Parallel 後，將使用其搜尋摘錄補充知識庫檢索。', + parallelApiKeyMessage: '請輸入你的 Parallel API Key', + queritApiKeyTip: + '選擇 Querit 後，將使用 Querit 的網路搜尋結果補充知識庫檢索。', + queritApiKeyMessage: '請輸入你的 Querit API Key', + serplyApiKeyTip: + '選擇 Serply 後，將使用 Serply 的網路搜尋結果補充知識庫檢索。', + serplyApiKeyMessage: '請輸入你的 Serply API Key', + youcomApiKeyTip: + '可選。You.com 在限速端點上無需 API Key 即可使用；填寫 Key 可解除限速。', + youcomApiKeyMessage: '可選 —— 留空則使用免費額度', + webSearchApiKeyHelp: '如何獲取？', crossLanguage: '跨語言搜尋', crossLanguageTip: `選擇一種或多種語言進行跨語言搜尋。如果沒有選擇語言，系統將使用原始查詢進行搜尋。 `, showChunkMetadata: '顯示區塊中繼資料', diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index ca57c72b67..125cea2773 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -1124,23 +1124,42 @@ NER：使用 spaCy NER 和基于规则的关键词提取来抽取 Entities 和 R reasoning: '推理', reasoningTip: '在问答过程中是否启用推理工作流，类似 Deepseek-R1 等模型所采用的方式。启用后，该功能允许模型访问外部知识，并借助思维链推理等技术逐步解决复杂问题。通过将问题分解为可处理的步骤，这种方法增强了模型提供准确回答的能力，从而在需要逻辑推理和多步思考的任务上表现更优。', - tavilyApiKeyTip: - '如果 API 密钥设置正确，它将利用 Tavily 进行网络搜索作为知识库的补充。', - tavilyApiKeyMessage: '请输入你的 Tavily API Key', webSearch: '网络搜索', webSearchProvider: '网络搜索服务', webSearchProviderTip: '选择启用联网搜索时使用的搜索服务。', webSearchProviderPlaceholder: '请选择网络搜索服务', + webSearchApiKeyRequired: + '所选服务必须填写 API Key —— 否则不会发起联网搜索，聊天框里也不会出现联网开关。', + // 密钥输入框的标签。{{provider}} 是 provider 的品牌名，故意不翻译， + // 因此一个模板即可覆盖全部 9 个 provider。 + webSearchApiKeyLabel: '{{provider}} API Key', + // 每个 provider 一组 Tip/Message，按 provider id 字典序排列。 + braveApiKeyTip: + '选择 Brave Search 后，将使用其搜索结果补充知识库检索。Brave 的所有端点都需要 Key。', + braveApiKeyMessage: '请输入你的 Brave Search API Key', + exaApiKeyTip: + '必填。选择 Exa 后，将使用其搜索结果补充知识库检索。即使是每月 1,000 次的免费额度，也仍然需要 Key。', + exaApiKeyMessage: '请输入你的 Exa API Key', + firecrawlApiKeyTip: + '选择 Firecrawl 后，将使用其搜索结果补充知识库检索。只取搜索摘要，不抓取整页。', + firecrawlApiKeyMessage: '请输入你的 Firecrawl API Key', + linkupApiKeyTip: '选择 Linkup 后，将使用其搜索结果补充知识库检索。', + linkupApiKeyMessage: '请输入你的 Linkup API Key', + parallelApiKeyTip: '选择 Parallel 后，将使用其搜索摘录补充知识库检索。', + parallelApiKeyMessage: '请输入你的 Parallel API Key', queritApiKeyTip: '选择 Querit 后，将使用 Querit 的网络搜索结果补充知识库检索。', queritApiKeyMessage: '请输入你的 Querit API Key', serplyApiKeyTip: '选择 Serply 后，将使用 Serply 的网络搜索结果补充知识库检索。', serplyApiKeyMessage: '请输入你的 Serply API Key', + tavilyApiKeyTip: + '如果 API 密钥设置正确，它将利用 Tavily 进行网络搜索作为知识库的补充。', + tavilyApiKeyMessage: '请输入你的 Tavily API Key', youcomApiKeyTip: '可选。You.com 在限速端点上无需 API Key 即可使用；填写 Key 可解除限速。', youcomApiKeyMessage: '可选 —— 留空则使用免费额度', - tavilyApiKeyHelp: '如何获取？', + webSearchApiKeyHelp: '如何获取？', crossLanguage: '跨语言搜索', crossLanguagePlaceholder: '请选择', crossLanguageTip: `选择一种或多种语言进行跨语言搜索。如果未选择任何语言，系统将使用原始查询进行搜索。`, @@ -2211,7 +2230,8 @@ NER：使用 spaCy NER 和基于规则的关键词提取来抽取 Entities 和 R includeHeadingContentTip: '启用后，每个分块仅保留标题路径和自身内容，与上级标题紧挨着的内容将作为一个独立的块保留。', chunkTokenCap: '分块 Token 上限', - chunkTokenCapTip:'每个分块允许的最大 Token 数。超过上限的分块将按照句子边界（中英文句号、感叹号、问号和换行符）重新拆分。设为 0 可禁用此上限。', + chunkTokenCapTip: + '每个分块允许的最大 Token 数。超过上限的分块将按照句子边界（中英文句号、感叹号、问号和换行符）重新拆分。设为 0 可禁用此上限。', rootAsHeading: '将首个分块设为全局上下文', rootAsHeadingTip: '将首个切片设为全局标题，以确保整个文档层级结构中拥有一致的上下文信息。该功能尤其适用于首段包含关键信息的简历。', diff --git a/web/src/pages/next-chats/chat-stream/utils.ts b/web/src/pages/next-chats/chat-stream/utils.ts index 42757a3973..f01e3fcc7c 100644 --- a/web/src/pages/next-chats/chat-stream/utils.ts +++ b/web/src/pages/next-chats/chat-stream/utils.ts @@ -18,6 +18,28 @@ export function mergeAnswerChunk( previousAnswer: string, chunk: CompletionChunk, ): string { + // A final chunk carries the COMPLETE answer (citations inserted, references + // attached server-side). It replaces the accumulated ANSWER — but NOT the + // thinking accumulated so far, which is the only record of the ReAct + // trajectory the user watched unfold (the final chunk ships the deliverable + // alone). A naive final is self-contained (`think + decorated answer`), so + // it replaces everything wholesale instead of duplicating the think span. + if (chunk.final === true && chunk.answer) { + if (chunk.answer.includes('')) { + return chunk.answer; + } + const closeTag = ''; + const closeIdx = previousAnswer.lastIndexOf(closeTag); + if (closeIdx !== -1) { + return ( + previousAnswer.slice(0, closeIdx + closeTag.length) + + '\n\n' + + chunk.answer + ); + } + return chunk.answer; + } + const currentAnswer = chunk.final && previousAnswer ? '' : chunk.answer || ''; let nextAnswer: string; @@ -32,7 +54,11 @@ export function mergeAnswerChunk( } if (chunk.end_to_think === true) { - nextAnswer = nextAnswer + ''; + // The closing marker must be followed by a blank line: the answer's first + // line lands right after it, and a heading (`## Candidate Matrix`) glued + // onto '' stops being a heading — markdown only parses `#` at + // line begin. + nextAnswer = nextAnswer + '\n\n'; } return nextAnswer; diff --git a/web/src/pages/next-chats/chat/app-settings/chat-settings.tsx b/web/src/pages/next-chats/chat/app-settings/chat-settings.tsx index fa2207ebcb..24bb7a801c 100644 --- a/web/src/pages/next-chats/chat/app-settings/chat-settings.tsx +++ b/web/src/pages/next-chats/chat/app-settings/chat-settings.tsx @@ -60,6 +60,11 @@ export function ChatSettings({ hasSingleChatBox }: ChatSettingsProps) { const form = useForm({ resolver: zodResolver(formSchema), + // shouldUnregister: false keeps every field mounted in the form state, which + // includes all nine per-provider *_api_key fields. Switching provider therefore + // KEEPS the keys already typed for the others and saves them with the dialog — + // deliberate (a user may switch back), but it is why the payload can carry more + // than one provider key. shouldUnregister: false, mode: 'onChange', defaultValues: { diff --git a/web/src/pages/next-chats/chat/app-settings/use-chat-setting-schema.tsx b/web/src/pages/next-chats/chat/app-settings/use-chat-setting-schema.tsx index c8a09ed782..08eae228be 100644 --- a/web/src/pages/next-chats/chat/app-settings/use-chat-setting-schema.tsx +++ b/web/src/pages/next-chats/chat/app-settings/use-chat-setting-schema.tsx @@ -13,6 +13,7 @@ import { topnSchema } from '@/components/top-n-item'; import { WebSearchProvider } from '@/constants/chat'; import { useTranslate } from '@/hooks/common-hooks'; import { z, ZodIssueCode } from 'zod'; +import { missingWebSearchApiKeyField } from '../web-search-api-key'; import { chatPromptKbIssues } from './validate-chat-prompt'; export function useChatSettingSchema() { @@ -34,15 +35,25 @@ export function useChatSettingSchema() { }), ) .optional(), - tavily_api_key: z.string().optional(), + brave_api_key: z.string().optional(), + exa_api_key: z.string().optional(), + firecrawl_api_key: z.string().optional(), + linkup_api_key: z.string().optional(), + parallel_api_key: z.string().optional(), querit_api_key: z.string().optional(), serply_api_key: z.string().optional(), + tavily_api_key: z.string().optional(), youcom_api_key: z.string().optional(), web_search_provider: z .enum([ - WebSearchProvider.Tavily, + WebSearchProvider.Brave, + WebSearchProvider.Exa, + WebSearchProvider.Firecrawl, + WebSearchProvider.Linkup, + WebSearchProvider.Parallel, WebSearchProvider.Querit, WebSearchProvider.Serply, + WebSearchProvider.Tavily, WebSearchProvider.YouCom, ]) .optional() @@ -84,6 +95,19 @@ export function useChatSettingSchema() { message: issue.message, }); } + + // A keyed provider selected without its key fails SILENTLY at runtime — + // the Internet switch never appears in the chat box — so block the save + // here instead. Keyless providers (You.com) are exempt. The rule lives in + // missingWebSearchApiKeyField so it can be unit-tested without the form. + const missingKeyField = missingWebSearchApiKeyField(value?.prompt_config); + if (missingKeyField) { + ctx.addIssue({ + code: ZodIssueCode.custom, + path: ['prompt_config', missingKeyField], + message: t('webSearchApiKeyRequired'), + }); + } }); return formSchema; diff --git a/web/src/pages/next-chats/chat/use-show-internet.test.ts b/web/src/pages/next-chats/chat/use-show-internet.test.ts deleted file mode 100644 index 1d4e28580c..0000000000 --- a/web/src/pages/next-chats/chat/use-show-internet.test.ts +++ /dev/null @@ -1,145 +0,0 @@ -import { WebSearchProvider } from '@/constants/chat'; -import type { PromptConfig } from '@/interfaces/database/chat'; -import { - getWebSearchApiKey, - getWebSearchProvider, - hasWebSearchProvider, -} from './web-search-api-key'; - -describe('getWebSearchProvider', () => { - it('does not select a provider for a new unconfigured dialog', () => { - expect(getWebSearchProvider({} as PromptConfig)).toBeUndefined(); - }); - - it('selects Tavily for a legacy dialog with a Tavily key', () => { - const promptConfig = { - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchProvider(promptConfig)).toBe(WebSearchProvider.Tavily); - }); -}); - -describe('getWebSearchApiKey', () => { - it('uses Tavily for dialogs saved before provider selection existed', () => { - const promptConfig = { - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBe('tvly-test'); - }); - - it('uses only the selected Querit key', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Querit, - querit_api_key: 'querit-test', - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBe('querit-test'); - }); - - it('uses only the selected Serply key', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Serply, - serply_api_key: 'serply-test', - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBe('serply-test'); - }); - - it('does not fall back to Tavily when Serply is selected without a key', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Serply, - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); - }); - - it('does not fall back to Tavily when Querit is selected without a key', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Querit, - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); - }); - - it('treats a whitespace-only key as unconfigured', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Querit, - querit_api_key: ' ', - } as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBe(''); - }); - - it('does not fall back to Tavily for an unsupported provider', () => { - const promptConfig = { - web_search_provider: 'unsupported', - tavily_api_key: 'tvly-test', - } as unknown as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); - }); - - it('treats a non-string key as unconfigured', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.Querit, - querit_api_key: 123, - } as unknown as PromptConfig; - - expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); - }); -}); - -describe('hasWebSearchProvider', () => { - it('is false for a new unconfigured dialog', () => { - expect(hasWebSearchProvider({} as PromptConfig)).toBe(false); - }); - - it('requires a key for providers that need one', () => { - expect( - hasWebSearchProvider({ - web_search_provider: WebSearchProvider.Querit, - } as PromptConfig), - ).toBe(false); - - expect( - hasWebSearchProvider({ - web_search_provider: WebSearchProvider.Querit, - querit_api_key: 'querit-test', - } as PromptConfig), - ).toBe(true); - }); - - it('is true for keyless You.com with no key configured', () => { - expect( - hasWebSearchProvider({ - web_search_provider: WebSearchProvider.YouCom, - } as PromptConfig), - ).toBe(true); - - expect( - hasWebSearchProvider({ - web_search_provider: WebSearchProvider.YouCom, - youcom_api_key: '', - } as PromptConfig), - ).toBe(true); - }); -}); - -describe('You.com key selection', () => { - it('uses only the selected You.com key', () => { - const promptConfig = { - web_search_provider: WebSearchProvider.YouCom, - youcom_api_key: 'ydc-test', - tavily_api_key: 'tvly-test', - } as PromptConfig; - - expect(getWebSearchProvider(promptConfig)).toBe(WebSearchProvider.YouCom); - expect(getWebSearchApiKey(promptConfig)).toBe('ydc-test'); - }); -}); diff --git a/web/src/pages/next-chats/chat/web-search-api-key.test.ts b/web/src/pages/next-chats/chat/web-search-api-key.test.ts new file mode 100644 index 0000000000..f5efc75053 --- /dev/null +++ b/web/src/pages/next-chats/chat/web-search-api-key.test.ts @@ -0,0 +1,348 @@ +import { WebSearchProvider } from '@/constants/chat'; +import type { PromptConfig } from '@/interfaces/database/chat'; +import { + getWebSearchApiKey, + getWebSearchApiKeyField, + getWebSearchProvider, + hasWebSearchProvider, + isWebSearchApiKeyRequired, + missingWebSearchApiKeyField, +} from './web-search-api-key'; + +describe('getWebSearchProvider', () => { + it('does not select a provider for a new unconfigured dialog', () => { + expect(getWebSearchProvider({} as PromptConfig)).toBeUndefined(); + }); + + it('selects Tavily for a legacy dialog with a Tavily key', () => { + const promptConfig = { + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchProvider(promptConfig)).toBe(WebSearchProvider.Tavily); + }); +}); + +describe('getWebSearchApiKey', () => { + it('uses Tavily for dialogs saved before provider selection existed', () => { + const promptConfig = { + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBe('tvly-test'); + }); + + it('uses only the selected Querit key', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Querit, + querit_api_key: 'querit-test', + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBe('querit-test'); + }); + + it('uses only the selected Serply key', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Serply, + serply_api_key: 'serply-test', + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBe('serply-test'); + }); + + it('does not fall back to Tavily when Serply is selected without a key', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Serply, + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); + }); + + it('does not fall back to Tavily when Querit is selected without a key', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Querit, + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); + }); + + it('treats a whitespace-only key as unconfigured', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Querit, + querit_api_key: ' ', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBe(''); + }); + + it('does not fall back to Tavily for an unsupported provider', () => { + const promptConfig = { + web_search_provider: 'unsupported', + tavily_api_key: 'tvly-test', + } as unknown as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); + }); + + it('treats a non-string key as unconfigured', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.Querit, + querit_api_key: 123, + } as unknown as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); + }); +}); + +describe('hasWebSearchProvider', () => { + it('is false for a new unconfigured dialog', () => { + expect(hasWebSearchProvider({} as PromptConfig)).toBe(false); + }); + + it('requires a key for providers that need one', () => { + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.Querit, + } as PromptConfig), + ).toBe(false); + + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.Querit, + querit_api_key: 'querit-test', + } as PromptConfig), + ).toBe(true); + }); + + it('is true for keyless You.com with no key configured', () => { + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.YouCom, + } as PromptConfig), + ).toBe(true); + + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.YouCom, + youcom_api_key: '', + } as PromptConfig), + ).toBe(true); + }); +}); + +describe('You.com key selection', () => { + it('uses only the selected You.com key', () => { + const promptConfig = { + web_search_provider: WebSearchProvider.YouCom, + youcom_api_key: 'ydc-test', + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchProvider(promptConfig)).toBe(WebSearchProvider.YouCom); + expect(getWebSearchApiKey(promptConfig)).toBe('ydc-test'); + }); +}); + +// Brave, Exa, Firecrawl, Linkup and Parallel all authenticate with a key and +// have no keyless path, so none of them may show up as "usable" without one. +// Exa has a free tier of 1,000 requests/month, but it is not keyless. +describe('keyed providers added in 2026-09', () => { + const cases = [ + { + provider: WebSearchProvider.Brave, + keyName: 'brave_api_key', + key: 'brave-test', + }, + { + provider: WebSearchProvider.Exa, + keyName: 'exa_api_key', + key: 'exa-test', + }, + { + provider: WebSearchProvider.Firecrawl, + keyName: 'firecrawl_api_key', + key: 'firecrawl-test', + }, + { + provider: WebSearchProvider.Linkup, + keyName: 'linkup_api_key', + key: 'linkup-test', + }, + { + provider: WebSearchProvider.Parallel, + keyName: 'parallel_api_key', + key: 'parallel-test', + }, + ] as const; + + it.each(cases)( + 'reads $provider from $keyName', + ({ provider, keyName, key }) => { + const promptConfig = { + web_search_provider: provider, + [keyName]: ` ${key} `, + tavily_api_key: 'tvly-test', + } as unknown as PromptConfig; + + expect(getWebSearchProvider(promptConfig)).toBe(provider); + expect(getWebSearchApiKey(promptConfig)).toBe(key); + expect(hasWebSearchProvider(promptConfig)).toBe(true); + }, + ); + + it.each(cases)('is unusable when $provider has no key', ({ provider }) => { + const promptConfig = { + web_search_provider: provider, + tavily_api_key: 'tvly-test', + } as PromptConfig; + + expect(getWebSearchApiKey(promptConfig)).toBeUndefined(); + expect(hasWebSearchProvider(promptConfig)).toBe(false); + }); +}); + +// The required marker and the save-time validation both read the same two +// helpers, so the field name and the required flag have to agree for every +// provider — a mismatch would mark a field required and then validate a +// different one. +describe('provider key field mapping', () => { + it('maps every provider to its own prompt_config key', () => { + expect(getWebSearchApiKeyField(WebSearchProvider.Brave)).toBe( + 'brave_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Exa)).toBe('exa_api_key'); + expect(getWebSearchApiKeyField(WebSearchProvider.Firecrawl)).toBe( + 'firecrawl_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Linkup)).toBe( + 'linkup_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Parallel)).toBe( + 'parallel_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Querit)).toBe( + 'querit_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Serply)).toBe( + 'serply_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.Tavily)).toBe( + 'tavily_api_key', + ); + expect(getWebSearchApiKeyField(WebSearchProvider.YouCom)).toBe( + 'youcom_api_key', + ); + }); + + it('has no key field when no provider is selected', () => { + expect(getWebSearchApiKeyField(undefined)).toBeUndefined(); + expect( + getWebSearchApiKeyField('' as unknown as WebSearchProvider), + ).toBeUndefined(); + }); + + it('requires a key for every provider except the keyless ones', () => { + expect(isWebSearchApiKeyRequired(WebSearchProvider.Brave)).toBe(true); + expect(isWebSearchApiKeyRequired(WebSearchProvider.Exa)).toBe(true); + expect(isWebSearchApiKeyRequired(WebSearchProvider.Tavily)).toBe(true); + expect(isWebSearchApiKeyRequired(WebSearchProvider.YouCom)).toBe(false); + expect(isWebSearchApiKeyRequired(undefined)).toBe(false); + }); +}); + +// Exa is NOT keyless despite its free tier: 1,000 requests/month comes with no +// credit card, but every request still carries a key, so a blank field must read +// as "not configured" exactly like the other keyed providers. +describe('Exa key requirement', () => { + it('is unusable with no key configured', () => { + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.Exa, + } as PromptConfig), + ).toBe(false); + + expect( + hasWebSearchProvider({ + web_search_provider: WebSearchProvider.Exa, + exa_api_key: ' ', + } as PromptConfig), + ).toBe(false); + }); +}); + +// Clearing the provider emits '' from SelectWithSearch (the schema accepts it), and +// undefined means the same thing. The field lookup and the required marker have to +// agree on that, or the form marks a field required while validating nothing. +describe('cleared provider value', () => { + const cleared = '' as unknown as WebSearchProvider; + + it('is not a provider for either helper', () => { + expect(getWebSearchApiKeyField(cleared)).toBeUndefined(); + expect(isWebSearchApiKeyRequired(cleared)).toBe(false); + }); +}); + +// The save-time rule. A keyed provider without a key fails SILENTLY at runtime — +// the Internet switch never appears — so the schema rejects the save; the rule +// lives in this pure helper so it can be tested without rendering the form. +describe('missingWebSearchApiKeyField', () => { + it('names the field when a keyed provider has no key', () => { + expect( + missingWebSearchApiKeyField({ + web_search_provider: WebSearchProvider.Firecrawl, + } as PromptConfig), + ).toBe('firecrawl_api_key'); + }); + + it('names the field when the key is only whitespace', () => { + expect( + missingWebSearchApiKeyField({ + web_search_provider: WebSearchProvider.Linkup, + linkup_api_key: ' ', + } as PromptConfig), + ).toBe('linkup_api_key'); + }); + + it('accepts a keyed provider once its key is set', () => { + expect( + missingWebSearchApiKeyField({ + web_search_provider: WebSearchProvider.Parallel, + parallel_api_key: 'parallel-test', + } as PromptConfig), + ).toBeUndefined(); + }); + + it('exempts the keyless provider', () => { + expect( + missingWebSearchApiKeyField({ + web_search_provider: WebSearchProvider.YouCom, + } as PromptConfig), + ).toBeUndefined(); + }); + + it('accepts an unconfigured new dialog', () => { + expect(missingWebSearchApiKeyField({} as PromptConfig)).toBeUndefined(); + expect(missingWebSearchApiKeyField(undefined)).toBeUndefined(); + }); + + it('accepts a legacy dialog that only carries a Tavily key', () => { + expect( + missingWebSearchApiKeyField({ + tavily_api_key: 'tvly-test', + } as PromptConfig), + ).toBeUndefined(); + }); + + it('accepts an unknown provider instead of guessing a field', () => { + expect( + missingWebSearchApiKeyField({ + web_search_provider: 'unsupported', + tavily_api_key: 'tvly-test', + } as unknown as PromptConfig), + ).toBeUndefined(); + }); +}); diff --git a/web/src/pages/next-chats/chat/web-search-api-key.ts b/web/src/pages/next-chats/chat/web-search-api-key.ts index b45752a0fb..1d648f68a9 100644 --- a/web/src/pages/next-chats/chat/web-search-api-key.ts +++ b/web/src/pages/next-chats/chat/web-search-api-key.ts @@ -2,20 +2,41 @@ import { KEYLESS_WEB_SEARCH_PROVIDERS, WebSearchProvider, } from '@/constants/chat'; -import type { PromptConfig } from '@/interfaces/database/chat'; -export function getWebSearchProvider(promptConfig?: PromptConfig) { +/** + * The subset of a dialog's prompt_config these helpers read — the web-search + * provider and its per-provider key slots. + * + * Structural rather than PromptConfig because the save-time form schema allows + * web_search_provider to be '' (a cleared select), which PromptConfig does not; + * with this shape the schema can call the helpers without casting. + */ +export type WebSearchKeyConfig = { + web_search_provider?: WebSearchProvider | ''; + tavily_api_key?: string; +}; + +const WEB_SEARCH_PROVIDERS: WebSearchProvider[] = + Object.values(WebSearchProvider); + +export function getWebSearchProvider(promptConfig?: WebSearchKeyConfig) { const provider = promptConfig?.web_search_provider; + // Derived from the enum instead of a hand-written union: a provider added to the + // enum — and therefore to the dropdown — used to be silently treated as + // unconfigured here until someone remembered to extend the list too, which + // disabled its Internet switch with no error. if ( - provider === WebSearchProvider.Tavily || - provider === WebSearchProvider.Querit || - provider === WebSearchProvider.Serply || - provider === WebSearchProvider.YouCom + provider !== undefined && + provider !== '' && + WEB_SEARCH_PROVIDERS.includes(provider) ) { return provider; } + // undefined means the dialog predates provider selection, which is where the + // legacy Tavily key comes in. '' is an explicit "no provider" (cleared select) + // and must NOT fall back to Tavily. if ( provider === undefined && typeof promptConfig?.tavily_api_key === 'string' && @@ -27,35 +48,90 @@ export function getWebSearchProvider(promptConfig?: PromptConfig) { return undefined; } -export function getWebSearchApiKey(promptConfig?: PromptConfig) { - const provider = getWebSearchProvider(promptConfig); - let apiKey: unknown; +// The prompt_config field each provider reads its key from. Kept as the single +// source of truth: the form uses it for the required marker, the schema uses it +// for validation, and the reader below uses it to fetch the value — three +// callers that used to each spell the mapping out again. +const webSearchApiKeyFields: Record = { + [WebSearchProvider.Brave]: 'brave_api_key', + [WebSearchProvider.Exa]: 'exa_api_key', + [WebSearchProvider.Firecrawl]: 'firecrawl_api_key', + [WebSearchProvider.Linkup]: 'linkup_api_key', + [WebSearchProvider.Parallel]: 'parallel_api_key', + [WebSearchProvider.Querit]: 'querit_api_key', + [WebSearchProvider.Serply]: 'serply_api_key', + [WebSearchProvider.Tavily]: 'tavily_api_key', + [WebSearchProvider.YouCom]: 'youcom_api_key', +}; - switch (provider) { - case WebSearchProvider.Tavily: - apiKey = promptConfig?.tavily_api_key; - break; - case WebSearchProvider.Querit: - apiKey = promptConfig?.querit_api_key; - break; - case WebSearchProvider.Serply: - apiKey = promptConfig?.serply_api_key; - break; - case WebSearchProvider.YouCom: - apiKey = promptConfig?.youcom_api_key; - break; - default: - return undefined; +// A cleared SelectWithSearch emits '' rather than undefined (the schema accepts +// both), so normalise it here: the field lookup and the required marker have to +// agree on what "no provider selected" means. +function normalizeProvider(provider?: WebSearchProvider | '') { + return provider === '' ? undefined : provider; +} + +// Reads a prompt_config field the app names at runtime (the per-provider key +// slots), which no static type can enumerate. +function fieldValue( + promptConfig: WebSearchKeyConfig | undefined, + field: string, +) { + return (promptConfig as unknown as Record | undefined)?.[ + field + ]; +} + +export function getWebSearchApiKeyField(provider?: WebSearchProvider | '') { + const normalized = normalizeProvider(provider); + return normalized ? webSearchApiKeyFields[normalized] : undefined; +} + +// Whether the selected provider must have a key before it can be used. A +// keyless provider answers on its own endpoint/tier and stays usable blank. +export function isWebSearchApiKeyRequired(provider?: WebSearchProvider | '') { + const normalized = normalizeProvider(provider); + return ( + normalized !== undefined && + !KEYLESS_WEB_SEARCH_PROVIDERS.includes(normalized) + ); +} + +export function getWebSearchApiKey(promptConfig?: WebSearchKeyConfig) { + const keyField = getWebSearchApiKeyField(getWebSearchProvider(promptConfig)); + if (!keyField) { + return undefined; } + const apiKey = fieldValue(promptConfig, keyField); return typeof apiKey === 'string' ? apiKey.trim() : undefined; } +/** + * Returns the prompt_config field of the selected provider's API key when the + * provider requires one and it is missing or blank — the reason a save has to be + * rejected. A keyed provider without a key fails SILENTLY at runtime (the Internet + * switch never appears), which is why the form blocks it up front. undefined means + * the configuration is acceptable. + * + * Pure and exported so the save-time rule is testable without rendering the form. + */ +export function missingWebSearchApiKeyField(promptConfig?: WebSearchKeyConfig) { + const provider = getWebSearchProvider(promptConfig); + const keyField = getWebSearchApiKeyField(provider); + if (!keyField || !isWebSearchApiKeyRequired(provider)) { + return undefined; + } + + const apiKey = fieldValue(promptConfig, keyField); + return typeof apiKey === 'string' && apiKey.trim() ? undefined : keyField; +} + /** * Whether web search is usable as configured. Most providers need a key; a * keyless provider is usable as soon as it is selected. */ -export function hasWebSearchProvider(promptConfig?: PromptConfig) { +export function hasWebSearchProvider(promptConfig?: WebSearchKeyConfig) { const provider = getWebSearchProvider(promptConfig); if (provider === undefined) {
elements. The +// alternative - matching GFM exactly - is to keep losing the text, and the indexed text is +// what this parser exists to preserve. +func balanceTableColumns(tableLines []string) []string { + if len(tableLines) < 2 { + return tableLines + } + cells := func(line string) int { + s := strings.TrimRight(line, "\n") + pipes, escaped := 0, false + for i := 0; i < len(s); i++ { + switch { + case escaped: + escaped = false + case s[i] == '\\': + escaped = true + case s[i] == '|': + pipes++ + } + } + trimmed := strings.TrimSpace(s) + if strings.HasPrefix(trimmed, "|") { + pipes-- + } + if strings.HasSuffix(trimmed, "|") { + pipes-- + } + if pipes < 0 { + return 0 + } + return pipes + 1 + } + + head := cells(tableLines[0]) + widest := head + for _, line := range tableLines[1:] { + if n := cells(line); n > widest { + widest = n + } + } + if widest <= head { + return tableLines + } + add := widest - head + + out := make([]string, len(tableLines)) + copy(out, tableLines) + pad := func(line string, separator bool) string { + nl := "" + if strings.HasSuffix(line, "\n") { + nl, line = "\n", strings.TrimRight(line, "\n") + } + cell := " |" + if separator { + cell = "---|" + } + line = strings.TrimRight(line, " \t") + if !strings.HasSuffix(line, "|") { + line += "|" + } + return line + strings.Repeat(cell, add) + nl + } + out[0] = pad(out[0], false) + out[1] = pad(out[1], true) + return out +} + // renderMarkdownTablesInlineText renders every GFM/HTML table inline as an // HTML block and returns the rewritten text. When no table is present the // input is returned unchanged. Unlike renderMarkdownTablesInline it always @@ -296,9 +392,21 @@ func walkMarkdownBlocksWithImages(doc ast.Node, out *[]map[string]any, flatten, ckType = "text" docTypeKwd = "text" case *ast.List: - txt = leafText(n) - ckType = "list" - docTypeKwd = "text" + // Python's _markdown keeps one item per list entry, marker included + // ("- Eat a light diet before the exam?"), so a list never collapses into one blob. + // Emit the same shape: the reconstructed marker plus the entry's own text, one item + // each. The boundaries matter beyond alignment - the marker lives in the list node + // rather than in any leaf literal, and concatenating entries glued their words + // ("ChinaChina"), which destroys the word boundaries the tokenizer and the retrieval + // index depend on. + for _, entry := range listEntries(n) { + *out = append(*out, map[string]any{ + "text": entry, + "doc_type_kwd": "text", + "ck_type": "list", + }) + } + continue case *ast.CodeBlock: txt = leafText(n) ckType = "code" @@ -571,6 +679,59 @@ func leafText(n ast.Node) string { return strings.TrimSpace(buf.String()) } +// listEntries renders a list's entries the way the source spells them: the marker first, then the +// entry text, one string per entry (Python's _markdown shape, see the en golden). The marker has to +// be reconstructed because the AST keeps it on the list node - BulletChar for bullet lists, +// Start/Delimiter for ordered ones - rather than in any leaf literal. +func listEntries(list *ast.List) []string { + entries := make([]string, 0, len(list.GetChildren())) + ordered := list.ListFlags&ast.ListTypeOrdered != 0 + for i, item := range list.GetChildren() { + text := strings.TrimSpace(leafText(item)) + if text == "" { + continue + } + var marker string + switch { + case ordered: + start := list.Start + if start == 0 { + start = 1 + } + delim := list.Delimiter + if delim == 0 { + delim = '.' + } + marker = fmt.Sprintf("%d%c ", start+i, delim) + default: + // The markdown parser records the bullet on the ITEM (parser/block.go sets + // ListItem.BulletChar), not on the list, so read it there first. + bullet := list.BulletChar + if li, ok := item.(*ast.ListItem); ok && li.BulletChar != 0 { + bullet = li.BulletChar + } + if bullet != 0 { + marker = string(bullet) + " " + } + } + entries = append(entries, marker+text) + } + return entries +} + +// writeListItems emits a list's items separated by a newline. Without the separator the item texts +// run together ("bullet onebullet two"), because the marker lives in the list structure rather than +// in the item's leaf text - the same defect that merged table rows showed as "ChinaChina" in the +// indexed chunks. +func writeListItems(list *ast.List, buf *bytes.Buffer) { + for i, item := range list.GetChildren() { + if i > 0 { + buf.WriteByte('\n') + } + walkLeaf(item, buf) + } +} + func walkLeaf(n ast.Node, buf *bytes.Buffer) { switch t := n.(type) { case *ast.Text: @@ -588,6 +749,43 @@ func walkLeaf(n ast.Node, buf *bytes.Buffer) { // markup in Content, the markdown-block path in Literal. Emit both. buf.Write(t.Literal) buf.Write(t.Content) + case *ast.Math: + // Inline math (`$...$`) keeps its payload on the node literal, not in a + // child Text node, so an unhandled ast.Math silently drops the whole + // span. Emit the literal so the text can never be lost, even if the + // MathJax extension is re-enabled (markdownNew masks it off because + // `$` is currency in real documents far more often than it is math). + buf.Write(t.Literal) + case *ast.MathBlock: + // Display math (`$$...$$`) stores its payload the same way. + buf.Write(t.Literal) + case *ast.HTMLSpan: + // Inline raw HTML (a that is glued to surrounding prose with no + // blank line around it, /,
, …) is an HTMLSpan: the + // markup AND the text inside it live on the node literal, with no child + // Text nodes. Dropping it deleted whole regions of a document - a 664k + // character CDC report lost 397k characters across 14 such spans. Emit + // the literal (and Content for the paths that use it) so the text + // survives verbatim, exactly as the source file spells it. + buf.Write(t.Literal) + buf.Write(t.Content) + case *ast.List: + // A list marker is structure, not a leaf: the item's own Text literal + // carries neither the bullet nor a trailing newline, so concatenating the + // items glued the last word of one item to the first word of the next + // ("ChinaChina", "HathiTrustThe"). That is not a cosmetic difference: it + // destroys the word boundaries the tokenizer and the retrieval index rely + // on, so a query for "China" can no longer match the merged token. Python's + // _markdown keeps the source lines, so a newline between items restores + // parity and keeps the words separate. + writeListItems(t, buf) + case *ast.ListItem: + for i, c := range t.GetChildren() { + if i > 0 { + buf.WriteByte('\n') + } + walkLeaf(c, buf) + } default: for _, c := range n.GetChildren() { walkLeaf(c, buf) diff --git a/internal/parser/parser/markdown_parser_test.go b/internal/parser/parser/markdown_parser_test.go index 8616b49e48..610c312c79 100644 --- a/internal/parser/parser/markdown_parser_test.go +++ b/internal/parser/parser/markdown_parser_test.go @@ -35,6 +35,146 @@ func TestMarkdownParser_ParseWithResult_Basic(t *testing.T) { } } +// TestMarkdownParser_ListItemsKeepWordBoundaries pins the shape of a list in the parsed output: +// one item per entry, marker included, and no gluing between entries. +// +// A list marker is structure, not a leaf: the marker lives on the list node (bullets on the item) +// and the entry's own text carries no marker and no trailing newline. Emitting the list as a single +// blob therefore glued the last word of one entry to the first word of the next ("GamesChina"), +// which destroys the word boundaries the tokenizer and the retrieval index depend on - a query for +// "China" can no longer match the merged token. Python's _markdown keeps one item per entry with the +// marker, so this pins the same shape. +func TestMarkdownParser_ListItemsKeepWordBoundaries(t *testing.T) { + ctx := t.Context() + p, err := NewMarkdownParser(GoMarkdown) + if err != nil { + t.Fatalf("NewMarkdownParser: %v", err) + } + md := "See also\n\n* China at the Asian Games\n* China at the Paralympics\n" + res := p.ParseWithResult(ctx, "list.md", []byte(md)) + if res.Err != nil { + t.Fatalf("ParseWithResult: %v", res.Err) + } + var lists []string + for _, item := range res.JSON { + if ck, _ := item["ck_type"].(string); ck == "list" { + text, _ := item["text"].(string) + lists = append(lists, text) + } + } + if len(lists) != 2 { + t.Fatalf("want one item per list entry, got %d: %#v", len(lists), lists) + } + for i, want := range []string{"* China at the Asian Games", "* China at the Paralympics"} { + if lists[i] != want { + t.Errorf("list item %d = %q, want %q", i, lists[i], want) + } + if strings.Contains(lists[i], "GamesChina") { + t.Errorf("list entries are glued together: %q", lists[i]) + } + } +} + +// TestMarkdownParser_OrderedListKeepsNumbers covers the ordered branch of the marker +// reconstruction: the number and its delimiter live on the list node, so they have to be rebuilt +// from Start/Delimiter instead of being read off a leaf. +// +// Two of the three cases start above 1 on purpose. A list that starts at 1 cannot tell "kept the +// numbers the source wrote" from "renumbered from 1", and this test used to run nothing but +// "1."/"2." - it passed while the parser renumbered every ordered list in a real corpus (a +// 2,244-entry list whose source numbers start at 996 was indexed as 1..N, which removes the source +// numbers from the index and made them count as unserved words). +func TestMarkdownParser_OrderedListKeepsNumbers(t *testing.T) { + ctx := t.Context() + p, err := NewMarkdownParser(GoMarkdown) + if err != nil { + t.Fatalf("NewMarkdownParser: %v", err) + } + cases := []struct { + name string + md string + want []string + }{ + { + name: "starts at ten", + md: "10. a\n11. b\n", + want: []string{"10. a", "11. b"}, + }, + { + name: "starts in the thousands", + md: "996. first entry\n\n997. second entry\n\n998. third entry\n", + want: []string{"996. first entry", "997. second entry", "998. third entry"}, + }, + { + name: "starts at one", + md: "1. first item\n2. second item\n", + want: []string{"1. first item", "2. second item"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + res := p.ParseWithResult(ctx, "ordered.md", []byte(tc.md)) + if res.Err != nil { + t.Fatalf("ParseWithResult: %v", res.Err) + } + var lists []string + for _, item := range res.JSON { + if ck, _ := item["ck_type"].(string); ck == "list" { + text, _ := item["text"].(string) + lists = append(lists, text) + } + } + if len(lists) != len(tc.want) { + t.Fatalf("want %d items, got %d: %#v", len(tc.want), len(lists), lists) + } + for i := range tc.want { + if lists[i] != tc.want[i] { + t.Errorf("ordered item %d = %q, want %q", i, lists[i], tc.want[i]) + } + } + }) + } +} + +// TestMarkdownParser_TableKeepsCellsBeyondTheHeader covers the deliberate deviation from GFM. +// +// A row may legitimately carry more cells than its header - an unescaped '|' inside a cell, +// e.g. a Wikipedia image spec "150x150px|alt=..." - and GFM ignores every cell beyond the +// header's column count, so the tail of such a row (the Year and Description columns of a +// World Heritage list) never reaches the index. Measured on the production index that cost +// 514 documents and 18,185 words no query could reach; the header is padded to the widest row +// instead. Fails on the old behaviour: the sentinel below is dropped. +func TestMarkdownParser_TableKeepsCellsBeyondTheHeader(t *testing.T) { + ctx := t.Context() + p, err := NewMarkdownParser(GoMarkdown) + if err != nil { + t.Fatalf("NewMarkdownParser: %v", err) + } + const sentinel = "DESCRIPTIONSENTINEL" + md := "| Site | Image | Year |\n" + + "|---|---|---|\n" + + "| Aachen Cathedral | 150x150px|alt=A Gothic building | 1978 | " + sentinel + " died in 814 |\n" + res := p.ParseWithResult(ctx, "table.md", []byte(md)) + if res.Err != nil { + t.Fatalf("ParseWithResult: %v", res.Err) + } + var joined strings.Builder + for _, item := range res.JSON { + text, _ := item["text"].(string) + joined.WriteString(text) + joined.WriteString("\n") + } + out := joined.String() + // The unescaped '|' inside the image cell becomes a cell boundary of its own, so the + // assertions are on cell contents: everything the row said has to be somewhere in the + // output, which is exactly what the column-count padding buys. + for _, want := range []string{"Aachen Cathedral", "150x150px", "alt=A Gothic building", "1978", sentinel} { + if !strings.Contains(out, want) { + t.Errorf("table cell content %q is missing from the parsed output:\n%s", want, out) + } + } +} + func TestMarkdownParser_ParseWithResult_EmptyInput(t *testing.T) { ctx := t.Context() p, _ := NewMarkdownParser(GoMarkdown) @@ -137,6 +277,97 @@ func TestMarkdownParser_ParseWithResult_RendersTableInline(t *testing.T) { } } +// TestMarkdownParser_ParseWithResult_DollarIsCurrencyNotMath guards the +// regression that made a 13.8k-character billing report parse into 1.4k +// characters: CommonExtensions enables MathJax, so `$...$` spans became +// ast.Math nodes whose payload is not a Text child, and the text walker +// dropped them. With two far-apart `$` that deletes everything in between. +// A `$` in an ingested document is far more often a price than it is math, so +// the text must survive verbatim, dollar signs included. +func TestMarkdownParser_ParseWithResult_DollarIsCurrencyNotMath(t *testing.T) { + ctx := t.Context() + p, _ := NewMarkdownParser(GoMarkdown) + md := "# Report\n\nFunding was provided by a $4,800 grant from the Sanctuary.\n\n" + + "Costs: 224.50 theodolite, 48.70 binoculars, $3.75 internet.\n\n" + + "Remaining funds for Winter 2002: $193.92\n" + res := p.ParseWithResult(ctx, "report.md", []byte(md)) + if res.Err != nil { + t.Fatalf("ParseWithResult: %v", res.Err) + } + var sb strings.Builder + for _, item := range res.JSON { + text, _ := item["text"].(string) + sb.WriteString(text) + sb.WriteString("\n") + } + got := sb.String() + // Every sentence of the source has to survive; the price table sentence is + // the one sitting between two '$' and used to be deleted whole. + for _, want := range []string{ + "Funding was provided by a $4,800 grant from the Sanctuary.", + "Costs: 224.50 theodolite, 48.70 binoculars, $3.75 internet.", + "Remaining funds for Winter 2002: $193.92", + } { + if !strings.Contains(got, want) { + t.Errorf("parsed text lost %q\n--- got ---\n%s", want, got) + } + } + // `$$...$$` display math must stay literal text as well. + mdDisplay := "Totals are shown as $$4,606.17 in the audit.\n" + resDisplay := p.ParseWithResult(ctx, "totals.md", []byte(mdDisplay)) + if resDisplay.Err != nil { + t.Fatalf("ParseWithResult: %v", resDisplay.Err) + } + var sbDisplay strings.Builder + for _, item := range resDisplay.JSON { + text, _ := item["text"].(string) + sbDisplay.WriteString(text) + } + if !strings.Contains(sbDisplay.String(), "4,606.17 in the audit") { + t.Errorf("display math span was dropped: %q", sbDisplay.String()) + } +} + +// TestMarkdownParser_ParseWithResult_InlineRawHTMLSurvives guards the second +// text-loss mechanism found in the same walker: inline raw HTML is an +// ast.HTMLSpan whose payload (markup and the text inside it) lives on the node +// literal, with no child Text nodes. A 664k-character report lost 397k +// characters across 14 such spans. +func TestMarkdownParser_ParseWithResult_InlineRawHTMLSurvives(t *testing.T) { + ctx := t.Context() + p, _ := NewMarkdownParser(GoMarkdown) + // No blank line around the
, so gomarkdown keeps it inline inside + // the paragraph (an HTMLSpan) instead of a block-level HTMLBlock. + md := "Lead-in sentence before the table. " + + "
" + + "
theodolite224.50
binoculars48.70
" + + "Trailing sentence after the table.\n" + res := p.ParseWithResult(ctx, "report.md", []byte(md)) + if res.Err != nil { + t.Fatalf("ParseWithResult: %v", res.Err) + } + var sb strings.Builder + for _, item := range res.JSON { + text, _ := item["text"].(string) + sb.WriteString(text) + sb.WriteString("\n") + } + got := sb.String() + for _, want := range []string{ + "Lead-in sentence before the table.", + "Trailing sentence after the table.", + } { + if !strings.Contains(got, want) { + t.Errorf("parsed text lost %q\n--- got ---\n%s", want, got) + } + } + // The inline HTML payload must not vanish either: the walker emits the raw + // markup, exactly as it stands in the source. + if !strings.Contains(got, "binoculars") { + t.Errorf("inline raw HTML payload was dropped: %q", got) + } +} + func TestMarkdownParser_ConfigureFromSetup(t *testing.T) { p, _ := NewMarkdownParser(GoMarkdown) p.ConfigureFromSetup(map[string]any{ diff --git a/internal/rag/agentic-rag/agentic_rag.go b/internal/rag/agentic-rag/agentic_rag.go index e455238087..b569de72c8 100644 --- a/internal/rag/agentic-rag/agentic_rag.go +++ b/internal/rag/agentic-rag/agentic_rag.go @@ -231,13 +231,13 @@ type RAGTools struct { // the runtime defaults. TopN int SimilarityThreshold float64 - // VectorSimilarityWeight is the vector leg's weight. A pointer so an explicit 0 — - // keyword-only, which is what the agentic retrieve channel uses — is distinguishable - // from "not configured". Nil selects DefaultAgenticVectorWeight (0). - VectorSimilarityWeight *float64 - // UsingEmbedding: when false (the agentic default) retrieval is keyword-only (vector - // weight 0); when true the embedder is engaged and VectorSimilarityWeight (default 0.7) - // applies. + // KeywordsSimilarityWeight is the keyword leg's weight. A pointer keeps an + // explicit zero (pure vector search) distinct from an unset value. + KeywordsSimilarityWeight *float64 + // UsingEmbedding is the Go spelling of Python RAGTools.retrieve's + // `using_embedding: bool = False` (agentic_rag.py:599). When false (the + // agentic default) retrieval is keyword-only; when true the embedder is + // engaged and KeywordsSimilarityWeight (default 0.3) applies. // // SCOPE: only the retrieve channel honours it. UsingEmbedding bool @@ -974,18 +974,18 @@ func searchDepsFor(ctx context.Context, deps RAGTools, req runtime.RunRequest, d DocTenantResolver: dbDocTenantResolver{}, Model: deps.Model, // the calculate tool writes its expression via the model DocScope: deps.DocScope, - // Configuration is the middle precedence level, between an explicit tool argument - // and the module defaults. - TopN: deps.TopN, - SimilarityThreshold: deps.SimilarityThreshold, - VectorSimilarityWeight: deps.VectorSimilarityWeight, - UsingEmbedding: deps.UsingEmbedding, - HasEmbedder: deps.HasEmbedder, - RerankCandidatesCount: deps.RerankCandidatesCount, - TopK: deps.TopK, - MetaDataFilter: deps.MetaDataFilter, - // RAGTools carries the KB objects and a tagger, so the tag boost can be - // computed. + // Python retrieve:614-646 — configuration is the middle precedence + // level, between an explicit tool argument and the module defaults. + TopN: deps.TopN, + SimilarityThreshold: deps.SimilarityThreshold, + KeywordsSimilarityWeight: deps.KeywordsSimilarityWeight, + UsingEmbedding: deps.UsingEmbedding, + HasEmbedder: deps.HasEmbedder, + RerankCandidatesCount: deps.RerankCandidatesCount, + TopK: deps.TopK, + MetaDataFilter: deps.MetaDataFilter, + // rank_feature (Python retrieve:668): RAGTools carries the KB objects + // and a tagger, mirroring rank_feature=label_question(question, self.kbs). KBs: deps.KBs, Tagger: deps.Tagger, // External embedding handle. When the caller supplied one on RAGTools it is used diff --git a/internal/rag/agentic-rag/agentic_rag_graph_test.go b/internal/rag/agentic-rag/agentic_rag_graph_test.go index c34ebcb1cc..5336d222a8 100644 --- a/internal/rag/agentic-rag/agentic_rag_graph_test.go +++ b/internal/rag/agentic-rag/agentic_rag_graph_test.go @@ -15,14 +15,15 @@ import ( "testing" "time" - "github.com/cloudwego/eino/compose" - "github.com/cloudwego/eino/schema" "ragflow/internal/engine" "ragflow/internal/engine/types" "ragflow/internal/entity" "ragflow/internal/rag/agentic-rag/runtime" "ragflow/internal/rag/agentic-rag/slots" "ragflow/internal/rag/prompts" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" ) // Test doubles @@ -745,19 +746,19 @@ func TestRunSlotResearchPassLogsSlotEvidenceBound(t *testing.T) { // leg is keyword-only (weight 0, top_n=60, query terms folded into the query), // the semantic-bypass leg has the vector weight on (top_n=30, plain query). type channelRetriever struct { - weights []float64 - topN []int - queries []string + keywordsWeights []float64 + topN []int + queries []string } func (r *channelRetriever) Retrieve(_ context.Context, req runtime.RetrieveRequest) ([]map[string]any, error) { - // Each leg names its weight explicitly; nil (the control was not supplied) + // Each leg names its keyword weight explicitly; nil (the control was not supplied) // is recorded as -1 so the assertions below catch a leg that drops it. weight := -1.0 - if req.VectorSimilarityWeight != nil { - weight = *req.VectorSimilarityWeight + if req.KeywordsSimilarityWeight != nil { + weight = *req.KeywordsSimilarityWeight } - r.weights = append(r.weights, weight) + r.keywordsWeights = append(r.keywordsWeights, weight) r.topN = append(r.topN, req.TopN) r.queries = append(r.queries, req.Query) return nil, nil @@ -777,13 +778,13 @@ func TestFanoutSearchIsDualChannel(t *testing.T) { FanoutSearch(ctx, deps, st, []string{"when was it built", "where located"}, 8, 60) // Two fan-outs x two channels. - if len(r.weights) != 4 { - t.Fatalf("expected 4 retrieve calls (2 fan-outs x 2 channels), got %d", len(r.weights)) + if len(r.keywordsWeights) != 4 { + t.Fatalf("expected 4 retrieve calls (2 fan-outs x 2 channels), got %d", len(r.keywordsWeights)) } for i := 0; i < 4; i += 2 { // Channel A: keyword-only weight, wide pool, query terms folded in. - if r.weights[i] != 0 { - t.Errorf("fan-out %d channel A: weight = %v, want 0 (BM25 keyword leg)", i/2, r.weights[i]) + if r.keywordsWeights[i] != 1 { + t.Errorf("fan-out %d channel A: keyword weight = %v, want 1 (BM25 keyword leg)", i/2, r.keywordsWeights[i]) } if r.topN[i] != fanoutBM25TopN { t.Errorf("fan-out %d channel A: TopN = %d, want %d", i/2, r.topN[i], fanoutBM25TopN) @@ -792,8 +793,8 @@ func TestFanoutSearchIsDualChannel(t *testing.T) { t.Errorf("fan-out %d channel A: query = %q, want the fan-out query (terms folded in)", i/2, r.queries[i]) } // Channel B: vector weight on, top_n=30 (narrowing bypassed). - if r.weights[i+1] <= 0 { - t.Errorf("fan-out %d channel B: weight = %v, want > 0 (semantic bypass leg)", i/2, r.weights[i+1]) + if r.keywordsWeights[i+1] >= 1 { + t.Errorf("fan-out %d channel B: keyword weight = %v, want < 1 (semantic bypass leg)", i/2, r.keywordsWeights[i+1]) } if r.topN[i+1] != fanoutHybridTopN { t.Errorf("fan-out %d channel B: TopN = %d, want %d", i/2, r.topN[i+1], fanoutHybridTopN) diff --git a/internal/rag/agentic-rag/runtime/tool_executor.go b/internal/rag/agentic-rag/runtime/tool_executor.go index 72fa740230..dee33a27be 100644 --- a/internal/rag/agentic-rag/runtime/tool_executor.go +++ b/internal/rag/agentic-rag/runtime/tool_executor.go @@ -1684,13 +1684,13 @@ func (r *RuntimeRetriever) Retrieve(ctx context.Context, req RetrieveRequest) ([ // stay nil so the retrieval service keeps its own default, while a zero // is a real override. Taking the address of a zero value here forced // "threshold 0 / full vector weight" onto every caller that omitted them. - SimilarityThreshold: req.SimilarityThreshold, - VectorSimilarityWeight: req.VectorSimilarityWeight, - DisableVectorLeg: req.DisableVectorLeg, - TenantID: req.TenantID, - RankFeature: req.RankFeature, - // ExcludeCompiled maps onto the runtime request's OnlyOriginalText (the - // "no compile_kwd" exclusion). + SimilarityThreshold: req.SimilarityThreshold, + KeywordsSimilarityWeight: req.KeywordsSimilarityWeight, + TenantID: req.TenantID, + RankFeature: &req.RankFeature, + // ExcludeCompiled maps Python hybrid_search's + // must_not={"exists":"compile_kwd"} onto the runtime request's + // OnlyOriginalText (the "no compile_kwd" exclusion). OnlyOriginalText: req.ExcludeCompiled, }) if err != nil { diff --git a/internal/rag/agentic-rag/runtime/tool_executor_test.go b/internal/rag/agentic-rag/runtime/tool_executor_test.go index 18272bf1fe..93fac48129 100644 --- a/internal/rag/agentic-rag/runtime/tool_executor_test.go +++ b/internal/rag/agentic-rag/runtime/tool_executor_test.go @@ -25,9 +25,10 @@ import ( "testing" "unicode/utf8" - "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/engine" + + "gorm.io/gorm" ) // stubRetrievalService returns a fixed set of child chunks so the runtime @@ -75,24 +76,24 @@ func TestRuntimeRetrieverPreservesUnsetControls(t *testing.T) { if _, err := r.Retrieve(context.Background(), RetrieveRequest{Query: "q", DatasetIDs: []string{"kb-1"}}); err != nil { t.Fatalf("Retrieve: %v", err) } - if got.SimilarityThreshold != nil || got.VectorSimilarityWeight != nil { + if got.SimilarityThreshold != nil || got.KeywordsSimilarityWeight != nil { t.Errorf("omitted controls = %v / %v, want nil so the service keeps its defaults", - got.SimilarityThreshold, got.VectorSimilarityWeight) + got.SimilarityThreshold, got.KeywordsSimilarityWeight) } - threshold, weight := 0.35, 0.3 + threshold, keywordsWeight := 0.35, 0.7 if _, err := r.Retrieve(context.Background(), RetrieveRequest{ - Query: "q", - SimilarityThreshold: &threshold, - VectorSimilarityWeight: &weight, + Query: "q", + SimilarityThreshold: &threshold, + KeywordsSimilarityWeight: &keywordsWeight, }); err != nil { t.Fatalf("Retrieve: %v", err) } if got.SimilarityThreshold == nil || *got.SimilarityThreshold != threshold { t.Errorf("SimilarityThreshold = %v, want %v", got.SimilarityThreshold, threshold) } - if got.VectorSimilarityWeight == nil || *got.VectorSimilarityWeight != weight { - t.Errorf("VectorSimilarityWeight = %v, want %v", got.VectorSimilarityWeight, weight) + if got.KeywordsSimilarityWeight == nil || *got.KeywordsSimilarityWeight != keywordsWeight { + t.Errorf("KeywordsSimilarityWeight = %v, want %v", got.KeywordsSimilarityWeight, keywordsWeight) } } @@ -110,9 +111,9 @@ func TestChunkAggRetrieveLeavesControlsUnset(t *testing.T) { t.Fatalf("chunks = %d, want the retrieved chunk", len(chunks)) } req := r.lastReq(t) - if req.SimilarityThreshold != nil || req.VectorSimilarityWeight != nil { + if req.SimilarityThreshold != nil || req.KeywordsSimilarityWeight != nil { t.Errorf("controls = %v / %v, want nil (zero is a valid value, not an unset marker)", - req.SimilarityThreshold, req.VectorSimilarityWeight) + req.SimilarityThreshold, req.KeywordsSimilarityWeight) } } @@ -473,14 +474,15 @@ func TestExecuteNarratesToolOutcome(t *testing.T) { t.Errorf("search-leg step missing %q; legs:\n%s", want, legs) } } - // The log keeps its own terse form: the same verb without the method clause and - // without a period. - if !strings.Contains(out, `[Hybrid search] Searching for "q"`) { + // The log keeps Python's own wording (search.py:122): it names the corpus, not + // the method, and ends without a period. + if !strings.Contains(out, `[Hybrid search] Searching the knowledge base for "q"`) { t.Errorf("the log must keep its searching line:\n%s", out) } // ...and the log form must not be what the block got: the two sentences differ - // by design (the log has no method clause and no period). - if strings.Contains(legs, `Searching for "q"`) { + // by design (the log names the corpus, the block names the method, and only the + // block ends in a period). + if strings.Contains(legs, `Searching the knowledge base for`) { t.Errorf("the think block got the log line instead of the sentence:\n%s", legs) } // The per-document breakdown the sentence is derived from stays in the diff --git a/internal/rag/agentic-rag/runtime/tool_search.go b/internal/rag/agentic-rag/runtime/tool_search.go index aa47669a4d..41438d5b37 100644 --- a/internal/rag/agentic-rag/runtime/tool_search.go +++ b/internal/rag/agentic-rag/runtime/tool_search.go @@ -150,8 +150,9 @@ type DocTenantResolver interface { // RetrieveRequest is one retrieval call. // -// Weight is the vector-similarity weight: 0.3 hybrid (default), 1.0 vector-only, -// 0.0 keyword-only (BM25). That is how the three search entry points below differ. +// Weight is the keyword-similarity weight: 0.7 hybrid (default), 0.0 vector-only, +// 1.0 keyword-only (BM25). This is how Python's three search entry points +// (hybrid_search / vector_search / bm25_search) differ. type RetrieveRequest struct { Query string DatasetIDs []string @@ -163,22 +164,10 @@ type RetrieveRequest struct { // (threshold 0 = no floor); nil means "not supplied", so the retriever keeps // its own default instead of being handed a zero override. SimilarityThreshold *float64 - // VectorSimilarityWeight carries the VECTOR leg's weight directly: 0.3 hybrid, 1.0 - // vector-only, 0.0 keyword-only. - // It is forwarded verbatim to the retrieval backend — NO inversion. The - // canvas-facing adapter field of the same name family - // (agent/tool RetrievalRequest.KeywordsSimilarityWeight) means the KEYWORD - // weight and is inverted there; conflating the two silently turned the - // agentic hybrid leg vector-dominant (0.3 -> 0.7) and the BM25 legs into - // pure-vector searches. - VectorSimilarityWeight *float64 - // DisableVectorLeg: embd_mdl=None: bm25_search / - // grep_search and RAGTools.retrieve with - // using_embedding=False (agentic_rag.py:retrieve) run NO dense leg at all — - // not even a weight-0 one, which would still constrain the candidate pool - // through the KNN similarity option. - DisableVectorLeg bool - TenantID string + // KeywordsSimilarityWeight is the keyword leg's weight. The retrieval + // adapter derives the complementary vector weight and pure-mode behavior. + KeywordsSimilarityWeight *float64 + TenantID string // MetaDataFilter restricts retrieval to chunks whose metadata matches // Nil means no filtering. MetaDataFilter map[string]any @@ -245,29 +234,31 @@ const ( ChannelRetrieve ) -// resolveVectorWeight is the weight each channel computes. The gate differs per channel -// — that is the whole point: +// resolveKeywordsSimilarityWeight mirrors the weight each Python entry point computes. The +// gate differs per channel — that is the whole point: // -// - ChannelGrep: always 0 (the grep leg has no vector leg). -// - ChannelHybrid: 0 when no embedder is available, otherwise VectorSimilarityWeight -// ?? 0.3. UsingEmbedding is IRRELEVANT here: this channel has no such parameter. -// - ChannelRetrieve: 0 unless UsingEmbedding, otherwise VectorSimilarityWeight ?? 0.7. -func resolveVectorWeight(deps SearchDeps, ch SearchChannel) float64 { +// - ChannelGrep: always 1 (grep_search has no vector leg). +// - ChannelHybrid: 1 when no embedder is configured (Python's `if embd_mdl`), +// otherwise KeywordsSimilarityWeight ?? 0.7. UsingEmbedding is IRRELEVANT +// here: Python's hybrid_search has no such parameter. +// - ChannelRetrieve: 1 unless UsingEmbedding (Python's using_embedding), +// otherwise KeywordsSimilarityWeight ?? 0.3. +func resolveKeywordsSimilarityWeight(deps SearchDeps, ch SearchChannel) float64 { switch ch { case ChannelGrep: - return 0 + return 1 case ChannelHybrid: // The configured weight when an embedder is available, else 0. HasEmbedder says // only that one is available — the retrieval service resolves the actual model. if !deps.HasEmbedder { - return 0 + return 1 } - return floatPtrOrDef(deps.VectorSimilarityWeight, HybridSearchDefaultVectorWeight) + return floatPtrOrDef(deps.KeywordsSimilarityWeight, 1-HybridSearchDefaultVectorWeight) default: if !deps.UsingEmbedding { - return 0 + return 1 } - return floatPtrOrDef(deps.VectorSimilarityWeight, DefaultHybridVectorWeight) + return floatPtrOrDef(deps.KeywordsSimilarityWeight, 1-DefaultHybridVectorWeight) } } @@ -334,20 +325,24 @@ type SearchDeps struct { // database and therefore stays self-contained for callers that cannot or do not want // to supply one. Embedder nlp.NavEmbedder - // Retrieval tuning: top_n, similarity threshold, vector similarity weight, rerank - // candidates and top_k. Zero falls back to this package's Default* constants. + // Retrieval tuning (Python RAGTools.retrieve: _setting(self, "top_n"), + // similarity_threshold, keywords_similarity_weight, rerank_candidates_count, + // top_k). Zero falls back to this package's Default* constants. TopN int SimilarityThreshold float64 - // VectorSimilarityWeight is a pointer so an explicit 0 (keyword-only, the - // agentic default) is distinguishable from "not configured", which falls - // back to DefaultAgenticVectorWeight. - VectorSimilarityWeight *float64 - // UsingEmbedding turns the embedded leg of the RETRIEVE channel on. When false (the - // agentic default) the vector leg is disabled and retrieval is keyword-only (vector - // weight 0). When true the embedder is engaged and the vector weight is applied: - // VectorSimilarityWeight if set, else DefaultHybridVectorWeight. The query embedder - // itself is supplied by the runtime retrieval service, so this carries the flag - // rather than a handle. + // KeywordsSimilarityWeight is a pointer so an explicit 0 (vector-only) is + // distinguishable from "not configured", which falls back to the channel's + // default keyword weight. + KeywordsSimilarityWeight *float64 + // UsingEmbedding is the Go spelling of Python RAGTools.retrieve's + // `using_embedding: bool = False` (agentic_rag.py:retrieve). When false (the + // agentic default), the vector leg is disabled and retrieval is keyword-only + // (vector weight 0) — exactly Python's `embd_mdl = None; vector_weight = 0`. + // When true, the embedder is engaged and the vector weight is applied: + // KeywordsSimilarityWeight if set, else 0.3 (the complement of Python's + // default vector weight of 0.7). Unlike Python, the + // query embedder is supplied by the runtime retrieval service, so Go carries + // the flag rather than an embd_mdl handle on the retrieve call itself. // // SCOPE: this flag governs ChannelRetrieve ONLY. The hybrid channel (search_chunks) // has no such parameter, so gating it here is what silently disabled the semantic @@ -492,8 +487,8 @@ func FanoutKeyedTerms(terms []string) []string { // similarity threshold and compiled-row exclusion are chosen by the CALLER, never by a // single shared "channel flag" that a global switch could silently corrupt. type searchOpts struct { - // weight is the VECTOR similarity weight. - weight float64 + // keywordsSimilarityWeight is the keyword leg's similarity weight. + keywordsSimilarityWeight float64 // threshold is the engine similarity floor. threshold float64 // excludeCompiled: must_not={"exists":"compile_kwd"}. @@ -507,19 +502,18 @@ type searchOpts struct { // "searching" line — the LOG form, terse and grep-friendly, with the keywords // suffix and no period: // - // hybrid → [Hybrid search] Searching for "q" (keywords: k) + // hybrid → [Hybrid search] Searching the knowledge base for "q" (keywords: k) // vector → [Vector search] Searching by meaning for "q" (keywords: k) // bm25 → [BM25 search] Searching by keyword for "q" (keywords: k) // grep → [Grep search] Keyword-first locate for "q" (NO keywords) // retrieve → [Retrieve] Searching for "q" (keywords: k) (Go-only tag) // - // The verbs started out Python-exact (search.py:122/:215/:252/:418), and the - // vector/bm25/grep ones still are. The hybrid and retrieve legs dropped "the - // knowledge base" from it (2026-09-16): every leg searches it, so the phrase - // distinguishes nothing a reader can act on, and the leg's own tag - // ("[BM25 search]") already names the tool. The shape a Python log diff is - // compared against is otherwise unchanged — same tag, same keywords suffix, same - // `"%q -> N chunk(s): …"` result line. + // Every verb is Python's own (search.py:122/:216/:254/:421), so a log diff + // against Python lines up — retrieve is the one Go-only tag, and it no longer + // carries "the knowledge base" (2026-09-16), which every leg searches and which + // therefore distinguishes nothing a reader can act on. The shape is otherwise + // unchanged — same tag, same keywords suffix, same `"%q -> N chunk(s): …"` + // result line. // // The think block gets thinkVerb through searchThinkLine instead — one sentence // family for every leg, with the method named and the keyword list left out. The @@ -542,7 +536,7 @@ type searchOpts struct { // SENTENCE: it names the method ("by meaning and keyword") because the // "[Hybrid search]" prefix is a developer's label, and it is the reason the two // projections differ at all. Every leg sets it explicitly — so the log side can - // be reworded (the hybrid and retrieve legs dropped "the knowledge base", above) + // be reworded (the Go-only retrieve leg dropped "the knowledge base", above) // without the trace moving with it; the fallback to logVerb only covers a future // leg that forgets to set one, where an identical sentence beats an empty one. thinkVerb string @@ -557,11 +551,6 @@ type searchOpts struct { // leg passes the snake_case "hybrid_search", while the vector and bm25 legs pass their // display labels "Vector search" / "BM25 search". Both forms are reproduced verbatim. narrowLabel string - // disableVector: embd_mdl=None: when set the request - // carries DisableVectorLeg and the backend runs the keyword-only branch. - // Grep (bm25_search) always sets it; hybrid/retrieve set it from their own - // embedder gates (HasEmbedder / UsingEmbedding). - disableVector bool // rankFeature opts the entry point into the question-type tag boost // (rank_feature). ONLY the retrieve channel passes the question-type tags; the three // search legs call the retriever WITHOUT them (the grep leg delegates to bm25). Before @@ -661,11 +650,12 @@ func runSearch(ctx context.Context, deps SearchDeps, p SearchParams, opts search } // The per-leg searching line: BOTH the bracket label and the verb vary by entry - // point — "[Hybrid search] Searching for", "[Vector search] Searching by meaning - // for", "[BM25 search] Searching by keyword for", "[Grep search] - // Keyword-first locate for" — and only the three keyword-carrying legs - // append the keywords (see searchOpts.logLabel / logVerb / logKeywords for which - // verbs are Python's own and which dropped "the knowledge base"). + // point — "[Hybrid search] Searching the knowledge base for", "[Vector search] + // Searching by meaning for", "[BM25 search] Searching by keyword for", "[Grep + // search] Keyword-first locate for" — and only the three keyword-carrying legs + // append the keywords (see searchOpts.logLabel / logVerb / logKeywords for the + // Python source of each verb, and for the one Go-only tag that dropped "the + // knowledge base"). // The suffix is printed only when non-empty: a keyword-less leg is a // deliberate state (fan-out Channel B is a narrow bypass, // agentic_rag_graph.py:_search_one; search_chunks takes no keywords at all, @@ -681,7 +671,7 @@ func runSearch(ctx context.Context, deps SearchDeps, p SearchParams, opts search // // The two audiences get different sentences from this one call site. The log // keeps Python's shape — label, keyword list and result line, plus Python's verb - // on the vector/bm25/grep legs (the hybrid and retrieve verbs deliberately dropped + // on every leg that has one (only the Go-only retrieve verb deliberately dropped // "the knowledge base"; see logVerb above for what a log diff can still be // compared against). The think block gets the sentence family every leg // shares, which names the method inside the sentence because "[BM25 search]" is @@ -730,22 +720,18 @@ func runSearch(ctx context.Context, deps SearchDeps, p SearchParams, opts search // cached, nothing narrowed yet). retrieve := func() ([]map[string]any, error) { return deps.Backend.Retrieve(ctx, RetrieveRequest{ - Query: effectiveQuery, - DatasetIDs: targetIDs, - DocScope: docScope, - TopN: topN, - TopK: intOrDef(deps.TopK, DefaultTopK), - RerankCandidatesCount: max(intOrDef(deps.RerankCandidatesCount, DefaultRerankCandidatesCount), topN), - SimilarityThreshold: &opts.threshold, - // VectorSimilarityWeight carries the VECTOR weight verbatim; the canvas - // adapter's keyword-weight inversion does NOT apply to this field. - // DisableVectorLeg is the hard switch: no dense leg at all. - VectorSimilarityWeight: &opts.weight, - DisableVectorLeg: opts.disableVector, - TenantID: deps.TenantID, - MetaDataFilter: deps.MetaDataFilter, - RankFeature: rankFeature, - ExcludeCompiled: opts.excludeCompiled, + Query: effectiveQuery, + DatasetIDs: targetIDs, + DocScope: docScope, + TopN: topN, + TopK: intOrDef(deps.TopK, DefaultTopK), + RerankCandidatesCount: max(intOrDef(deps.RerankCandidatesCount, DefaultRerankCandidatesCount), topN), + SimilarityThreshold: &opts.threshold, + KeywordsSimilarityWeight: &opts.keywordsSimilarityWeight, + TenantID: deps.TenantID, + MetaDataFilter: deps.MetaDataFilter, + RankFeature: rankFeature, + ExcludeCompiled: opts.excludeCompiled, }) } chunks, err := retrieve() @@ -810,18 +796,15 @@ func runSearch(ctx context.Context, deps SearchDeps, p SearchParams, opts search // weight when one is, else 0); default weight 0.3. Compiled rows are excluded. func HybridSearch(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[string]any, []map[string]any) { return runSearch(ctx, deps, p, searchOpts{ - weight: resolveVectorWeight(deps, ChannelHybrid), - threshold: floatOrDef(deps.SimilarityThreshold, DefaultSimilarityThreshold), - excludeCompiled: true, - promoteChildren: true, - // The real embedder is passed even at weight 0; the dense leg is dropped only when - // none is configured. - disableVector: !deps.HasEmbedder, - logLabel: "Hybrid search", - logVerb: "Searching for", - logKeywords: true, - thinkVerb: "Searching by meaning and keyword for", - cache: true, // Python's search_cache is read/written here only (:136/:207) + keywordsSimilarityWeight: resolveKeywordsSimilarityWeight(deps, ChannelHybrid), + threshold: floatOrDef(deps.SimilarityThreshold, DefaultSimilarityThreshold), + excludeCompiled: true, + promoteChildren: true, + logLabel: "Hybrid search", + logVerb: "Searching the knowledge base for", + logKeywords: true, + cache: true, // Python's search_cache is read/written here only (:136/:207) + thinkVerb: "Searching by meaning and keyword for", // Python passes the SNAKE_CASE tag to _narrow_or_keep on this leg // (search.py:hybrid_search) although its own lines say "Hybrid search" — the // inconsistency is Python's and is reproduced verbatim. @@ -836,15 +819,15 @@ func VectorSearch(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[s return nil, nil } return runSearch(ctx, deps, p, searchOpts{ - weight: 1.0, - threshold: VectorSearchDefaultSimilarityThreshold, - excludeCompiled: true, - promoteChildren: true, - logLabel: "Vector search", - logVerb: "Searching by meaning for", - logKeywords: true, - thinkVerb: "Searching by meaning for", - narrowLabel: "Vector search", + keywordsSimilarityWeight: 0, + threshold: VectorSearchDefaultSimilarityThreshold, + excludeCompiled: true, + promoteChildren: true, + logLabel: "Vector search", + logVerb: "Searching by meaning for", + logKeywords: true, + thinkVerb: "Searching by meaning for", + narrowLabel: "Vector search", }) } @@ -853,17 +836,15 @@ func VectorSearch(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[s // rows are excluded. func BM25Search(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[string]any, []map[string]any) { return runSearch(ctx, deps, p, searchOpts{ - weight: 0, - threshold: BM25SearchDefaultSimilarityThreshold, - excludeCompiled: true, - promoteChildren: true, - // Keyword only, NO dense leg at all — not even a weight-0 one. - disableVector: true, - logLabel: "BM25 search", - logVerb: "Searching by keyword for", - logKeywords: true, - thinkVerb: "Searching by keyword for", - narrowLabel: "BM25 search", + keywordsSimilarityWeight: 1, + threshold: BM25SearchDefaultSimilarityThreshold, + excludeCompiled: true, + promoteChildren: true, + logLabel: "BM25 search", + logVerb: "Searching by keyword for", + logKeywords: true, + thinkVerb: "Searching by keyword for", + narrowLabel: "BM25 search", }) } @@ -1260,14 +1241,14 @@ func GrepSearch(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[str // UsingEmbedding and does NOT exclude compiled rows. func RetrieveSearch(ctx context.Context, deps SearchDeps, p SearchParams) ([]map[string]any, []map[string]any) { return runSearch(ctx, deps, p, searchOpts{ - weight: resolveVectorWeight(deps, ChannelRetrieve), - threshold: floatOrDef(deps.SimilarityThreshold, DefaultSimilarityThreshold), - excludeCompiled: false, - promoteChildren: true, - // No dense leg without the flag. - disableVector: !deps.UsingEmbedding, - // This pass has its own identity line: the tag stays human-readable and the extra - // lines stay on, because this is the only trace the low-mode direct pass leaves. + keywordsSimilarityWeight: resolveKeywordsSimilarityWeight(deps, ChannelRetrieve), + threshold: floatOrDef(deps.SimilarityThreshold, DefaultSimilarityThreshold), + excludeCompiled: false, + promoteChildren: true, + // Go-only identity: Python's L1 RAGTools.retrieve + // logs nothing of its own, so there is no Python string to mirror. The + // tag stays human-readable and the extra lines stay on, because this is + // the only trace the low-mode direct pass leaves. logLabel: "Retrieve", logVerb: "Searching for", logKeywords: true, diff --git a/internal/rag/agentic-rag/runtime/tool_search_test.go b/internal/rag/agentic-rag/runtime/tool_search_test.go index c97203fc70..25e6ce382f 100644 --- a/internal/rag/agentic-rag/runtime/tool_search_test.go +++ b/internal/rag/agentic-rag/runtime/tool_search_test.go @@ -569,35 +569,31 @@ func TestVectorSearchBailsWithoutEmbedder(t *testing.T) { } } -// TestVectorSearchWeightIsOne: the pure-vector leg -// carries weight 1.0 and excludes compiled rows. -func TestVectorSearchWeightIsOne(t *testing.T) { +// TestVectorSearchUsesZeroKeywordWeight mirrors Python vector_search: the +// pure-vector leg carries keyword weight 0 and excludes compiled rows. +func TestVectorSearchUsesZeroKeywordWeight(t *testing.T) { r := &stubRetriever{chunks: []map[string]any{{"content": "x"}}} deps, _ := newTestSearchDeps(r) deps.HasEmbedder = true VectorSearch(context.Background(), deps, SearchParams{Question: "q"}) req := r.lastReq(t) - if got := ptrFloat(t, req.VectorSimilarityWeight); got != 1.0 { - t.Errorf("vector search weight = %v, want 1.0", got) + if got := ptrFloat(t, req.KeywordsSimilarityWeight); got != 0 { + t.Errorf("vector search keyword weight = %v, want 0", got) } if !req.ExcludeCompiled { t.Error("vector search must exclude compiled rows") } } -// TestBM25SearchUsesZeroWeight: keyword-only, vector -// weight unconditionally 0, threshold 0.0, excludes compiled rows. -func TestBM25SearchUsesZeroWeight(t *testing.T) { +// TestBM25SearchUsesFullKeywordWeight mirrors Python bm25_search: keyword-only, +// keyword weight 1, threshold 0.0, excludes compiled rows. +func TestBM25SearchUsesFullKeywordWeight(t *testing.T) { r := &stubRetriever{chunks: []map[string]any{{"content": "x"}}} deps, _ := newTestSearchDeps(r) BM25Search(context.Background(), deps, SearchParams{Question: "q"}) req := r.lastReq(t) - if got := ptrFloat(t, req.VectorSimilarityWeight); got != 0 { - t.Errorf("bm25 search weight = %v, want 0", got) - } - // No dense leg at all. - if !req.DisableVectorLeg { - t.Error("bm25 search must disable the vector leg") + if got := ptrFloat(t, req.KeywordsSimilarityWeight); got != 1 { + t.Errorf("bm25 search keyword weight = %v, want 1", got) } if got := ptrFloat(t, req.SimilarityThreshold); got != 0 { t.Errorf("bm25 search threshold = %v, want 0", got) @@ -607,18 +603,15 @@ func TestBM25SearchUsesZeroWeight(t *testing.T) { } } -// TestGrepSearchDelegatesToBM25: it is bm25_search -// with a keyword-only (weight 0) leg and compiled-row exclusion. +// TestGrepSearchDelegatesToBM25 mirrors Python grep_search: it is bm25_search +// with a keyword-only (weight 1) leg and compiled-row exclusion. func TestGrepSearchDelegatesToBM25(t *testing.T) { r := &stubRetriever{chunks: []map[string]any{{"content": "x"}}} deps, _ := newTestSearchDeps(r) GrepSearch(context.Background(), deps, SearchParams{Question: "q", Keywords: "kw"}) req := r.lastReq(t) - if got := ptrFloat(t, req.VectorSimilarityWeight); got != 0 { - t.Errorf("grep search weight = %v, want 0", got) - } - if !req.DisableVectorLeg { - t.Error("grep search must disable the vector leg (embd_mdl=None)") + if got := ptrFloat(t, req.KeywordsSimilarityWeight); got != 1 { + t.Errorf("grep search keyword weight = %v, want 1", got) } if !req.ExcludeCompiled { t.Error("grep search must exclude compiled rows") @@ -925,8 +918,8 @@ func TestGrepSearchDerivesKeywordsHint(t *testing.T) { if want := "who made Culdcept? who made Culdcept"; req.Query != want { t.Errorf("query = %q, want %q (question + derived terms)", req.Query, want) } - if ptrFloat(t, req.VectorSimilarityWeight) != 0 || !req.ExcludeCompiled { - t.Error("grep must stay keyword-only (weight 0) and exclude compiled rows") + if ptrFloat(t, req.KeywordsSimilarityWeight) != 1 || !req.ExcludeCompiled { + t.Error("grep must stay keyword-only (weight 1) and exclude compiled rows") } // The derived hint also drives the narrowing stage: a prose candidate whose // sentences miss those terms is dropped, while a >=3-row pipe table is kept @@ -970,21 +963,17 @@ func TestSearchCacheIsHybridOnly(t *testing.T) { } } -// TestHybridSearchExcludesCompiledAndWeightsThreeTenths: vector weight 0.3 when an -// embedder is configured, compiled rows excluded. -func TestHybridSearchExcludesCompiledAndWeightsThreeTenths(t *testing.T) { +// TestHybridSearchExcludesCompiledAndUsesSevenTenthsKeywordWeight mirrors Python +// hybrid_search: keyword weight 0.7 when an embedder is configured, compiled +// rows excluded. +func TestHybridSearchExcludesCompiledAndUsesSevenTenthsKeywordWeight(t *testing.T) { r := &stubRetriever{chunks: []map[string]any{{"content": "x"}}} deps, _ := newTestSearchDeps(r) deps.HasEmbedder = true HybridSearch(context.Background(), deps, SearchParams{Question: "q"}) req := r.lastReq(t) - if got := ptrFloat(t, req.VectorSimilarityWeight); got != HybridSearchDefaultVectorWeight { - t.Errorf("hybrid search weight = %v, want %v", got, HybridSearchDefaultVectorWeight) - } - // With an embedder configured the real handle is passed: the dense leg RUNS at - // weight 0.3. - if req.DisableVectorLeg { - t.Error("hybrid search with an embedder must keep the vector leg") + if got := ptrFloat(t, req.KeywordsSimilarityWeight); got != 1-HybridSearchDefaultVectorWeight { + t.Errorf("hybrid search keyword weight = %v, want %v", got, 1-HybridSearchDefaultVectorWeight) } if !req.ExcludeCompiled { t.Error("hybrid search must exclude compiled rows") @@ -1003,8 +992,8 @@ func TestRetrieveSearchDoesNotExcludeCompiled(t *testing.T) { if req.ExcludeCompiled { t.Error("retrieve search must NOT exclude compiled rows") } - if got := ptrFloat(t, req.VectorSimilarityWeight); got != DefaultHybridVectorWeight { - t.Errorf("retrieve search weight = %v, want %v", got, DefaultHybridVectorWeight) + if got := ptrFloat(t, req.KeywordsSimilarityWeight); got != 1-DefaultHybridVectorWeight { + t.Errorf("retrieve search keyword weight = %v, want %v", got, 1-DefaultHybridVectorWeight) } } @@ -1172,52 +1161,55 @@ func TestRetrievalDefaultsUseIntOrDef(t *testing.T) { } } -func TestResolveVectorWeightRetrieveMirrorsUsingEmbedding(t *testing.T) { - // UsingEmbedding off → keyword-only (weight 0); on → 0.7 default or the configured - // override. - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: false}, ChannelRetrieve); got != 0 { - t.Fatalf("using_embedding=false → %v, want 0 (keyword-only)", got) +func TestResolveKeywordsSimilarityWeightRetrieveMirrorsUsingEmbedding(t *testing.T) { + // Python RAGTools.retrieve(using_embedding: bool = False). + // Off → keyword-only (weight 1); on → 0.3 default or the configured override. + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: false}, ChannelRetrieve); got != 1 { + t.Fatalf("using_embedding=false → %v, want 1 (keyword-only)", got) } - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: true}, ChannelRetrieve); got != DefaultHybridVectorWeight { - t.Fatalf("using_embedding=true → %v, want %v", got, DefaultHybridVectorWeight) + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: true}, ChannelRetrieve); got != 1-DefaultHybridVectorWeight { + t.Fatalf("using_embedding=true → %v, want %v", got, 1-DefaultHybridVectorWeight) } override := 0.5 - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: true, VectorSimilarityWeight: &override}, ChannelRetrieve); got != 0.5 { + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: true, KeywordsSimilarityWeight: &override}, ChannelRetrieve); got != 0.5 { t.Fatalf("using_embedding=true with override → %v, want 0.5", got) } } -func TestResolveVectorWeightHybridDefaultsToThreeTenths(t *testing.T) { - // The hybrid leg defaults the vector weight to 0.3, unlike the retrieve channel's 0.7. - if got := resolveVectorWeight(SearchDeps{HasEmbedder: true}, ChannelHybrid); got != HybridSearchDefaultVectorWeight { - t.Fatalf("hybrid → %v, want %v", got, HybridSearchDefaultVectorWeight) +func TestResolveKeywordsSimilarityWeightHybridDefaultsToSevenTenths(t *testing.T) { + // Python hybrid_search defaults the vector weight to 0.3, hence its + // keyword weight is 0.7. + // (_DEFAULT_HYBRID_VECTOR_WEIGHT,), unlike RAGTools.retrieve's 0.7. + if got := resolveKeywordsSimilarityWeight(SearchDeps{HasEmbedder: true}, ChannelHybrid); got != 1-HybridSearchDefaultVectorWeight { + t.Fatalf("hybrid → %v, want %v", got, 1-HybridSearchDefaultVectorWeight) } - if got := resolveVectorWeight(SearchDeps{HasEmbedder: false}, ChannelHybrid); got != 0 { - t.Fatalf("hybrid with no embedder → %v, want 0", got) + if got := resolveKeywordsSimilarityWeight(SearchDeps{HasEmbedder: false}, ChannelHybrid); got != 1 { + t.Fatalf("hybrid with no embedder → %v, want 1 (Python: `if embd_mdl`)", got) } } -// TestResolveVectorWeightHybridIgnoresUsingEmbedding is the regression guard for the -// channel-granularity bug: the hybrid leg has NO using_embedding parameter, so gating it -// on that flag disabled the semantic leg for search_chunks — losing recall of passages -// sharing no surface words. The gate for this channel is the embedder, nothing else. -func TestResolveVectorWeightHybridIgnoresUsingEmbedding(t *testing.T) { - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: false, HasEmbedder: true}, ChannelHybrid); got != HybridSearchDefaultVectorWeight { +// TestResolveKeywordsSimilarityWeightHybridIgnoresUsingEmbedding is the regression guard for +// the channel-granularity bug: Python's hybrid_search has NO using_embedding +// parameter (search.py:hybrid_search, :143-145), so gating it on that flag disabled the +// semantic leg for search_chunks — losing recall of passages sharing no surface +// words. The gate for this channel is the embedder, nothing else. +func TestResolveKeywordsSimilarityWeightHybridIgnoresUsingEmbedding(t *testing.T) { + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: false, HasEmbedder: true}, ChannelHybrid); got != 1-HybridSearchDefaultVectorWeight { t.Fatalf("hybrid with using_embedding=false → %v, want %v: the vector leg "+ - "must NOT depend on using_embedding", got, HybridSearchDefaultVectorWeight) + "must NOT depend on using_embedding", got, 1-HybridSearchDefaultVectorWeight) } } -// TestResolveVectorWeightGrepIsAlwaysZero: the grep channel and the retrieve / grep_* -// session tools are keyword-only and have no vector leg at all, so no flag can turn one -// on. -func TestResolveVectorWeightGrepIsAlwaysZero(t *testing.T) { - override := 0.9 - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: true, HasEmbedder: true}, ChannelGrep); got != 0 { - t.Fatalf("grep → %v, want 0 (grep_search has no vector leg)", got) +// TestResolveKeywordsSimilarityWeightGrepIsAlwaysOne pins Python grep_search: the retrieve +// and grep_* session tools are keyword-only and have no vector leg at all, so no +// flag can turn one on. +func TestResolveKeywordsSimilarityWeightGrepIsAlwaysOne(t *testing.T) { + override := 0.1 + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: true, HasEmbedder: true}, ChannelGrep); got != 1 { + t.Fatalf("grep → %v, want 1 (grep_search has no vector leg)", got) } - if got := resolveVectorWeight(SearchDeps{UsingEmbedding: true, HasEmbedder: true, VectorSimilarityWeight: &override}, ChannelGrep); got != 0 { - t.Fatalf("grep with override → %v, want 0 (grep_search has no vector leg)", got) + if got := resolveKeywordsSimilarityWeight(SearchDeps{UsingEmbedding: true, HasEmbedder: true, KeywordsSimilarityWeight: &override}, ChannelGrep); got != 1 { + t.Fatalf("grep with override → %v, want 1 (grep_search has no vector leg)", got) } } diff --git a/internal/server/config/log_config.go b/internal/server/config/log_config.go index b6d4406a75..c11997866f 100644 --- a/internal/server/config/log_config.go +++ b/internal/server/config/log_config.go @@ -16,7 +16,11 @@ package config -import "github.com/spf13/viper" +import ( + "github.com/spf13/viper" + + "ragflow/internal/common" +) // LogConfig logging configuration. // @@ -42,18 +46,19 @@ type LogConfig struct { MaxSize int `mapstructure:"max_size"` // MB before rotation; default 100 MaxBackups int `mapstructure:"max_backups"` // retained rotated files; default 10 MaxAge int `mapstructure:"max_age"` // days; default 30 - Compress bool `mapstructure:"compress"` // gzip rotated files; default false + Compress *bool `mapstructure:"compress"` // gzip rotated files; nil = project default (true) } func (c *Config) ParseLogConfig(v *viper.Viper) error { - // Default Log config + // Default Log config. Rotation defaults come from common so they stay + // in sync with the logger's own fallback values. c.log.Level = "info" c.log.Format = "json" c.log.Path = "logs" - c.log.MaxSize = 100 * 1024 * 1024 // 1024MB - c.log.MaxBackups = 10 - c.log.MaxAge = 30 - c.log.Compress = false + c.log.MaxSize = common.DefaultLogMaxSizeMB + c.log.MaxBackups = common.DefaultLogMaxBackups + c.log.MaxAge = common.DefaultLogMaxAgeDays + c.log.Compress = boolPtr(common.DefaultLogCompress) if !v.IsSet("log") { return nil @@ -88,12 +93,19 @@ func (c *Config) ParseLogConfig(v *viper.Viper) error { } if sub.IsSet("compress") { - c.log.Compress = sub.GetBool("compress") + v := sub.GetBool("compress") + c.log.Compress = &v } return nil } +// boolPtr returns a pointer to the given bool, used to populate *bool config +// fields so that "not set" (nil) is distinguishable from an explicit value. +func boolPtr(b bool) *bool { + return &b +} + func (c *Config) GetLogConfig() LogConfig { return c.log } diff --git a/internal/service/chat.go b/internal/service/chat.go index c565bfde4f..1a488838d1 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -1318,6 +1318,36 @@ func (s *ChatService) GetChat(ctx context.Context, userID string, chatID string) // Resolve kb_ids to kb_names kbNames, datasetIDs := s.getDatasetNamesAndIDs(ctx, chat.KBIDs) + // Normalize fields that the frontend chat-setting form schema requires to + // be present and valid. The Python API returns these defaults; the Go + // port previously omitted them (rerank_candidates_count=0, + // reference_metadata=null), which made the form invalid and silently + // blocked Save (no request sent). Mirror Python's defaults without + // touching the persisted DB row. + if chat.RerankCandidatesCount <= 0 { + chat.RerankCandidatesCount = 64 + } + if chat.PromptConfig != nil { + refMeta, ok := chat.PromptConfig["reference_metadata"].(map[string]interface{}) + if !ok || refMeta == nil { + refMeta = map[string]interface{}{} + } + if _, hasInclude := refMeta["include"]; !hasInclude { + refMeta["include"] = false + } + if _, hasFields := refMeta["fields"]; !hasFields { + refMeta["fields"] = []interface{}{} + } + chat.PromptConfig["reference_metadata"] = refMeta + + // refine_multiturn is required by the frontend schema (z.boolean()), + // but the DB column omits it (Python defaults to false). Default it + // so the chat-setting form validates and Save can be submitted. + if _, hasRefine := chat.PromptConfig["refine_multiturn"]; !hasRefine { + chat.PromptConfig["refine_multiturn"] = false + } + } + return &GetChatResponse{ Chat: chat, DatasetIDs: datasetIDs, diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index e9081447fe..4894069e26 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -221,7 +221,8 @@ func (s *ChatPipelineService) AsyncChat( return nil, fmt.Errorf("the last content of this conversation is not from user") } - // No KBs & no web search → fast-path to LLM-only chat. + // Resolve what this conversation can reach BEFORE dispatching: whether it + // has knowledge bases, and whether web search is enabled. hasKBs := false for _, raw := range chat.KBIDs { if id, ok := raw.(string); ok && id != "" { @@ -238,6 +239,7 @@ func (s *ChatPipelineService) AsyncChat( zap.Bool("enabled", useWebSearch)) } + // No KBs & no web search → fast-path to LLM-only chat. if !hasKBs && !useWebSearch { return s.AsyncChatSolo(ctx, userID, chat, messages, stream, kwargs) } @@ -1987,7 +1989,7 @@ func (s *ChatPipelineService) tavilyRetrieve(ctx context.Context, apiKey, questi req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+apiKey) - client := &http.Client{Timeout: 30 * time.Second} + client := tavilyWebSearchHTTPClient resp, err := client.Do(req) if err != nil { return nil, fmt.Errorf("tavily: do request: %w", err) diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go index 776acd06c9..55748ea758 100644 --- a/internal/service/chat_pipeline_test.go +++ b/internal/service/chat_pipeline_test.go @@ -102,6 +102,38 @@ func TestAsyncChat_EmptyMessages(t *testing.T) { } } +// TestBuildChatConfig_GenerationConfigReachesModel pins that the dialog's LLM +// setting and per-request overrides (temperature, top_p, max_tokens, thinking, +// stop) reach the model driver, and that request-level overrides win. +func TestBuildChatConfig_GenerationConfigReachesModel(t *testing.T) { + chat := dialForTest("llm-1") + chat.LLMSetting = entity.JSONMap{ + "temperature": 0.7, + "top_p": 0.9, + "max_tokens": 512, + "thinking": true, + "stop": []interface{}{"\n", "END"}, + } + // Request-level overrides win over dialog values. + cfg := BuildChatConfig(chat, map[string]interface{}{"temperature": 0.3}) + + if cfg.Temperature == nil || *cfg.Temperature != 0.3 { + t.Fatalf("Temperature: want request override 0.3, got %v", cfg.Temperature) + } + if cfg.TopP == nil || *cfg.TopP != 0.9 { + t.Fatalf("TopP: want dialog 0.9, got %v", cfg.TopP) + } + if cfg.MaxTokens == nil || *cfg.MaxTokens != 512 { + t.Fatalf("MaxTokens: want 512, got %v", cfg.MaxTokens) + } + if cfg.Thinking == nil || !*cfg.Thinking { + t.Fatalf("Thinking: want true, got %v", cfg.Thinking) + } + if cfg.Stop == nil || len(*cfg.Stop) != 2 { + t.Fatalf("Stop: want [\"\\n\", \"END\"], got %v", cfg.Stop) + } +} + // --- P1 Timer + decorateAnswer tests (P0/P1/P7 surface) --- // TestDecorateAnswer_TimerFormatAlwaysEmitted pins the Markdown diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index e3d5bafda3..142e1b7c5f 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -1373,6 +1373,10 @@ func (s *ChatSessionService) ChatCompletions( } } + // Correlate every log line this request emits (retrieval, model calls) + // with the conversation turn's session id. + ctx = common.WithSessionID(ctx, sessionID) + common.Info("ChatCompletions started") // --- 1. Normalize messages --- @@ -1421,6 +1425,10 @@ func (s *ChatSessionService) ChatCompletions( return fail(err) } sessionID = session.ID + // Fresh-session requests carry no session_id, so the correlation + // tag has to be (re)applied with the id the server just + // allocated. + ctx = common.WithSessionID(ctx, sessionID) } if storeHistoryMessages { @@ -1513,7 +1521,6 @@ func (s *ChatSessionService) ChatCompletions( } } } - if legacy { if result.Final { if strings.Contains(result.Answer, "**ERROR**") { @@ -1523,6 +1530,12 @@ func (s *ChatSessionService) ChatCompletions( } sendOrCancel(fmt.Sprintf("data:%s\n\n", sseMarshalChunk(sanitizeJSONFloats(ans).(map[string]interface{}), chatID))) } + // Turn compaction: progressive persistence has been + // writing every delta into the assistant message. The + // next turn re-enters AsyncChat with this session as its + // history, so the stored message must end up holding the + // turn's final answer alone. + s.compactSessionAssistant(session, result.Answer, messageID) finalLegacyAnswer = s.structureAnswer(session, result.Answer, messageID, sessionID, reference) continue } @@ -1530,13 +1543,18 @@ func (s *ChatSessionService) ChatCompletions( fullAnswer.WriteString("") } else if result.EndToThink { fullAnswer.WriteString("") - } else if result.Answer != "" { - fullAnswer.WriteString(result.Answer) - } else if result.Reasoning != "" { - // In-think reasoning text between the markers (tool and - // harness paths) — keep it inside the block. + } + if result.Reasoning != "" { + // Same as the non-legacy branch: there is no reasoning + // channel in this protocol, so thinking rides in `answer` + // between the start/end flags. fullAnswer.WriteString(result.Reasoning) } + if result.Answer != "" { + // Marker and text can arrive together (see the note in the + // non-legacy branch): never trade one for the other. + fullAnswer.WriteString(result.Answer) + } if session != nil { s.appendAssistantToSession(session, fullAnswer.String(), messageID) } @@ -1555,6 +1573,12 @@ func (s *ChatSessionService) ChatCompletions( sendOrCancel(fmt.Sprintf("data:%s\n\n", sseMarshalChunk(sanitizeJSONFloats(ans).(map[string]interface{}), chatID))) } else { if result.Final { + // Turn compaction, same as the legacy branch: progressive + // persistence wrote every delta into the assistant + // message. The next user input re-enters AsyncChat with + // this session as history, so the stored message must + // hold the final answer alone. + s.compactSessionAssistant(session, result.Answer, messageID) if strings.Contains(result.Answer, "**ERROR**") { ans := s.structureAnswer(session, result.Answer, messageID, sessionID, reference) if chatID != "" { @@ -1582,25 +1606,39 @@ func (s *ChatSessionService) ChatCompletions( deltaAnswer := "" if result.StartToThink { fullAnswer.WriteString("") + deltaAnswer = "" } else if result.EndToThink { fullAnswer.WriteString("") - } else if result.Answer != "" { - fullAnswer.WriteString(result.Answer) - deltaAnswer = result.Answer - } else if result.Reasoning != "" { - // In-think reasoning text arriving between the - // StartToThink / EndToThink markers (the tool path and the - // harness high/ultra path both deliver it via Reasoning so - // OpenAI-compat can map it to reasoning_content). Treat it - // as think-block content so the chat UI shows the - // reasoning instead of dropping it. + deltaAnswer = "" + } + if result.Reasoning != "" { + // The native protocol has NO reasoning channel: the UI + // rebuilds the whole assistant message from `answer` + // alone, wrapping the stretch between start_to_think and + // end_to_think in . Reasoning must therefore + // travel in `answer` too — dropping it here is what left + // the think panel empty while the run was in progress. fullAnswer.WriteString(result.Reasoning) - deltaAnswer = result.Reasoning + deltaAnswer += result.Reasoning + } + if result.Answer != "" { + // A marker and the text it delimits can ride on the SAME + // result: the first content delta after a think block is + // the EndToThink result. Treating the marker as an + // alternative to the text drops that delta, which is how an + // answer came out starting mid-sentence. + fullAnswer.WriteString(result.Answer) + deltaAnswer += result.Answer } if session != nil { s.appendAssistantToSession(session, fullAnswer.String(), messageID) } ans := s.structureAnswer(session, deltaAnswer, messageID, sessionID, reference) + // Citations ship ONLY on the final event: the intermediate + // deltas have nothing to cite yet, and an empty + // `reference: {chunks: []}` on every chunk just buries the + // real payload. + delete(ans, "reference") ans["start_to_think"] = result.StartToThink ans["end_to_think"] = result.EndToThink // The structured twin of a reasoning step rides the chunk that @@ -1884,6 +1922,25 @@ func (s *ChatSessionService) appendAssistantToSession(session *entity.ChatSessio session.Message, _ = json.Marshal(messages) } +// compactSessionAssistant rewrites the session's stored assistant message to +// the turn's FINAL answer, which is what the next turn must see. +// +// Streaming persistence (appendAssistantToSession on every delta) deliberately +// keeps partial text in the session so a client that refreshes mid-stream still +// sees what was produced. The next user input re-enters AsyncChat with this +// session as its history, so the stored message must carry the turn's final +// answer rather than every intermediate delta. +// +// A turn therefore ends compacted to question + final answer. A blank final +// (an error result, or a run that produced nothing) leaves the streamed text +// standing rather than blanking the message the user can already see. +func (s *ChatSessionService) compactSessionAssistant(session *entity.ChatSession, final, messageID string) { + if session == nil || strings.TrimSpace(final) == "" { + return + } + s.appendAssistantToSession(session, final, messageID) +} + // getSessionMessagesAsSlice returns the session's messages as a slice of maps. func (s *ChatSessionService) getSessionMessagesAsSlice(session *entity.ChatSession) []map[string]interface{} { if session == nil { @@ -1954,8 +2011,12 @@ func (s *ChatSessionService) checkTenantLLMAPIKey(ctx context.Context, tenantID, // sseAnswerChunk has deterministic JSON field order matching Python's structure_answer output. type sseAnswerChunk struct { - Answer string `json:"answer"` - Reference map[string]interface{} `json:"reference"` + Answer string `json:"answer"` + // Reference is omitted from the JSON when the producer did not set one: + // intermediate streaming deltas have nothing to cite yet, and only the + // final event carries the citation payload (set explicitly by the + // pipeline's Final branch). sseMarshalChunk leaves it nil in that case. + Reference map[string]interface{} `json:"reference,omitempty"` AudioBinary interface{} `json:"audio_binary"` Prompt string `json:"prompt"` CreatedAt float64 `json:"created_at"` @@ -2054,9 +2115,17 @@ func sanitizeJSONFloats(v interface{}) interface{} { // sseMarshalChunk converts an answer map to the ordered sseAnswerChunk struct // and marshals it with Python-compatible JSON formatting (spaces, field order). func sseMarshalChunk(ans map[string]interface{}, chatID string) string { - ref, _ := ans["reference"].(map[string]interface{}) - if ref == nil { - ref = map[string]interface{}{"chunks": []interface{}{}} + // Reference is emitted only when the producer set one: the delta branch + // deletes the key so intermediate chunks stay reference-free, and silently + // re-adding an empty `{"chunks": []}` here would undo that — every + // intermediate chunk would carry a citation payload it cannot back. + ref := map[string]interface{}{"chunks": []interface{}{}} + if raw, hasRef := ans["reference"]; hasRef { + if m, ok := raw.(map[string]interface{}); ok && m != nil { + ref = m + } + } else { + ref = nil // key absent → omit from the JSON entirely (omitempty) } answer, _ := ans["answer"].(string) prompt, _ := ans["prompt"].(string) diff --git a/internal/service/deep_researcher.go b/internal/service/deep_researcher.go index 1c2611a027..5137f4813f 100644 --- a/internal/service/deep_researcher.go +++ b/internal/service/deep_researcher.go @@ -433,7 +433,7 @@ func (dr *DeepResearcher) tavilyRetrieve(ctx context.Context, apiKey, query stri } req.Header.Set("Content-Type", "application/json") - client := &http.Client{Timeout: 15 * time.Second} + client := tavilyDeepResearchHTTPClient resp, err := client.Do(req) if err != nil { return nil, fmt.Errorf("tavily call: %w", err) diff --git a/internal/service/ingestion_log_retention.go b/internal/service/ingestion_log_retention.go index 950c003de0..11c7f81552 100644 --- a/internal/service/ingestion_log_retention.go +++ b/internal/service/ingestion_log_retention.go @@ -91,6 +91,12 @@ func (s *IngestionTaskService) foldIngestionRun(ctx context.Context, pipelineLog if maxRows < 2 { return ingestionLogFoldResult{}, errors.New("fold ingestion run requires max rows >= 2") } + if dao.DB == nil { + // This runs on the ingestor's terminal-log path, so a nil handle used to + // crash a worker goroutine (and, via the panic-recovery handler, the + // whole process) instead of degrading to a logged error. + return ingestionLogFoldResult{}, errors.New("fold ingestion run: nil database") + } var result ingestionLogFoldResult err := dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 7da99eeedf..30229b6ca4 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -424,6 +424,13 @@ func (s *IngestionTaskService) Remove(ctx context.Context, taskID string, userID } func (s *IngestionTaskService) GetTask(ctx context.Context, taskID string) (*entity.IngestionTask, error) { + if dao.DB == nil { + // Every task-status write funnels through here (MarkFailed/Stopped/ + // Completed), and that includes the panic-recovery path in the ingestor + // worker. Dereferencing a nil handle there turns a recovered task panic + // into a panicking recovery handler, which kills the worker process. + return nil, errors.New("ingestion task: nil database") + } task, err := s.ingestionTaskDAO.GetByID(ctx, dao.DB, taskID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -465,6 +472,14 @@ func (s *IngestionTaskService) ReloadAndValidateRunIdentity(ctx context.Context, return task, nil } +// GetTaskByDocument returns the document's latest ingestion task. A document +// may be parsed multiple times over its lifetime; the newest task (by +// create_time) is the current parse round, so doc-level run state derives from +// it alone. +func (s *IngestionTaskService) GetTaskByDocument(ctx context.Context, documentID string) (*entity.IngestionTask, error) { + return s.ingestionTaskDAO.GetByDocumentID(ctx, dao.DB, documentID) +} + func validateTransition(from, to string) error { switch from { case common.CREATED: diff --git a/internal/service/model_service.go b/internal/service/model_service.go index cfcea2ebe1..5b54289f45 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3825,3 +3825,64 @@ func (m *ModelProviderService) isImage2TextLLM(ctx context.Context, tenantID, ll } return false } + +// ChatModelRef identifies one chat-capable tenant model together with its +// human-readable coordinates (provider instance), so callers can log and +// reason about failover chains. Ref is the tenant_model.id that +// ResolveModelConfig / GetChatModelConfig accepts verbatim. +type ChatModelRef struct { + Ref string + ModelName string + InstanceName string + ProviderName string +} + +// ListTenantChatModelRefs enumerates every ACTIVE chat-capable model the +// tenant owns across all provider instances. Resolution of each ref is the +// caller's business — a single broken entry here is not an error for the +// whole enumeration. +func (m *ModelProviderService) ListTenantChatModelRefs(ctx context.Context, tenantID string) ([]ChatModelRef, error) { + providers, err := m.modelProviderDAO.GetByTenantID(ctx, dao.DB, tenantID) + if err != nil { + return nil, err + } + if len(providers) == 0 { + return nil, nil + } + providerIDs := make([]string, 0, len(providers)) + providerInfoByID := make(map[string]*entity.TenantModelProvider, len(providers)) + for _, p := range providers { + providerIDs = append(providerIDs, p.ID) + providerInfoByID[p.ID] = p + } + instances, err := m.modelInstanceDAO.GetByProviderIDs(ctx, dao.DB, providerIDs) + if err != nil { + return nil, err + } + if len(instances) == 0 { + return nil, nil + } + instanceIDs := make([]string, 0, len(instances)) + instanceInfoByID := make(map[string]*entity.TenantModelInstance, len(instances)) + for _, inst := range instances { + instanceIDs = append(instanceIDs, inst.ID) + instanceInfoByID[inst.ID] = inst + } + models, err := m.modelDAO.GetActiveModelsByProviderAndInstanceIDsAndType( + ctx, dao.DB, providerIDs, instanceIDs, int(entity.ModelTypeChat)) + if err != nil { + return nil, err + } + refs := make([]ChatModelRef, 0, len(models)) + for _, rec := range models { + ref := ChatModelRef{Ref: rec.ID, ModelName: rec.ModelName} + if inst := instanceInfoByID[rec.InstanceID]; inst != nil { + ref.InstanceName = inst.InstanceName + } + if prov := providerInfoByID[rec.ProviderID]; prov != nil { + ref.ProviderName = prov.ProviderName + } + refs = append(refs, ref) + } + return refs, nil +} diff --git a/internal/service/model_solver.go b/internal/service/model_solver.go index 454816add1..c76c2ad7c3 100644 --- a/internal/service/model_solver.go +++ b/internal/service/model_solver.go @@ -289,7 +289,17 @@ func (s *ModelSolver) resolveModel(ctx context.Context, tenantID string, modelTy modelEntity, err := s.lookupTenantModel(ctx, tenantID, modelRef) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return s.resolveCompositeModel(ctx, tenantID, modelType, modelRef) + // A bare ref that is BOTH unknown as a tenant model id AND rejected as + // a composite name is almost always a DANGLING REFERENCE — a knowledge + // base or chat pointing at a tenant_model row that was deleted. Say so + // explicitly: the underlying "provider name missing in model name: + // " reads like a naming-format mistake and sends the operator to + // look at the model's name instead of at the row that no longer exists. + composite, compositeErr := s.resolveCompositeModel(ctx, tenantID, modelType, modelRef) + if compositeErr != nil && !strings.Contains(modelRef, "@") { + return nil, fmt.Errorf("model %q is neither a tenant model id (no tenant_model row) nor a valid composite name — the reference is dangling: %w", modelRef, compositeErr) + } + return composite, compositeErr } return nil, err } diff --git a/internal/service/nlp/reranker.go b/internal/service/nlp/reranker.go index 1b5ad0adad..89f73a9e72 100644 --- a/internal/service/nlp/reranker.go +++ b/internal/service/nlp/reranker.go @@ -934,6 +934,7 @@ func applyRankFeatureScoresForIDs(ids []string, field map[string]map[string]inte // - qb: QueryBuilder for token processing // - rankFeature: rank feature weights (e.g., {"pagerank_fea": 10.0}) func RerankWithKNN( + ctx context.Context, chunks []map[string]interface{}, ids []string, field map[string]map[string]interface{}, @@ -948,7 +949,7 @@ func RerankWithKNN( return []float64{}, []float64{}, []float64{} } - common.Info("RerankWithKNN started", zap.Int("chunkCount", len(ids)), zap.Float64("tkWeight", tkWeight), zap.Float64("vtWeight", vtWeight)) + common.InfoCtx(ctx, "RerankWithKNN started", zap.Int("chunkCount", len(ids)), zap.Float64("tkWeight", tkWeight), zap.Float64("vtWeight", vtWeight)) // Normalize important_kwd - Python checks if it's a string and wraps in list // for i in sres.ids: @@ -971,7 +972,7 @@ func RerankWithKNN( if qb != nil { _, keywords = qb.Question(query, "qa", 0.6) } - common.Info("RerankWithKNN keywords", zap.Any("keywords", keywords)) + common.InfoCtx(ctx, "RerankWithKNN keywords", zap.Any("keywords", keywords)) // Build token lists matching Python's OrderedDict approach insTw := make([][]string, 0, len(ids)) @@ -1010,14 +1011,14 @@ func RerankWithKNN( // Calculate token similarity tsim = TokenSimilarity(keywords, insTw, qb) - common.Info("RerankWithKNN tsim", zap.Float64s("tsim", tsim)) + common.InfoCtx(ctx, "RerankWithKNN tsim", zap.Float64s("tsim", tsim)) // Build vector similarity from knnScores - matches Python's np.array([knn_scores.get(chunk_id, 0.0) for chunk_id in sres.ids]) vsim = make([]float64, len(ids)) for i, chunkID := range ids { vsim[i] = knnScores[chunkID] // Returns 0.0 if not found (Go map default) } - common.Debug("RerankWithKNN knnScores", zap.Int("knnScoreCount", len(knnScores)), zap.Float64s("vsim", vsim), zap.Strings("ids", ids), zap.Float64s("knnScores", func() []float64 { + common.InfoCtx(ctx, "RerankWithKNN knnScores", zap.Int("knnScoreCount", len(knnScores)), zap.Float64s("vsim", vsim), zap.Strings("ids", ids), zap.Float64s("knnScores", func() []float64 { scores := make([]float64, 0, len(knnScores)) for _, id := range ids { if s, ok := knnScores[id]; ok { @@ -1035,9 +1036,9 @@ func RerankWithKNN( // Apply rank feature scores (tag_score * 10 + pagerank) sim = applyRankFeatureScoresForIDs(ids, field, sim, rankFeature) - common.Info("RerankWithKNN rankFeatureScores", zap.Any("rankFeature", rankFeature), zap.Any("simAfterRank", sim)) + common.InfoCtx(ctx, "RerankWithKNN rankFeatureScores", zap.Any("rankFeature", rankFeature), zap.Any("simAfterRank", sim)) - common.Info("RerankWithKNN completed", zap.Int("outputChunks", len(sim))) + common.InfoCtx(ctx, "RerankWithKNN completed", zap.Int("outputChunks", len(sim))) return sim, tsim, vsim } diff --git a/internal/service/nlp/retrieval.go b/internal/service/nlp/retrieval.go index 25ccbdac0e..cec0583224 100644 --- a/internal/service/nlp/retrieval.go +++ b/internal/service/nlp/retrieval.go @@ -67,6 +67,12 @@ type RetrievalRequest struct { Highlight *bool AllowDenseFallback *bool Filter map[string]interface{} + // VectorOnly, when true, runs the dense leg alone: no text match expression + // and no fusion expression reach the engine, so the search is a pure kNN + // whose filter is the scope conditions only. Set by callers that need to + // bridge a wording gap the query's own words cannot cross; the scoring pass + // then uses the engine's kNN score directly instead of re-scoring tokens. + VectorOnly bool } // RetrievalResult result from retrieval search @@ -86,7 +92,7 @@ type RetrievalResult struct { // - Build chunks // - Build document aggregation if specified func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) (*RetrievalResult, error) { - common.Info("Retrieval START", zap.String("question", req.Question), zap.Int("page", req.Page), zap.Int("pageSize", req.PageSize)) + common.InfoCtx(ctx, "Retrieval START", zap.String("question", req.Question), zap.Int("page", req.Page), zap.Int("pageSize", req.PageSize)) if req.Question == "" { return &RetrievalResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}, Total: 0}, nil } @@ -102,7 +108,15 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) req.SimilarityThreshold = func() *float64 { v := 0.2; return &v }() } if req.VectorSimilarityWeight == nil { - req.VectorSimilarityWeight = func() *float64 { v := 0.3; return &v }() + // A vector-only search is scored by vector similarity alone: no tokens + // of the query are reliable signals there (the query is a description, + // not a term list), so the default 0.3 vector weight would let token + // coincidence re-rank a pure-vector result set. + if req.VectorOnly { + req.VectorSimilarityWeight = func() *float64 { v := 1.0; return &v }() + } else { + req.VectorSimilarityWeight = func() *float64 { v := 0.3; return &v }() + } } if req.RankFeature == nil { req.RankFeature = &map[string]float64{"pagerank_fea": 10.0} @@ -129,7 +143,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) if req.RerankModel != nil && req.Page != 1 { return nil, fmt.Errorf("Pagination is not supported when rerank_mdl is specified. Please set page=1 to retrieve the top %d results.", pageSize) } - common.Debug("Retrieval rerank candidate params", zap.Int("page", req.Page), zap.Int("pageSize", pageSize), zap.Int("rerankCandidatesCount", rerankCandidatesCount)) + common.DebugCtx(ctx, "Retrieval rerank candidate params", zap.Int("page", req.Page), zap.Int("pageSize", pageSize), zap.Int("rerankCandidatesCount", rerankCandidatesCount)) // Execute search via Search() searchReq := &RetrievalSearchRequest{ @@ -146,6 +160,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) VectorSimilarityWeight: req.VectorSimilarityWeight, Highlight: req.Highlight, AllowDenseFallback: req.AllowDenseFallback, + VectorOnly: req.VectorOnly, Filter: req.Filter, } searchResult, err := s.Search(ctx, searchReq) @@ -189,6 +204,16 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) if *req.VectorSimilarityWeight <= 0 { postThreshold = 0.0 } + // A VECTOR-ONLY search has already been floored by the engine's own kNN + // similarity filter, so re-applying a threshold here buys nothing — and + // costs everything when the score is the one this pipeline reads back: a + // similarity that comes out 0 (absent from the hit's map, say) sends the + // whole candidate set through the floor and the tool reports "0 hits" for a + // corpus that certainly has neighbours. The engine's ranking is the truth + // for this leg; keep it. + if req.VectorOnly { + postThreshold = 0.0 + } // Get valid indices where score >= postThreshold validIdx := make([]int, 0) @@ -214,7 +239,7 @@ func (s *RetrievalService) Retrieval(ctx context.Context, req *RetrievalRequest) } pageIdx = validIdx[begin:end] } - common.Info("Pagination result info", zap.Int("totalValid", len(validIdx)), zap.Int("begin", begin), + common.InfoCtx(ctx, "Pagination result info", zap.Int("totalValid", len(validIdx)), zap.Int("begin", begin), zap.Int("end", end), zap.Int("chunkCount", len(pageIdx)), zap.Float64("postThreshold", postThreshold)) total := int64(len(validIdx)) @@ -465,6 +490,45 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval return sim, tsim, vsim, nil } + if req.VectorOnly { + // The engine's kNN search already ranked these chunks by vector + // similarity, so the second-pass KNNScores round trip would recompute + // exactly the same numbers — one extra search request per query for + // nothing. Report the engine score as the similarity, term similarity + // as zero (there is no term signal in a vector-only request), and let + // the caller's sort keep the engine order. + // + // The score is read from the CHUNK maps, not from searchResult.Field: + // Field carries only the requested SelectFields (the engine's GetFields + // copies exactly those), so a "_score" lookup there yields nothing and + // every similarity would come out 0 — which the caller's threshold then + // filters away wholesale (observed as a 0-hit pure-vector search over a + // corpus that certainly holds neighbours). + scores := make(map[string]float64, len(searchResult.Chunks)) + for _, chunk := range searchResult.Chunks { + id, _ := chunk["id"].(string) + if id == "" { + if alt, ok := chunk["_id"].(string); ok { + id = alt + } + } + if id == "" { + continue + } + for _, key := range []string{"_score", "SCORE", "SIMILARITY"} { + if score, ok := chunk[key].(float64); ok { + scores[id] = score + break + } + } + } + sim := make([]float64, len(searchResult.IDs)) + for i, id := range searchResult.IDs { + sim[i] = scores[id] + } + return sim, make([]float64, len(sim)), sim, nil + } + knnResult, err := s.docEngine.KNNScores(ctx, searchResult.Chunks, searchResult.QueryVector, len(searchResult.IDs)) if err != nil { common.Warn("KNNScores failed for ES, falling back to local computation", zap.Error(err)) @@ -483,6 +547,7 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval } knnScores := s.docEngine.GetScores(knnResult) sim, tsim, vsim := RerankWithKNN( + ctx, searchResult.Chunks, searchResult.IDs, searchResult.Field, @@ -499,17 +564,21 @@ func (s *RetrievalService) scoreSearchResult(ctx context.Context, req *Retrieval // RetrievalSearchRequest is the request struct for RetrievalService.Search() type RetrievalSearchRequest struct { - Question string - TenantIDs []string - KbIDs []string - DocIDs []string - KNNTopK int - KNNNumCandidates int - Page int - PageSize int - Sort bool - Highlight *bool - SimilarityThreshold float64 + Question string + TenantIDs []string + KbIDs []string + DocIDs []string + KNNTopK int + KNNNumCandidates int + Page int + PageSize int + Sort bool + Highlight *bool + SimilarityThreshold float64 + // VectorOnly runs the dense leg alone (see RetrievalRequest.VectorOnly): + // no text match expression and no fusion expression are sent, so the + // engine's kNN filter carries the scope conditions only. + VectorOnly bool RankFeature map[string]float64 Filter map[string]interface{} EmbeddingModel *models.EmbeddingModel @@ -538,10 +607,23 @@ func buildRetrievalFusionExpr(docEngineType string, topn int, vectorSimilarityWe return buildInfinityFusionExpr(topn, vectorSimilarityWeight) } + // The caller weight MUST reach the engine: the previous hardcoded + // "0.05,0.95" silently discarded it, so every hybrid search on the ES + // backend ran at 95% vector / 5% keyword regardless of what the tool or + // the model asked for - and chunk.go derives the BM25 boost from this + // same value, so the lexical leg was effectively switched off. + vectorWeight := 0.3 + if vectorSimilarityWeight != nil { + vectorWeight = *vectorSimilarityWeight + } + termWeight := math.Round((1.0-vectorWeight)*10000) / 10000 + return &types.FusionExpr{ - Method: "weighted_sum", - TopN: topn, - FusionParams: map[string]interface{}{"weights": "0.001,1"}, + Method: "weighted_sum", + TopN: topn, + FusionParams: map[string]interface{}{ + "weights": fmt.Sprintf("%g,%g", termWeight, vectorWeight), + }, } } @@ -595,15 +677,18 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque src := []string{ "docnm_kwd", "content_ltks", "kb_id", "img_id", "title_tks", "important_kwd", "position_int", "doc_id", "chunk_order_int", "page_num_int", "top_int", "create_timestamp_flt", "knowledge_graph_kwd", - "question_kwd", "question_tks", "doc_type_kwd", - "available_int", "content_with_weight", "mom_id", "pagerank_fea", "tag_feas", "row_id()", + // Fields the projection below reads off each chunk. This list IS the ES + // _source filter, so omitting one returns it empty even though the + // document carries it: content_with_weight is the text the caller + // displays, doc_type_kwd marks image/table chunks, mom_id is what promotes + // child fragments to their parent, row_id() the table row identity. + "content_with_weight", "doc_type_kwd", "mom_id", "row_id()", "_score", } kwds := make(map[string]struct{}) - // Build base engine request with common fields - // Note: RankFeature is NOT set here, it's set per-call where needed + // Build base engine request with common fields. searchRequest := &types.SearchRequest{ IndexNames: buildIndexNames(req.TenantIDs), KbIDs: req.KbIDs, @@ -613,7 +698,6 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque SelectFields: src, } - // engineResult holds the result from docEngine.Search() (types.SearchResult) // queryVector tracks the query vector for reranking var engineResult *types.SearchResult var queryVector []float64 @@ -663,8 +747,16 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque } denseTemplate := cloneDenseExpr(matchDense) - // Execute search with fusion - fusionExpr := buildRetrievalFusionExpr(s.docEngine.GetType(), knnTopK, req.VectorSimilarityWeight) + // Execute search with fusion — unless this is a VECTOR-ONLY + // request, which must reach the engine as a dense-only search: an + // empty text expression is what keeps the BM25 clause out of the + // kNN's filter (see the engine's query builder), and no fusion + // expression means no weighted-sum step over a leg that does not + // exist. + var fusionExpr *types.FusionExpr + if !req.VectorOnly { + fusionExpr = buildRetrievalFusionExpr(s.docEngine.GetType(), knnTopK, req.VectorSimilarityWeight) + } // Build source with vector column for ES searchSrc := make([]string, len(searchRequest.SelectFields)) @@ -674,10 +766,14 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque } searchRequest.SelectFields = searchSrc - if matchText == nil { - searchRequest.MatchExprs = []any{matchDense} + if req.VectorOnly || matchText == nil { + // Dense leg alone. The engine then builds a kNN query whose + // filter is the scope conditions only — no BM25 clause, so the + // result set is not restricted to chunks that contain the + // query's words. + searchRequest.MatchExprs = []interface{}{matchDense} } else { - searchRequest.MatchExprs = []any{matchText, matchDense, fusionExpr} + searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr} } searchRequest.RankFeature = req.RankFeature @@ -688,7 +784,7 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque // If result is empty, retry with relaxed conditions if engineResult.Total == 0 { _, hasDocIDFilter := filters["doc_id"] - if matchText == nil { + if req.VectorOnly || matchText == nil { if *req.AllowDenseFallback { common.Debug("Retrieval dense-only fallback after empty initial search") matchDense = cloneDenseExpr(denseTemplate) @@ -728,10 +824,10 @@ func (s *RetrievalService) Search(ctx context.Context, req *RetrievalSearchReque matchText, _ := GetQueryBuilder().Question(req.Question, "qa", 0.1) matchDense = cloneDenseExpr(denseTemplate) matchDense.ExtraOptions["similarity"] = 0.17 - if matchText == nil { - searchRequest.MatchExprs = []any{matchDense} + if req.VectorOnly || matchText == nil { + searchRequest.MatchExprs = []interface{}{matchDense} } else { - searchRequest.MatchExprs = []any{matchText, matchDense, fusionExpr} + searchRequest.MatchExprs = []interface{}{matchText, matchDense, fusionExpr} } searchRequest.RankFeature = req.RankFeature diff --git a/internal/service/nlp/retrieval_test.go b/internal/service/nlp/retrieval_test.go index 6e3542d6a7..3756b96861 100644 --- a/internal/service/nlp/retrieval_test.go +++ b/internal/service/nlp/retrieval_test.go @@ -514,9 +514,16 @@ func assertFusionWeights(t *testing.T, request *types.SearchRequest, want string func TestBuildRetrievalFusionExprKeepsPythonWeightsOutsideInfinity(t *testing.T) { expr := buildRetrievalFusionExpr(string(engine.EngineElasticsearch), 10, float64Ptr(0.8)) - // Python Dealer.search's non-Infinity branch (rag/nlp/search.py:265). - if got := expr.FusionParams["weights"]; got != "0.001,1" { - t.Fatalf("expected Elasticsearch weights=0.001,1, got %v", got) + // Elasticsearch must honour the caller weight exactly like Infinity does: + // it used to be hardcoded to "0.05,0.95", which silently discarded it. + if got := expr.FusionParams["weights"]; got != "0.2,0.8" { + t.Fatalf("expected Elasticsearch weights=0.2,0.8, got %v", got) + } + + // nil weight falls back to the documented default (0.3 vector). + expr = buildRetrievalFusionExpr(string(engine.EngineElasticsearch), 10, nil) + if got := expr.FusionParams["weights"]; got != "0.7,0.3" { + t.Fatalf("expected default Elasticsearch weights=0.7,0.3, got %v", got) } } @@ -551,7 +558,9 @@ func TestSearchKeepsPythonFusionWeightForElasticsearch(t *testing.T) { if !ok { t.Fatalf("expected third match expression to be FusionExpr, got %T", docEngine.searchRequest.MatchExprs[2]) } - if got := fusionExpr.FusionParams["weights"]; got != "0.001,1" { - t.Fatalf("expected Elasticsearch weights=0.001,1, got %v", got) + // Elasticsearch must honour the caller weight (0.8 vector -> "0.2,0.8") + // instead of the legacy hardcoded "0.05,0.95". + if got := fusionExpr.FusionParams["weights"]; got != "0.2,0.8" { + t.Fatalf("expected Elasticsearch weights=0.2,0.8, got %v", got) } } diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index 8703e6709d..a7fbb7e662 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -339,7 +339,6 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID if docIDsStr != "" { chatKwargs["doc_ids"] = docIDsStr } - asyncResults, asyncErr := s.pipeline.AsyncChat(ctx, userID, dialog, filteredMessages, openaiReq.Stream, chatKwargs) if asyncErr != nil { s.writeDataError(c, asyncErr.Error()) @@ -368,13 +367,13 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID for result := range asyncResults { lastResult = result - if result.StartToThink || result.EndToThink { - // Think markers only toggle routing state; no SSE event - // emitted. Matches Python's _stream_chat_completion_sse - // which ignores start_to_think/end_to_think flags and - // never emits "" or "" as content. - continue - } + // Think markers only toggle routing state — Python never emits + // ""/"" as content — but the result carrying a + // marker can also carry the first delta of the text it delimits + // (the first content delta after a think block IS the EndToThink + // result). The event must therefore fall through to the emission + // below instead of being dropped: dropping it loses that text, + // which is what left an answer opening mid-sentence. if result.Final { finalContent := strings.TrimSpace(result.Answer) diff --git a/internal/service/web_search_provider.go b/internal/service/web_search_provider.go index 5cb7272b3d..8faf9970e3 100644 --- a/internal/service/web_search_provider.go +++ b/internal/service/web_search_provider.go @@ -29,29 +29,67 @@ import ( "time" ) +// Providers, endpoints and clients are listed alphabetically by provider id +// (brave, exa, firecrawl, linkup, parallel, querit, serply, tavily, youcom) so +// a new provider has exactly one obvious place in each list. Tavily's endpoint +// lives with its retrieval code in chat_pipeline.go, which is why it has no +// entry here. const ( - webSearchProviderTavily = "tavily" - webSearchProviderQuerit = "querit" - webSearchProviderSerply = "serply" - webSearchProviderYouCom = "youcom" - queritWebSearchEndpoint = "https://api.querit.ai/v1/search" - serplyWebSearchEndpoint = "https://api.serply.io/v1/search/" + webSearchProviderBrave = "brave" + webSearchProviderExa = "exa" + webSearchProviderFirecrawl = "firecrawl" + webSearchProviderLinkup = "linkup" + webSearchProviderParallel = "parallel" + webSearchProviderQuerit = "querit" + webSearchProviderSerply = "serply" + webSearchProviderTavily = "tavily" + webSearchProviderYouCom = "youcom" + + braveWebSearchEndpoint = "https://api.search.brave.com/res/v1/web/search" + exaWebSearchEndpoint = "https://api.exa.ai/search" + // v2 is Firecrawl's current search shape; the v1 endpoint is deprecated. + firecrawlWebSearchEndpoint = "https://api.firecrawl.dev/v2/search" + linkupWebSearchEndpoint = "https://api.linkup.so/v1/search" + parallelWebSearchEndpoint = "https://api.parallel.ai/v1/search" + queritWebSearchEndpoint = "https://api.querit.ai/v1/search" + serplyWebSearchEndpoint = "https://api.serply.io/v1/search/" // You.com serves the same response shape from two endpoints. The keyless // one is rate-limited but needs no credentials; the keyed one lifts those // limits. The keyless endpoint rejects an X-API-Key header, so the endpoint // and the headers are always chosen together. youComWebSearchEndpoint = "https://api.you.com/v1/search" youComKeylessWebSearchEndpoint = "https://api.you.com/v1/agents/search" - youComWebSearchResultCount = 6 // Identifies RAGFlow to You.com. On the keyless endpoint there is no key to // attribute traffic to, so this is the only signal available. youComWebSearchUserAgent = "RAGFlow youdotcom-integration/infiniflow-ragflow" + // webSearchResultCount is how many hits every provider is asked for. Six + // keeps the web block in a prompt the size of one corpus chunk set — the + // point is to give the model something to cite, not to mirror a full SERP. + webSearchResultCount = 6 + // exaWebSearchMaxCharacters caps the page text Exa returns per hit: Exa + // bills for extracted content and an uncapped page would swamp the prompt. + exaWebSearchMaxCharacters = 2000 + // webSearchMaxResponseBytes caps how much of a provider's response body is + // read: six hits of JSON are a few KiB, so 4 MiB is far above any legitimate + // answer while still bounding what a misbehaving provider can make us allocate. + webSearchMaxResponseBytes = 4 << 20 ) var ( - queritWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} - serplyWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} - youComWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + braveWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + exaWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + firecrawlWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + linkupWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + parallelWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + queritWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + serplyWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + youComWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + // Tavily is reached from two call sites (the chat pipeline and the deep + // researcher) and they used different timeouts; keeping one client per + // call site preserves both while allowing connection reuse — a fresh + // http.Client per request on a hot path shares no connections at all. + tavilyWebSearchHTTPClient = &http.Client{Timeout: 30 * time.Second} + tavilyDeepResearchHTTPClient = &http.Client{Timeout: 15 * time.Second} ) type webSearchProviderConfig struct { @@ -74,16 +112,31 @@ func resolveWebSearchProvider(promptConfig map[string]interface{}) *webSearchPro } apiKeyField := "" - // You.com is usable with no credentials at all; every other provider here - // requires a key before it can be selected. + // keyOptional marks providers usable with no key. Only You.com: it serves a + // dedicated keyless endpoint (a different path, same response shape, rate + // limited per source IP, and it rejects an X-API-Key header — see + // youComEndpointFor). A key moves to the keyed endpoint and lifts the limit. + // + // Free tier != keyless: Exa's free 1,000 requests/month still requires a key + // on every call, so it does not get this carve-out. keyOptional := false switch provider { - case webSearchProviderTavily: - apiKeyField = "tavily_api_key" + case webSearchProviderBrave: + apiKeyField = "brave_api_key" + case webSearchProviderExa: + apiKeyField = "exa_api_key" + case webSearchProviderFirecrawl: + apiKeyField = "firecrawl_api_key" + case webSearchProviderLinkup: + apiKeyField = "linkup_api_key" + case webSearchProviderParallel: + apiKeyField = "parallel_api_key" case webSearchProviderQuerit: apiKeyField = "querit_api_key" case webSearchProviderSerply: apiKeyField = "serply_api_key" + case webSearchProviderTavily: + apiKeyField = "tavily_api_key" case webSearchProviderYouCom: apiKeyField = "youcom_api_key" keyOptional = true @@ -102,17 +155,81 @@ func resolveWebSearchProvider(promptConfig map[string]interface{}) *webSearchPro } } +// retrieveWebSearch dispatches one query to the configured provider for the chat +// pipeline. func (s *ChatPipelineService) retrieveWebSearch( ctx context.Context, provider *webSearchProviderConfig, question string, +) (map[string]interface{}, error) { + return retrieveWebSearchWithTavily(ctx, provider, question, s.tavilyRetrieve) +} + +// retrieveWebSearch dispatches one query to the configured provider for the deep +// researcher. +func (dr *DeepResearcher) retrieveWebSearch( + ctx context.Context, + provider *webSearchProviderConfig, + query string, +) (map[string]interface{}, error) { + return retrieveWebSearchWithTavily(ctx, provider, query, dr.tavilyRetrieve) +} + +// retrieveWebSearchWithTavily dispatches one web-search query to the configured +// provider. The chat pipeline and the deep researcher share it: the only difference +// between their former near-identical copies was which tavilyRetrieve receiver they +// called, and Tavily is the one provider implemented outside this file, so that call +// is the parameter. +func retrieveWebSearchWithTavily( + ctx context.Context, + provider *webSearchProviderConfig, + question string, + tavilyRetrieve func(context.Context, string, string) (map[string]interface{}, error), ) (map[string]interface{}, error) { if provider == nil { return nil, fmt.Errorf("web search provider is not configured") } switch provider.Provider { - case webSearchProviderTavily: - return s.tavilyRetrieve(ctx, provider.APIKey, question) + case webSearchProviderBrave: + return retrieveBraveWebSearch( + ctx, + braveWebSearchHTTPClient, + braveWebSearchEndpoint, + provider.APIKey, + question, + ) + case webSearchProviderExa: + return retrieveExaWebSearch( + ctx, + exaWebSearchHTTPClient, + exaWebSearchEndpoint, + provider.APIKey, + question, + ) + case webSearchProviderFirecrawl: + return retrieveFirecrawlWebSearch( + ctx, + firecrawlWebSearchHTTPClient, + firecrawlWebSearchEndpoint, + provider.APIKey, + question, + ) + case webSearchProviderLinkup: + return retrieveLinkupWebSearch( + ctx, + linkupWebSearchHTTPClient, + linkupWebSearchEndpoint, + provider.APIKey, + question, + ) + case webSearchProviderParallel: + return retrieveParallelWebSearch( + ctx, + parallelWebSearchHTTPClient, + parallelWebSearchEndpoint, + provider.APIKey, + question, + ) case webSearchProviderQuerit: return retrieveQueritWebSearch( ctx, @@ -129,6 +246,8 @@ func (s *ChatPipelineService) retrieveWebSearch( provider.APIKey, question, ) + case webSearchProviderTavily: + return tavilyRetrieve(ctx, provider.APIKey, question) case webSearchProviderYouCom: return retrieveYouComWebSearch( ctx, @@ -142,44 +261,483 @@ func (s *ChatPipelineService) retrieveWebSearch( } } -func (dr *DeepResearcher) retrieveWebSearch( +// webSearchHit is what every provider's response collapses into: the three +// fields a model can quote and cite. +type webSearchHit struct { + Title string + URL string + Content string +} + +// webSearchPayload turns web hits into the {chunks, doc_aggs} map retrieval +// expects — the same shape a corpus chunk would have had, because everything +// downstream (the prompt, the citation markers, the reference payload) only +// knows that shape. A web hit has no knowledge base, no positions and no +// embedding, and its "document" is the URL itself; idPrefix keeps chunk_id and +// doc_id scoped per provider so two providers cannot collide on one URL. +// +// Hits with no usable text are dropped: a result the model cannot quote is +// noise in the context and a citation it cannot defend. +func webSearchPayload(idPrefix string, hits []webSearchHit) map[string]interface{} { + chunks := make([]map[string]interface{}, 0, len(hits)) + docAggs := make([]interface{}, 0, len(hits)) + for _, hit := range hits { + content := strings.TrimSpace(hit.Content) + if content == "" || hit.URL == "" { + continue + } + // The cap lives HERE, after the filter: a provider's hit list can carry + // unusable entries (You.com merges web and news results), and stopping the + // collection loop on a raw count would skip the usable hits behind them. + if len(chunks) >= webSearchResultCount { + break + } + chunkID := idPrefix + "-" + hit.URL + chunks = append(chunks, map[string]interface{}{ + "chunk_id": chunkID, + "content_ltks": tokenizeText(content), + "content_with_weight": content, + "doc_id": chunkID, + "docnm_kwd": hit.Title, + "kb_id": []interface{}{}, + "important_kwd": []interface{}{}, + "image_id": "", + "similarity": float64(1), + "vector_similarity": float64(1), + "term_similarity": float64(0), + "vector": []float64{}, + "positions": []interface{}{}, + "url": hit.URL, + }) + docAggs = append(docAggs, map[string]interface{}{ + "doc_name": hit.Title, + "doc_id": chunkID, + "count": 1, + "url": hit.URL, + }) + } + return map[string]interface{}{ + "chunks": chunks, + "doc_aggs": docAggs, + } +} + +// webSearchRequest performs one provider call and returns the response body. +// Providers differ only in method, headers and body — the status check and the +// read are identical everywhere, and a request that never reached the provider +// must be reported the same way regardless of which one was called. Callers +// wrap the error with their own provider name. +func webSearchRequest( ctx context.Context, - provider *webSearchProviderConfig, + client *http.Client, + method string, + endpoint string, + headers map[string]string, + body io.Reader, +) ([]byte, error) { + request, err := http.NewRequestWithContext(ctx, method, endpoint, body) + if err != nil { + return nil, fmt.Errorf("new request: %w", err) + } + for name, value := range headers { + request.Header.Set(name, value) + } + response, err := client.Do(request) + if err != nil { + return nil, fmt.Errorf("do request: %w", err) + } + defer response.Body.Close() + + if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { + return nil, fmt.Errorf("status %d", response.StatusCode) + } + // Cap the body: a misbehaving or compromised provider must not be able to make + // the server allocate an unbounded buffer. webSearchMaxResponseBytes is far above + // the largest legitimate response (six hits, and Exa's excerpts are capped at + // 2000 characters each), and exceeding it is an error rather than a silent + // truncation that would surface later as an unparseable JSON body. + responseBody, err := io.ReadAll(io.LimitReader(response.Body, webSearchMaxResponseBytes+1)) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + if len(responseBody) > webSearchMaxResponseBytes { + return nil, fmt.Errorf("response body exceeds %d bytes", webSearchMaxResponseBytes) + } + return responseBody, nil +} + +// --- Brave Search ----------------------------------------------------------- + +type braveWebSearchResult struct { + URL string `json:"url"` + Title string `json:"title"` + Description string `json:"description"` +} + +type braveWebSearchResponse struct { + Web struct { + Results []braveWebSearchResult `json:"results"` + } `json:"web"` +} + +func decodeBraveWebSearchResults(responseBody []byte) ([]braveWebSearchResult, error) { + var decoded braveWebSearchResponse + if err := json.Unmarshal(responseBody, &decoded); err != nil { + return nil, fmt.Errorf("brave: decode response: %w", err) + } + return decoded.Web.Results, nil +} + +func retrieveBraveWebSearch( + ctx context.Context, + client *http.Client, + endpoint string, + apiKey string, query string, ) (map[string]interface{}, error) { - if provider == nil { - return nil, fmt.Errorf("web search provider is not configured") + parameters := url.Values{} + parameters.Set("q", query) + parameters.Set("count", strconv.Itoa(webSearchResultCount)) + + responseBody, err := webSearchRequest(ctx, client, http.MethodGet, + endpoint+"?"+parameters.Encode(), map[string]string{ + "Accept": "application/json", + "X-Subscription-Token": apiKey, + }, nil) + if err != nil { + return nil, fmt.Errorf("brave: %w", err) } - switch provider.Provider { - case webSearchProviderTavily: - return dr.tavilyRetrieve(ctx, provider.APIKey, query) - case webSearchProviderQuerit: - return retrieveQueritWebSearch( - ctx, - queritWebSearchHTTPClient, - queritWebSearchEndpoint, - provider.APIKey, - query, - ) - case webSearchProviderSerply: - return retrieveSerplyWebSearch( - ctx, - serplyWebSearchHTTPClient, - serplyWebSearchEndpoint, - provider.APIKey, - query, - ) - case webSearchProviderYouCom: - return retrieveYouComWebSearch( - ctx, - youComWebSearchHTTPClient, - youComEndpointFor(provider.APIKey), - provider.APIKey, - query, - ) - default: - return nil, fmt.Errorf("unsupported web search provider %q", provider.Provider) + results, err := decodeBraveWebSearchResults(responseBody) + if err != nil { + return nil, err } + + hits := make([]webSearchHit, 0, len(results)) + for _, result := range results { + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: result.Description, + }) + } + return webSearchPayload("brave", hits), nil +} + +// --- Exa -------------------------------------------------------------------- + +type exaWebSearchResult struct { + Title string `json:"title"` + URL string `json:"url"` + Text string `json:"text"` + Summary string `json:"summary"` + Highlights []string `json:"highlights"` +} + +func decodeExaWebSearchResults(responseBody []byte) ([]exaWebSearchResult, error) { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(responseBody, &envelope); err != nil { + return nil, fmt.Errorf("exa: decode response: %w", err) + } + if envelope == nil { + return nil, fmt.Errorf("exa: response must be an object") + } + + resultsValue, exists := envelope["results"] + if !exists || strings.TrimSpace(string(resultsValue)) == "null" { + return []exaWebSearchResult{}, nil + } + var results []exaWebSearchResult + if err := json.Unmarshal(resultsValue, &results); err != nil { + return nil, fmt.Errorf("exa: response field results must be an array: %w", err) + } + return results, nil +} + +// exaContent prefers the page text Exa extracted, then its summary, then the +// matched highlights — whichever the response actually carries. +func exaContent(result exaWebSearchResult) string { + if text := strings.TrimSpace(result.Text); text != "" { + return text + } + if summary := strings.TrimSpace(result.Summary); summary != "" { + return summary + } + passages := make([]string, 0, len(result.Highlights)) + for _, highlight := range result.Highlights { + if strings.TrimSpace(highlight) != "" { + passages = append(passages, highlight) + } + } + return strings.Join(passages, "\n") +} + +func retrieveExaWebSearch( + ctx context.Context, + client *http.Client, + endpoint string, + apiKey string, + query string, +) (map[string]interface{}, error) { + requestBody, err := json.Marshal(map[string]interface{}{ + "query": query, + "numResults": webSearchResultCount, + // Cap the extracted text: a full page would blow the web block of the + // prompt, and the model only needs enough to quote. + "contents": map[string]interface{}{ + "text": map[string]interface{}{"maxCharacters": exaWebSearchMaxCharacters}, + }, + }) + if err != nil { + return nil, fmt.Errorf("exa: marshal request: %w", err) + } + + // Exa authenticates every request, including on its free tier — there is no + // unauthenticated search path, so the header is unconditional here. + responseBody, err := webSearchRequest(ctx, client, http.MethodPost, endpoint, map[string]string{ + "Accept": "application/json", + "Content-Type": "application/json", + "x-api-key": apiKey, + }, bytes.NewReader(requestBody)) + if err != nil { + return nil, fmt.Errorf("exa: %w", err) + } + results, err := decodeExaWebSearchResults(responseBody) + if err != nil { + return nil, err + } + + hits := make([]webSearchHit, 0, len(results)) + for _, result := range results { + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: exaContent(result), + }) + } + return webSearchPayload("exa", hits), nil +} + +// --- Firecrawl -------------------------------------------------------------- + +type firecrawlWebSearchResult struct { + URL string `json:"url"` + Title string `json:"title"` + Description string `json:"description"` + Markdown string `json:"markdown"` +} + +type firecrawlWebSearchResponse struct { + Data struct { + Web []firecrawlWebSearchResult `json:"web"` + } `json:"data"` +} + +func decodeFirecrawlWebSearchResults(responseBody []byte) ([]firecrawlWebSearchResult, error) { + var decoded firecrawlWebSearchResponse + if err := json.Unmarshal(responseBody, &decoded); err != nil { + return nil, fmt.Errorf("firecrawl: decode response: %w", err) + } + return decoded.Data.Web, nil +} + +// firecrawlContent prefers the snippet the search already returned. Markdown is +// only present when the caller asked the endpoint to scrape the hits, which +// costs extra credits per result — this integration does not. +func firecrawlContent(result firecrawlWebSearchResult) string { + if description := strings.TrimSpace(result.Description); description != "" { + return description + } + return strings.TrimSpace(result.Markdown) +} + +func retrieveFirecrawlWebSearch( + ctx context.Context, + client *http.Client, + endpoint string, + apiKey string, + query string, +) (map[string]interface{}, error) { + requestBody, err := json.Marshal(map[string]interface{}{ + "query": query, + "limit": webSearchResultCount, + }) + if err != nil { + return nil, fmt.Errorf("firecrawl: marshal request: %w", err) + } + + responseBody, err := webSearchRequest(ctx, client, http.MethodPost, endpoint, map[string]string{ + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": "Bearer " + apiKey, + }, bytes.NewReader(requestBody)) + if err != nil { + return nil, fmt.Errorf("firecrawl: %w", err) + } + results, err := decodeFirecrawlWebSearchResults(responseBody) + if err != nil { + return nil, err + } + + hits := make([]webSearchHit, 0, len(results)) + for _, result := range results { + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: firecrawlContent(result), + }) + } + return webSearchPayload("firecrawl", hits), nil +} + +// --- Linkup ----------------------------------------------------------------- + +type linkupWebSearchResult struct { + Type string `json:"type"` + Name string `json:"name"` + URL string `json:"url"` + Content string `json:"content"` +} + +func decodeLinkupWebSearchResults(responseBody []byte) ([]linkupWebSearchResult, error) { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(responseBody, &envelope); err != nil { + return nil, fmt.Errorf("linkup: decode response: %w", err) + } + if envelope == nil { + return nil, fmt.Errorf("linkup: response must be an object") + } + + resultsValue, exists := envelope["results"] + if !exists || strings.TrimSpace(string(resultsValue)) == "null" { + return []linkupWebSearchResult{}, nil + } + var results []linkupWebSearchResult + if err := json.Unmarshal(resultsValue, &results); err != nil { + return nil, fmt.Errorf("linkup: response field results must be an array: %w", err) + } + return results, nil +} + +func retrieveLinkupWebSearch( + ctx context.Context, + client *http.Client, + endpoint string, + apiKey string, + query string, +) (map[string]interface{}, error) { + requestBody, err := json.Marshal(map[string]interface{}{ + "q": query, + "depth": "standard", + "outputType": "searchResults", + }) + if err != nil { + return nil, fmt.Errorf("linkup: marshal request: %w", err) + } + + responseBody, err := webSearchRequest(ctx, client, http.MethodPost, endpoint, map[string]string{ + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": "Bearer " + apiKey, + }, bytes.NewReader(requestBody)) + if err != nil { + return nil, fmt.Errorf("linkup: %w", err) + } + results, err := decodeLinkupWebSearchResults(responseBody) + if err != nil { + return nil, err + } + + hits := make([]webSearchHit, 0, len(results)) + for _, result := range results { + hits = append(hits, webSearchHit{ + Title: result.Name, + URL: result.URL, + Content: result.Content, + }) + } + return webSearchPayload("linkup", hits), nil +} + +// --- Parallel --------------------------------------------------------------- + +type parallelWebSearchResult struct { + URL string `json:"url"` + Title string `json:"title"` + Excerpts []string `json:"excerpts"` +} + +func decodeParallelWebSearchResults(responseBody []byte) ([]parallelWebSearchResult, error) { + var envelope map[string]json.RawMessage + if err := json.Unmarshal(responseBody, &envelope); err != nil { + return nil, fmt.Errorf("parallel: decode response: %w", err) + } + if envelope == nil { + return nil, fmt.Errorf("parallel: response must be an object") + } + + resultsValue, exists := envelope["results"] + if !exists || strings.TrimSpace(string(resultsValue)) == "null" { + return []parallelWebSearchResult{}, nil + } + var results []parallelWebSearchResult + if err := json.Unmarshal(resultsValue, &results); err != nil { + return nil, fmt.Errorf("parallel: response field results must be an array: %w", err) + } + return results, nil +} + +// parallelContent joins the excerpts Parallel extracted for a page. It returns +// nothing when a result carries no passage at all — such a hit gives the model +// a URL to cite but no text to quote. +func parallelContent(result parallelWebSearchResult) string { + passages := make([]string, 0, len(result.Excerpts)) + for _, excerpt := range result.Excerpts { + if strings.TrimSpace(excerpt) != "" { + passages = append(passages, excerpt) + } + } + return strings.Join(passages, "\n") +} + +func retrieveParallelWebSearch( + ctx context.Context, + client *http.Client, + endpoint string, + apiKey string, + query string, +) (map[string]interface{}, error) { + // search_queries carries the literal terms; objective is the same question + // in natural language and steers which of the matched pages come back. + requestBody, err := json.Marshal(map[string]interface{}{ + "search_queries": []string{query}, + "objective": query, + }) + if err != nil { + return nil, fmt.Errorf("parallel: marshal request: %w", err) + } + + responseBody, err := webSearchRequest(ctx, client, http.MethodPost, endpoint, map[string]string{ + "Accept": "application/json", + "Content-Type": "application/json", + "x-api-key": apiKey, + }, bytes.NewReader(requestBody)) + if err != nil { + return nil, fmt.Errorf("parallel: %w", err) + } + results, err := decodeParallelWebSearchResults(responseBody) + if err != nil { + return nil, err + } + + hits := make([]webSearchHit, 0, len(results)) + for _, result := range results { + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: parallelContent(result), + }) + } + return webSearchPayload("parallel", hits), nil } type queritWebSearchResult struct { @@ -197,75 +755,35 @@ func retrieveQueritWebSearch( ) (map[string]interface{}, error) { requestBody, err := json.Marshal(map[string]interface{}{ "query": query, - "count": 6, + "count": webSearchResultCount, "chunksPerDoc": 1, }) if err != nil { return nil, fmt.Errorf("querit: marshal request: %w", err) } - request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(requestBody)) + responseBody, err := webSearchRequest(ctx, client, http.MethodPost, endpoint, map[string]string{ + "Accept": "application/json", + "Authorization": "Bearer " + apiKey, + "Content-Type": "application/json", + }, bytes.NewReader(requestBody)) if err != nil { - return nil, fmt.Errorf("querit: new request: %w", err) - } - request.Header.Set("Accept", "application/json") - request.Header.Set("Authorization", "Bearer "+apiKey) - request.Header.Set("Content-Type", "application/json") - - response, err := client.Do(request) - if err != nil { - return nil, fmt.Errorf("querit: do request: %w", err) - } - defer response.Body.Close() - - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return nil, fmt.Errorf("querit: status %d", response.StatusCode) - } - - responseBody, err := io.ReadAll(response.Body) - if err != nil { - return nil, fmt.Errorf("querit: read response: %w", err) + return nil, fmt.Errorf("querit: %w", err) } results, err := decodeQueritWebSearchResults(responseBody) if err != nil { return nil, err } - chunks := make([]map[string]interface{}, 0, len(results)) - docAggs := make([]interface{}, 0, len(results)) + hits := make([]webSearchHit, 0, len(results)) for _, result := range results { - if result.Snippet == "" { - continue - } - chunkID := "querit-" + result.URL - chunks = append(chunks, map[string]interface{}{ - "chunk_id": chunkID, - "content_ltks": tokenizeText(result.Snippet), - "content_with_weight": result.Snippet, - "doc_id": chunkID, - "docnm_kwd": result.Title, - "kb_id": []interface{}{}, - "important_kwd": []interface{}{}, - "image_id": "", - "similarity": float64(1), - "vector_similarity": float64(1), - "term_similarity": float64(0), - "vector": []float64{}, - "positions": []interface{}{}, - "url": result.URL, - }) - docAggs = append(docAggs, map[string]interface{}{ - "doc_name": result.Title, - "doc_id": chunkID, - "count": 1, - "url": result.URL, + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: result.Snippet, }) } - - return map[string]interface{}{ - "chunks": chunks, - "doc_aggs": docAggs, - }, nil + return webSearchPayload("querit", hits), nil } func decodeQueritWebSearchResults(responseBody []byte) ([]queritWebSearchResult, error) { @@ -319,73 +837,33 @@ func retrieveSerplyWebSearch( ) (map[string]interface{}, error) { parameters := url.Values{} parameters.Set("q", query) - parameters.Set("num", "6") + parameters.Set("num", strconv.Itoa(webSearchResultCount)) - request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint+"?"+parameters.Encode(), nil) - if err != nil { - return nil, fmt.Errorf("serply: new request: %w", err) - } - request.Header.Set("Accept", "application/json") - request.Header.Set("X-Api-Key", apiKey) // Serply sits behind Cloudflare, which rejects requests without an // explicit User-Agent, so always send one. - request.Header.Set("User-Agent", "ragflow-web-search") - - response, err := client.Do(request) + responseBody, err := webSearchRequest(ctx, client, http.MethodGet, + endpoint+"?"+parameters.Encode(), map[string]string{ + "Accept": "application/json", + "X-Api-Key": apiKey, + "User-Agent": "ragflow-web-search", + }, nil) if err != nil { - return nil, fmt.Errorf("serply: do request: %w", err) - } - defer response.Body.Close() - - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return nil, fmt.Errorf("serply: status %d", response.StatusCode) - } - - responseBody, err := io.ReadAll(response.Body) - if err != nil { - return nil, fmt.Errorf("serply: read response: %w", err) + return nil, fmt.Errorf("serply: %w", err) } results, err := decodeSerplyWebSearchResults(responseBody) if err != nil { return nil, err } - chunks := make([]map[string]interface{}, 0, len(results)) - docAggs := make([]interface{}, 0, len(results)) + hits := make([]webSearchHit, 0, len(results)) for _, result := range results { - description := strings.TrimSpace(result.Description) - if description == "" { - continue - } - chunkID := "serply-" + result.Link - chunks = append(chunks, map[string]interface{}{ - "chunk_id": chunkID, - "content_ltks": tokenizeText(description), - "content_with_weight": description, - "doc_id": chunkID, - "docnm_kwd": result.Title, - "kb_id": []interface{}{}, - "important_kwd": []interface{}{}, - "image_id": "", - "similarity": float64(1), - "vector_similarity": float64(1), - "term_similarity": float64(0), - "vector": []float64{}, - "positions": []interface{}{}, - "url": result.Link, - }) - docAggs = append(docAggs, map[string]interface{}{ - "doc_name": result.Title, - "doc_id": chunkID, - "count": 1, - "url": result.Link, + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.Link, + Content: result.Description, }) } - - return map[string]interface{}{ - "chunks": chunks, - "doc_aggs": docAggs, - }, nil + return webSearchPayload("serply", hits), nil } func decodeSerplyWebSearchResults(responseBody []byte) ([]serplyWebSearchResult, error) { @@ -458,34 +936,23 @@ func retrieveYouComWebSearch( apiKey string, query string, ) (map[string]interface{}, error) { - request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil) - if err != nil { - return nil, fmt.Errorf("youcom: new request: %w", err) - } - queryParams := request.URL.Query() + queryParams := url.Values{} queryParams.Set("query", query) - queryParams.Set("count", strconv.Itoa(youComWebSearchResultCount)) - request.URL.RawQuery = queryParams.Encode() + queryParams.Set("count", strconv.Itoa(webSearchResultCount)) - request.Header.Set("Accept", "application/json") - request.Header.Set("User-Agent", youComWebSearchUserAgent) + headers := map[string]string{ + "Accept": "application/json", + "User-Agent": youComWebSearchUserAgent, + } + // The keyless endpoint rejects an X-API-Key header, so the key rides only + // when one is configured. if trimmedKey := strings.TrimSpace(apiKey); trimmedKey != "" { - request.Header.Set("X-API-Key", trimmedKey) + headers["X-API-Key"] = trimmedKey } - - response, err := client.Do(request) + responseBody, err := webSearchRequest(ctx, client, http.MethodGet, + endpoint+"?"+queryParams.Encode(), headers, nil) if err != nil { - return nil, fmt.Errorf("youcom: do request: %w", err) - } - defer response.Body.Close() - - if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices { - return nil, fmt.Errorf("youcom: status %d", response.StatusCode) - } - - responseBody, err := io.ReadAll(response.Body) - if err != nil { - return nil, fmt.Errorf("youcom: read response: %w", err) + return nil, fmt.Errorf("youcom: %w", err) } var decoded youComWebSearchResponse @@ -499,43 +966,13 @@ func retrieveYouComWebSearch( merged = append(merged, decoded.Results.Web...) merged = append(merged, decoded.Results.News...) - chunks := make([]map[string]interface{}, 0, len(merged)) - docAggs := make([]interface{}, 0, len(merged)) + hits := make([]webSearchHit, 0, len(merged)) for _, result := range merged { - if len(chunks) >= youComWebSearchResultCount { - break - } - content := youComContent(result) - if content == "" { - continue - } - chunkID := "youcom-" + result.URL - chunks = append(chunks, map[string]interface{}{ - "chunk_id": chunkID, - "content_ltks": tokenizeText(content), - "content_with_weight": content, - "doc_id": chunkID, - "docnm_kwd": result.Title, - "kb_id": []interface{}{}, - "important_kwd": []interface{}{}, - "image_id": "", - "similarity": float64(1), - "vector_similarity": float64(1), - "term_similarity": float64(0), - "vector": []float64{}, - "positions": []interface{}{}, - "url": result.URL, - }) - docAggs = append(docAggs, map[string]interface{}{ - "doc_name": result.Title, - "doc_id": chunkID, - "count": 1, - "url": result.URL, + hits = append(hits, webSearchHit{ + Title: result.Title, + URL: result.URL, + Content: youComContent(result), }) } - - return map[string]interface{}{ - "chunks": chunks, - "doc_aggs": docAggs, - }, nil + return webSearchPayload("youcom", hits), nil } diff --git a/internal/service/web_search_provider_test.go b/internal/service/web_search_provider_test.go index cbbf46e26f..bc4c6aec0b 100644 --- a/internal/service/web_search_provider_test.go +++ b/internal/service/web_search_provider_test.go @@ -134,9 +134,9 @@ func TestResolveWebSearchProviderRequiresKeyForSelectedProvider(t *testing.T) { name string config map[string]interface{} }{ - {name: "tavily", config: map[string]interface{}{"web_search_provider": "tavily"}}, {name: "querit", config: map[string]interface{}{"web_search_provider": "querit"}}, {name: "serply", config: map[string]interface{}{"web_search_provider": "serply"}}, + {name: "tavily", config: map[string]interface{}{"web_search_provider": "tavily"}}, { name: "serply does not fall back to tavily", config: map[string]interface{}{ @@ -407,6 +407,31 @@ func TestResolveWebSearchProviderSelectsYouComWithoutAKey(t *testing.T) { } } +// Exa has a free tier of 1,000 requests/month, but it is NOT keyless: every +// REST call must carry a key, so selecting Exa without one leaves web search +// unconfigured. (Only You.com is keyless — see the carve-out below.) +func TestResolveWebSearchProviderRequiresExaKeyForFreeTier(t *testing.T) { + if got := resolveWebSearchProvider(map[string]interface{}{ + "web_search_provider": "exa", + }); got != nil { + t.Fatalf("provider without a key = %+v, want nil", got) + } + + provider := resolveWebSearchProvider(map[string]interface{}{ + "web_search_provider": "exa", + "exa_api_key": "exa-test", + }) + if provider == nil { + t.Fatal("provider is nil") + } + if provider.Provider != webSearchProviderExa { + t.Fatalf("provider = %q, want %q", provider.Provider, webSearchProviderExa) + } + if provider.APIKey != "exa-test" { + t.Fatalf("api key = %q, want %q", provider.APIKey, "exa-test") + } +} + func TestResolveWebSearchProviderTrimsOptionalYouComKey(t *testing.T) { provider := resolveWebSearchProvider(map[string]interface{}{ "web_search_provider": "youcom", @@ -424,7 +449,7 @@ func TestResolveWebSearchProviderTrimsOptionalYouComKey(t *testing.T) { func TestResolveWebSearchProviderStillRequiresKeysForKeyedProviders(t *testing.T) { // The You.com carve-out must not relax any other provider. - for _, provider := range []string{"tavily", "querit", "serply"} { + for _, provider := range []string{"querit", "serply", "tavily"} { t.Run(provider, func(t *testing.T) { if got := resolveWebSearchProvider(map[string]interface{}{ "web_search_provider": provider, @@ -624,8 +649,8 @@ func TestRetrieveYouComWebSearchCapsMergedSections(t *testing.T) { // `count` applies per section, so the merged list is trimmed back to 6. chunks := result["chunks"].([]map[string]interface{}) - if len(chunks) != youComWebSearchResultCount { - t.Fatalf("chunks = %d, want %d", len(chunks), youComWebSearchResultCount) + if len(chunks) != webSearchResultCount { + t.Fatalf("chunks = %d, want %d", len(chunks), webSearchResultCount) } } @@ -639,3 +664,268 @@ func TestRetrieveYouComWebSearchRejectsErrorStatuses(t *testing.T) { t.Fatal("expected an error for a non-2xx status") } } + +// Brave, Exa, Firecrawl, Linkup and Parallel all authenticate with a key and +// ship no keyless path, so selecting one without a key must leave web search +// unconfigured rather than silently degrade. (Exa has a free tier, but the key +// is still mandatory — see its own test.) +func TestResolveWebSearchProviderSelectsKeyedProviders(t *testing.T) { + cases := []struct { + provider string + apiKeyName string + apiKey string + }{ + {provider: "brave", apiKeyName: "brave_api_key", apiKey: "brave-test"}, + {provider: "exa", apiKeyName: "exa_api_key", apiKey: "exa-test"}, + {provider: "firecrawl", apiKeyName: "firecrawl_api_key", apiKey: "firecrawl-test"}, + {provider: "linkup", apiKeyName: "linkup_api_key", apiKey: "linkup-test"}, + {provider: "parallel", apiKeyName: "parallel_api_key", apiKey: "parallel-test"}, + } + + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + provider := resolveWebSearchProvider(map[string]interface{}{ + "web_search_provider": tc.provider, + tc.apiKeyName: " " + tc.apiKey + " ", + // A configured Tavily key must not be picked up instead. + "tavily_api_key": "tvly-test", + }) + if provider == nil { + t.Fatal("provider is nil") + } + if provider.Provider != tc.provider { + t.Fatalf("provider = %q, want %q", provider.Provider, tc.provider) + } + if provider.APIKey != tc.apiKey { + t.Fatalf("api key = %q, want %q", provider.APIKey, tc.apiKey) + } + + if got := resolveWebSearchProvider(map[string]interface{}{ + "web_search_provider": tc.provider, + }); got != nil { + t.Fatalf("provider without a key = %+v, want nil", got) + } + }) + } +} + +// assertWebSearchChunk checks the one chunk every provider must produce: it +// carries the hit's URL twice (as the chunk/document id suffix and as `url`) so +// the citation machinery can resolve it. +func assertWebSearchChunk(t *testing.T, result map[string]interface{}, wantID, wantTitle, wantContent string) { + t.Helper() + chunks, ok := result["chunks"].([]map[string]interface{}) + if !ok || len(chunks) != 1 { + t.Fatalf("chunks = %#v, want one chunk", result["chunks"]) + } + if chunks[0]["chunk_id"] != wantID { + t.Fatalf("chunk_id = %#v, want %q", chunks[0]["chunk_id"], wantID) + } + if chunks[0]["docnm_kwd"] != wantTitle { + t.Fatalf("title = %#v, want %q", chunks[0]["docnm_kwd"], wantTitle) + } + if chunks[0]["content_with_weight"] != wantContent { + t.Fatalf("content = %#v, want %q", chunks[0]["content_with_weight"], wantContent) + } + if chunks[0]["url"] != "https://example.com/ragflow" { + t.Fatalf("url = %#v", chunks[0]["url"]) + } + if chunks[0]["similarity"] != float64(1) { + t.Fatalf("similarity = %#v, want 1", chunks[0]["similarity"]) + } + aggs, ok := result["doc_aggs"].([]interface{}) + if !ok || len(aggs) != 1 { + t.Fatalf("doc_aggs = %#v, want one aggregate", result["doc_aggs"]) + } +} + +func TestRetrieveParallelWebSearchSendsKeyAndJoinsExcerpts(t *testing.T) { + var requestBody map[string]interface{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("x-api-key"); got != "parallel-test" { + t.Errorf("x-api-key = %q, want %q", got, "parallel-test") + } + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request: %v", err) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [{ + "url": "https://example.com/ragflow", + "title": "RAGFlow", + "excerpts": ["An open-source", "RAG engine."] + }] + }`)) + })) + defer server.Close() + + result, err := retrieveParallelWebSearch(t.Context(), server.Client(), server.URL, "parallel-test", "What is RAGFlow?") + if err != nil { + t.Fatalf("retrieve Parallel web search: %v", err) + } + queries, ok := requestBody["search_queries"].([]interface{}) + if !ok || len(queries) != 1 || queries[0] != "What is RAGFlow?" { + t.Fatalf("search_queries = %#v, want the question", requestBody["search_queries"]) + } + if requestBody["objective"] != "What is RAGFlow?" { + t.Fatalf("objective = %#v", requestBody["objective"]) + } + assertWebSearchChunk(t, result, "parallel-https://example.com/ragflow", "RAGFlow", "An open-source\nRAG engine.") +} + +func TestRetrieveBraveWebSearchSendsSubscriptionToken(t *testing.T) { + var requestQuery url.Values + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("X-Subscription-Token"); got != "brave-test" { + t.Errorf("X-Subscription-Token = %q, want %q", got, "brave-test") + } + requestQuery = r.URL.Query() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "web": {"results": [{ + "url": "https://example.com/ragflow", + "title": "RAGFlow", + "description": "An open-source RAG engine." + }]} + }`)) + })) + defer server.Close() + + result, err := retrieveBraveWebSearch(t.Context(), server.Client(), server.URL, "brave-test", "What is RAGFlow?") + if err != nil { + t.Fatalf("retrieve Brave web search: %v", err) + } + if got := requestQuery.Get("q"); got != "What is RAGFlow?" { + t.Fatalf("q = %q, want %q", got, "What is RAGFlow?") + } + if got := requestQuery.Get("count"); got != "6" { + t.Fatalf("count = %q, want 6", got) + } + assertWebSearchChunk(t, result, "brave-https://example.com/ragflow", "RAGFlow", "An open-source RAG engine.") +} + +func TestRetrieveExaWebSearchCapsExtractedText(t *testing.T) { + var requestBody map[string]interface{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("x-api-key"); got != "exa-test" { + t.Errorf("x-api-key = %q, want %q", got, "exa-test") + } + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request: %v", err) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [{ + "url": "https://example.com/ragflow", + "title": "RAGFlow", + "text": "An open-source RAG engine." + }] + }`)) + })) + defer server.Close() + + result, err := retrieveExaWebSearch(t.Context(), server.Client(), server.URL, "exa-test", "What is RAGFlow?") + if err != nil { + t.Fatalf("retrieve Exa web search: %v", err) + } + if requestBody["query"] != "What is RAGFlow?" { + t.Fatalf("query = %#v", requestBody["query"]) + } + if requestBody["numResults"] != float64(6) { + t.Fatalf("numResults = %#v, want 6", requestBody["numResults"]) + } + // Uncapped page text would swamp the prompt and is billed per character. + contents, ok := requestBody["contents"].(map[string]interface{}) + if !ok { + t.Fatalf("contents = %#v, want an object", requestBody["contents"]) + } + text, ok := contents["text"].(map[string]interface{}) + if !ok || text["maxCharacters"] != float64(exaWebSearchMaxCharacters) { + t.Fatalf("contents.text = %#v, want a maxCharacters cap", contents["text"]) + } + assertWebSearchChunk(t, result, "exa-https://example.com/ragflow", "RAGFlow", "An open-source RAG engine.") +} + +func TestRetrieveLinkupWebSearchAsksForSearchResults(t *testing.T) { + var requestBody map[string]interface{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer linkup-test" { + t.Errorf("Authorization = %q, want %q", got, "Bearer linkup-test") + } + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request: %v", err) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "results": [{ + "type": "text", + "name": "RAGFlow", + "url": "https://example.com/ragflow", + "content": "An open-source RAG engine." + }] + }`)) + })) + defer server.Close() + + result, err := retrieveLinkupWebSearch(t.Context(), server.Client(), server.URL, "linkup-test", "What is RAGFlow?") + if err != nil { + t.Fatalf("retrieve Linkup web search: %v", err) + } + if requestBody["q"] != "What is RAGFlow?" { + t.Fatalf("q = %#v", requestBody["q"]) + } + // searchResults, not sourcedAnswer: the caller wants hits to cite, not a + // synthesized answer the model would have to take on trust. + if requestBody["outputType"] != "searchResults" { + t.Fatalf("outputType = %#v, want searchResults", requestBody["outputType"]) + } + assertWebSearchChunk(t, result, "linkup-https://example.com/ragflow", "RAGFlow", "An open-source RAG engine.") +} + +func TestRetrieveFirecrawlWebSearchUsesSearchSnippets(t *testing.T) { + var requestBody map[string]interface{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer firecrawl-test" { + t.Errorf("Authorization = %q, want %q", got, "Bearer firecrawl-test") + } + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request: %v", err) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "success": true, + "data": {"web": [{ + "url": "https://example.com/ragflow", + "title": "RAGFlow", + "description": "An open-source RAG engine." + }]} + }`)) + })) + defer server.Close() + + result, err := retrieveFirecrawlWebSearch(t.Context(), server.Client(), server.URL, "firecrawl-test", "What is RAGFlow?") + if err != nil { + t.Fatalf("retrieve Firecrawl web search: %v", err) + } + if requestBody["query"] != "What is RAGFlow?" { + t.Fatalf("query = %#v", requestBody["query"]) + } + if requestBody["limit"] != float64(6) { + t.Fatalf("limit = %#v, want 6", requestBody["limit"]) + } + assertWebSearchChunk(t, result, "firecrawl-https://example.com/ragflow", "RAGFlow", "An open-source RAG engine.") +} + +// A hit with no text is dropped rather than shipped as an empty citation. +func TestWebSearchPayloadSkipsHitsWithoutContent(t *testing.T) { + payload := webSearchPayload("exa", []webSearchHit{ + {Title: "No text", URL: "https://example.com/empty", Content: " "}, + {Title: "No url", URL: "", Content: "orphan text"}, + {Title: "RAGFlow", URL: "https://example.com/ragflow", Content: " An open-source RAG engine. "}, + }) + assertWebSearchChunk(t, payload, "exa-https://example.com/ragflow", "RAGFlow", "An open-source RAG engine.") +} diff --git a/internal/tokenizer/asset_dir_test.go b/internal/tokenizer/asset_dir_test.go new file mode 100644 index 0000000000..bcdc6c3e1f --- /dev/null +++ b/internal/tokenizer/asset_dir_test.go @@ -0,0 +1,152 @@ +// +// 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 + +// MODEL_ASSETS_DIR is how a deployment points at a mounted model-asset tree instead of +// relying on the working directory. The test runs in a child process with the working +// directory moved to a temporary tree, because a counter that a test in this process +// already loaded would otherwise hide the effect of the environment variable. + +import ( + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + + "ragflow/internal/common" +) + +const assetDirChildEnv = "RAGFLOW_TOKENIZER_ASSET_DIR_CHILD" + +// realSPMAsset is the asset shipped by ragflow_deps/download_go_deps.py, relative to the +// repository root. +var realSPMAsset = filepath.Join("ragflow_deps", "huggingface.co", "BAAI", "bge-m3", "sentencepiece.bpe.model") + +func TestModelAssetsDirIsHonoured(t *testing.T) { + if os.Getenv(assetDirChildEnv) != "" { + t.Skip("child process") + } + source := filepath.Join("..", "..", realSPMAsset) + if _, err := os.Stat(source); err != nil { + t.Skipf("the SPM asset is not present here (%v); run ragflow_deps/download_go_deps.py", err) + } + + dir := t.TempDir() + target := filepath.Join(dir, filepath.FromSlash("huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model")) + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + copyAssetFile(t, source, target) + + cmd := exec.Command(os.Args[0], "-test.run=^TestModelAssetsDirIsHonouredChild$", "-test.v") + cmd.Dir = dir // nothing to find by walking up from here + cmd.Env = append(os.Environ(), + assetDirChildEnv+"=1", + common.EnvModelAssetsDir+"="+dir, + ) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("child failed: %v\n%s", err, out) + } + if !strings.Contains(string(out), "asset dir child: ok") { + t.Fatalf("child did not confirm the override:\n%s", out) + } +} + +// TestModelAssetsDirIsHonoured is the child half: with MODEL_ASSETS_DIR pointing at the +// temporary tree and an unrelated working directory, the counter must still load, and +// report where it loaded from. +func TestModelAssetsDirIsHonouredChild(t *testing.T) { + if os.Getenv(assetDirChildEnv) == "" { + t.Skip("parent process") + } + dir, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + if !CounterExact(CounterXLMRSentence) { + t.Fatalf("%s=%s was not honoured: the xlmr-spm counter is unavailable", common.EnvModelAssetsDir, dir) + } + counter, ok := CounterByID(CounterXLMRSentence) + if !ok { + t.Fatal("xlmr-spm missing after a successful load check") + } + pathAware, ok := counter.(interface{ SourcePath() string }) + if !ok { + t.Fatal("the counter does not report its source path") + } + source := pathAware.SourcePath() + if !strings.HasPrefix(source, dir) { + t.Fatalf("counter loaded from %q, expected it under %q", source, dir) + } + if !strings.Contains(source, filepath.Join("huggingface.co", "BAAI", "bge-m3")) { + t.Fatalf("unexpected asset layout in %q", source) + } + for _, status := range CounterStatuses() { + if status.ID == CounterXLMRSentence { + if !status.Available || status.Source != source { + t.Fatalf("status for xlmr-spm = %+v, want available from %q", status, source) + } + t.Logf("asset dir child: ok (loaded %s)", source) + return + } + } + t.Fatal("xlmr-spm missing from CounterStatuses") +} + +func TestCounterStatusesCoverEveryCounter(t *testing.T) { + statuses := CounterStatuses() + want := []string{CounterCL100K, CounterXLMRSentence, CounterBERTWordPiece, CounterQwenBPE, CounterLlamaBPE} + seen := make(map[string]CounterStatus, len(statuses)) + for _, status := range statuses { + seen[status.ID] = status + } + for _, id := range want { + status, ok := seen[id] + if !ok { + t.Errorf("CounterStatuses is missing %s", id) + continue + } + // The report is what the startup log prints, so it must agree with the + // question the ingest path asks. + if status.Available != CounterExact(id) { + t.Errorf("%s: status.Available=%v but CounterExact=%v", id, status.Available, CounterExact(id)) + } + if status.Available && status.Source == "" { + t.Logf("note: %s is available but reports no source path", id) + } + } +} + +func copyAssetFile(t *testing.T, source, target string) { + t.Helper() + in, err := os.Open(source) + if err != nil { + t.Fatalf("open %s: %v", source, err) + } + defer in.Close() + out, err := os.Create(target) + if err != nil { + t.Fatalf("create %s: %v", target, err) + } + defer out.Close() + if _, err := io.Copy(out, in); err != nil { + t.Fatalf("copy %s: %v", source, err) + } +} diff --git a/internal/tokenizer/bpe.go b/internal/tokenizer/bpe.go new file mode 100644 index 0000000000..ac9702f9ab --- /dev/null +++ b/internal/tokenizer/bpe.go @@ -0,0 +1,860 @@ +// +// 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 + +// Byte-pair-encoding counters: the Qwen and Llama/Mistral embedding families. +// +// Two flavours share one merge engine, because they differ only in how text is +// prepared before merging. Both are read from the model's `tokenizer.json` +// (vocab + merges), and both are verified sample-by-sample against that same +// file through scripts/gen_tokenizer_oracle.py: +// +// llama-bpe (e5-mistral, Mistral/Llama SentencePiece-BPE) +// normalizer Sequence(Prepend("▁"), Replace(" " -> "▁")) +// pre_tokenizer none - the whole text is one merge unit +// model BPE, byte_fallback=true, fuse_unk=true, unk="" +// +// qwen-bpe (Qwen3-Embedding) +// normalizer NFC +// pre_tokenizer Sequence(Split(, isolated), ByteLevel(use_regex=false)) +// model BPE over the GPT-2 byte alphabet, no unk, no byte_fallback +// +// The Qwen pre-tokenizer regex cannot be compiled by Go's regexp (RE2): it uses a +// negative lookahead (`\s+(?!\S)`). It is implemented as an explicit scanner in +// pretokenizeQwen below, which is what the oracle test exists to keep honest. + +import ( + "container/heap" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +// normNFC is the normalizer step Qwen's tokenizer.json declares. +func normNFC(s string) string { return norm.NFC.String(s) } + +// mergeCacheLimit bounds the per-piece memo. Pieces repeat heavily across chunks +// ("▁the", "▁and", ...), so a small cache removes most of the merge work. +const mergeCacheLimit = 8192 + +// bpeFlavor selects the text preparation used before merging. +type bpeFlavor int + +const ( + // bpeFlavorLlama is SentencePiece-BPE: prepend+replace to ▁, merge the whole + // text, unknown runes fall back to <0xXX> byte pieces. + bpeFlavorLlama bpeFlavor = iota + // bpeFlavorQwen is byte-level BPE: NFC, split by regex, map bytes through the + // GPT-2 alphabet, then merge inside each split piece. + bpeFlavorQwen +) + +// bpeModel is a tokenizer.json BPE model plus the flavor it is served with. +type bpeModel struct { + vocab map[string]int32 + merges map[string]int + unk string + + byteFallback bool + byteLevel bool + flavor bpeFlavor + + mu sync.Mutex + memo map[string]int + + // names is the id -> token map, built on first use (only the oracle test asks + // for token texts). + nameOnce sync.Once + names map[int32]string +} + +// tokenizerJSON is the slice of tokenizer.json this package reads. Everything +// else (post_processors, decoders) cannot change a token count. +type tokenizerJSON struct { + Model struct { + Type string `json:"type"` + Vocab map[string]int32 `json:"vocab"` + Merges json.RawMessage `json:"merges"` + UnkToken *string `json:"unk_token"` + ByteFallback bool `json:"byte_fallback"` + } `json:"model"` + AddedTokens []struct { + Content string `json:"content"` + } `json:"added_tokens"` +} + +// parseMerges accepts both spellings seen in the wild: an array of "a b" strings +// (HF) and an array of ["a","b"] pairs. +func parseMerges(raw json.RawMessage) (map[string]int, error) { + if len(raw) == 0 { + return map[string]int{}, nil + } + var asStrings []string + if err := json.Unmarshal(raw, &asStrings); err == nil { + out := make(map[string]int, len(asStrings)) + for i, m := range asStrings { + out[m] = i + } + return out, nil + } + var asPairs [][]string + if err := json.Unmarshal(raw, &asPairs); err != nil { + return nil, fmt.Errorf("bpe: unsupported merges format: %w", err) + } + out := make(map[string]int, len(asPairs)) + for i, pair := range asPairs { + if len(pair) != 2 { + return nil, fmt.Errorf("bpe: merge %d has %d parts", i, len(pair)) + } + out[pair[0]+" "+pair[1]] = i + } + return out, nil +} + +func loadBPEModel(path string, flavor bpeFlavor) (*bpeModel, error) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("bpe: reading %s: %w", path, err) + } + var parsed tokenizerJSON + if err := json.Unmarshal(raw, &parsed); err != nil { + return nil, fmt.Errorf("bpe: parsing %s: %w", path, err) + } + if !strings.EqualFold(parsed.Model.Type, "BPE") { + return nil, fmt.Errorf("bpe: %s declares model type %q, want BPE", path, parsed.Model.Type) + } + merges, err := parseMerges(parsed.Model.Merges) + if err != nil { + return nil, err + } + if len(parsed.Model.Vocab) == 0 { + return nil, fmt.Errorf("bpe: %s has an empty vocabulary", path) + } + m := &bpeModel{ + vocab: parsed.Model.Vocab, + merges: merges, + byteFallback: parsed.Model.ByteFallback, + byteLevel: flavor == bpeFlavorQwen, + flavor: flavor, + memo: make(map[string]int), + } + if parsed.Model.UnkToken != nil { + m.unk = *parsed.Model.UnkToken + } + return m, nil +} + +// normalize prepares text the way the served tokenizer does. +func (m *bpeModel) normalize(text string) string { + if text == "" { + return "" + } + switch m.flavor { + case bpeFlavorQwen: + // normalizer: NFC + return normNFC(text) + default: + // normalizer: Sequence(Prepend("▁"), Replace(" " -> "▁")) + s := strings.ReplaceAll(text, " ", "\u2581") + return "\u2581" + s + } +} + +// countTokens counts the merge units of text. +func (m *bpeModel) countTokens(text string) int { + if text == "" { + return 0 + } + normalized := m.normalize(text) + if normalized == "" { + return 0 + } + pieces := []string{normalized} + if m.flavor == bpeFlavorQwen { + pieces = pretokenizeQwen(normalized) + } + total := 0 + for _, piece := range pieces { + total += m.countPiece(piece) + } + return total +} + +// countPiece is the BPE merge loop for one unit. +func (m *bpeModel) countPiece(piece string) int { + if piece == "" { + return 0 + } + if cached, ok := m.memoGet(piece); ok { + return cached + } + var symbols []string + if m.byteLevel { + // Byte-level BPE starts from the GPT-2 alphabet, one symbol per byte. + for i := 0; i < len(piece); i++ { + symbols = append(symbols, string(byteToRune[piece[i]])) + } + } else { + for _, r := range piece { + symbols = append(symbols, string(r)) + } + } + if len(symbols) > 1 { + symbols = m.merge(symbols) + } + count := 0 + for _, sym := range symbols { + if _, ok := m.vocab[sym]; ok { + count++ + continue + } + // A symbol the vocabulary does not contain: byte fallback emits one + // <0xXX> piece per UTF-8 byte of the symbol, otherwise it is a single + // unknown token. + if m.byteFallback && len(sym) > 0 { + count += len(sym) + continue + } + count++ + } + m.memoPut(piece, count) + return count +} + +// tokenIDs returns the exact token ids, for the oracle test's id-level comparison. +// It deliberately does not use the count memo: the memo stores numbers, and a +// count-neutral segmentation bug is exactly what the id comparison exists to find. +func (m *bpeModel) tokenIDs(text string) []int32 { + if text == "" { + return nil + } + normalized := m.normalize(text) + if normalized == "" { + return nil + } + pieces := []string{normalized} + if m.flavor == bpeFlavorQwen { + pieces = pretokenizeQwen(normalized) + } + var ids []int32 + for _, piece := range pieces { + ids = append(ids, m.pieceIDs(piece)...) + } + return ids +} + +// tokenNames returns the token text of every id, which is what the served +// tokenizer reports and what the oracle test compares. +func (m *bpeModel) tokenNames(text string) []string { + ids := m.tokenIDs(text) + if len(ids) == 0 { + return nil + } + m.nameOnce.Do(func() { + m.names = make(map[int32]string, len(m.vocab)) + for token, id := range m.vocab { + m.names[id] = token + } + }) + names := make([]string, 0, len(ids)) + for _, id := range ids { + names = append(names, m.names[id]) + } + return names +} + +// pieceIDs is countPiece with the ids kept: same symbol construction, same merge +// loop, same byte fallback. +func (m *bpeModel) pieceIDs(piece string) []int32 { + if piece == "" { + return nil + } + var symbols []string + if m.byteLevel { + for i := 0; i < len(piece); i++ { + symbols = append(symbols, string(byteToRune[piece[i]])) + } + } else { + for _, r := range piece { + symbols = append(symbols, string(r)) + } + } + if len(symbols) > 1 { + symbols = m.merge(symbols) + } + ids := make([]int32, 0, len(symbols)) + for _, sym := range symbols { + if id, ok := m.vocab[sym]; ok { + ids = append(ids, id) + continue + } + if m.byteFallback && len(sym) > 0 { + if m.byteLevel { + // A byte-level vocabulary contains every byte symbol, so this + // branch is unreachable for the shipped byte-level models; map + // back to the original byte if it ever is reached. + if r, ok := singleRune(sym); ok { + if b, ok := runeToByte[r]; ok { + ids = append(ids, m.bytePieceID(b)) + continue + } + } + ids = append(ids, m.unkID()) + continue + } + // Non-byte-level fallback: one <0xXX> token per UTF-8 byte. + for i := 0; i < len(sym); i++ { + ids = append(ids, m.bytePieceID(sym[i])) + } + continue + } + ids = append(ids, m.unkID()) + } + return ids +} + +// bytePieceID is the id of the "<0xXX>" token for one byte, or unk when the +// vocabulary has no byte-fallback token for it. +func (m *bpeModel) bytePieceID(b byte) int32 { + if id, ok := m.vocab[fmt.Sprintf("<0x%02X>", b)]; ok { + return id + } + return m.unkID() +} + +// unkID is the id of the vocabulary's unknown token, or -1 when it has none. +func (m *bpeModel) unkID() int32 { + if m.unk != "" { + if id, ok := m.vocab[m.unk]; ok { + return id + } + } + return -1 +} + +// singleRune reports the only rune of s, when s holds exactly one. +func singleRune(s string) (rune, bool) { + r, size := utf8.DecodeRuneInString(s) + if size == 0 || size != len(s) { + return 0, false + } + return r, true +} + +// merge applies BPE merges in the canonical order: repeatedly take the +// lowest-rank adjacent pair and merge EVERY occurrence of it before moving to the +// next rank. Ranks are unique per pair, so "every occurrence of the lowest rank" +// is well defined. +// +// Two implementations were tried and rejected before this one, both caught by the +// oracle test: +// +// - "merge one occurrence, then re-evaluate": lets a newly created pair be +// merged before the remaining occurrences of the current pair, which changes +// the segmentation. On a run of 32 'a's it returns 9 tokens where the model +// returns 7. +// - "rescan the whole unit for the lowest rank each round": correct but O(n^2) +// — a 78 KB chunk (the Mistral flavour merges a whole chunk as one unit) took +// nine seconds. +// +// This version keeps the symbols in a doubly linked list and the candidate pairs +// in a min-heap keyed by rank, and drains the heap per rank: O(n log n). Stale +// entries are detected with per-symbol version counters, and candidates that +// belong to a LOWER rank than the round in progress are deferred until the round +// ends, which is what makes the result identical to the batch-per-rank definition. +func (m *bpeModel) merge(symbols []string) []string { + n := len(symbols) + if n < 2 { + return symbols + } + prev := make([]int32, n) + next := make([]int32, n) + version := make([]int32, n) + alive := make([]bool, n) + for i := 0; i < n; i++ { + prev[i] = int32(i - 1) + next[i] = int32(i + 1) + alive[i] = true + } + next[n-1] = -1 + + h := make(mergeHeap, 0, n) + pushPair := func(left int32) { + if left < 0 { + return + } + right := next[left] + if right < 0 { + return + } + if rank, ok := m.merges[symbols[left]+" "+symbols[right]]; ok { + heap.Push(&h, mergeEntry{rank: rank, left: left, right: right, leftVer: version[left], rightVer: version[right]}) + } + } + for i := int32(0); i < int32(n); i++ { + pushPair(i) + } + heap.Init(&h) + + roundRank := -1 + var deferred []mergeEntry + for h.Len() > 0 || len(deferred) > 0 { + if h.Len() == 0 { + // The heap drained while this round's cheaper-than-roundRank candidates + // were still deferred, so nothing is left to trigger them. Re-queue them + // instead of dropping them: a dropped candidate leaves the piece split + // into more symbols than the merge table allows, which over-counts - the + // direction that makes a provider reject a request. A training-consistent + // table never reaches this state (a merged pair's rank is always higher + // than the rank that created it); a hand-written or converted one can. + for _, d := range deferred { + heap.Push(&h, d) + } + deferred = deferred[:0] + roundRank = -1 + } + entry := heap.Pop(&h).(mergeEntry) + if !alive[entry.left] || !alive[entry.right] || + version[entry.left] != entry.leftVer || version[entry.right] != entry.rightVer || + next[entry.left] != entry.right { + continue + } + if roundRank >= 0 { + if entry.rank > roundRank { + // The current rank is exhausted; this candidate starts the next + // round. Re-queue it along with anything deferred from this round. + heap.Push(&h, entry) + for _, d := range deferred { + heap.Push(&h, d) + } + deferred = deferred[:0] + roundRank = -1 + continue + } + if entry.rank < roundRank { + // A merge created a pair that is cheaper than the round in + // progress. The batch definition says the round finishes first. + deferred = append(deferred, entry) + continue + } + } + roundRank = entry.rank + + // Merge right into left. + symbols[entry.left] += symbols[entry.right] + version[entry.left]++ + alive[entry.right] = false + rightNext := next[entry.right] + next[entry.left] = rightNext + if rightNext >= 0 { + prev[rightNext] = entry.left + } + pushPair(prev[entry.left]) + pushPair(entry.left) + } + + out := make([]string, 0, n) + for i := 0; i < n; i++ { + if alive[i] && prev[i] < 0 { + for cur := int32(i); cur >= 0; cur = next[cur] { + out = append(out, symbols[cur]) + } + break + } + } + if len(out) == 0 { + // Defensive: never report zero tokens for non-empty symbols. + return symbols + } + return out +} + +// mergeEntry is one candidate merge of the adjacent symbols left|right, carrying +// the symbol versions it was computed from so staleness is detectable. +type mergeEntry struct { + rank int + left, right int32 + leftVer, rightVer int32 +} + +// mergeHeap is a min-heap over mergeEntry: by rank, then by position. +// +// The position tie-break is not cosmetic. A rank identifies a single pair, and +// every occurrence of it has that same rank, so the heap must hand them back +// leftmost-first: that is what the batch definition does (one left-to-right pass +// merging non-overlapping occurrences). Ordering by rank alone lets a middle +// occurrence be merged first, which then makes both of its neighbours +// unmergeable. On "aaaa" that yields [a, aa, a] instead of [aa, aa]. +type mergeHeap []mergeEntry + +func (h mergeHeap) Len() int { return len(h) } +func (h mergeHeap) Less(i, j int) bool { + if h[i].rank != h[j].rank { + return h[i].rank < h[j].rank + } + return h[i].left < h[j].left +} +func (h mergeHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *mergeHeap) Push(x any) { *h = append(*h, x.(mergeEntry)) } + +func (h *mergeHeap) Pop() any { + old := *h + n := len(old) + item := old[n-1] + *h = old[:n-1] + return item +} + +func (m *bpeModel) memoGet(piece string) (int, bool) { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.memo[piece] + return v, ok +} + +func (m *bpeModel) memoPut(piece string, count int) { + m.mu.Lock() + defer m.mu.Unlock() + if len(m.memo) >= mergeCacheLimit { + // Cheap eviction: the cache is an optimisation, not a correctness + // structure, so clearing it beats tracking recency. + m.memo = make(map[string]int, mergeCacheLimit) + } + m.memo[piece] = count +} + +// trimToLimit returns the longest rune prefix with at most limit tokens. +func (m *bpeModel) trimToLimit(text string, limit int) string { + if limit <= 0 { + return "" + } + if m.countTokens(text) <= limit { + return text + } + runes := []rune(text) + lo, hi := 0, len(runes) + for lo < hi { + mid := (lo + hi + 1) / 2 + if m.countTokens(string(runes[:mid])) <= limit { + lo = mid + } else { + hi = mid - 1 + } + } + return string(runes[:lo]) +} + +// --------------------------------------------------------------------------- +// Qwen pre-tokenization +// --------------------------------------------------------------------------- + +// pretokenizeQwen reproduces the Split behavior (behavior="isolated") of the +// Qwen pre-tokenizer regex: +// +// (?i:'s|'t|'re|'ve|'m|'ll|'d) +// |[^\r\n\p{L}\p{N}]?\p{L}+ +// |\p{N} +// | ?[^\s\p{L}\p{N}]+[\r\n]* +// |\s*[\r\n]+ +// |\s+(?!\S) +// |\s+ +// +// Alternatives are tried in order at each position, exactly like the regex +// engine's leftmost-first match, and every match is its own piece. +func pretokenizeQwen(s string) []string { + var out []string + rs := []rune(s) + i := 0 + for i < len(rs) { + if piece, next, ok := matchQwenAlternative(rs, i); ok && next > i { + out = append(out, piece) + i = next + continue + } + // No alternative matched (cannot happen for the pattern above, but be + // safe): emit the character alone so counting never stalls. + out = append(out, string(rs[i])) + i++ + } + return out +} + +func matchQwenAlternative(rs []rune, i int) (string, int, bool) { + // 1. contractions, case-insensitive + for _, suffix := range []string{"'s", "'t", "'re", "'ve", "'m", "'ll", "'d"} { + if hasRunePrefixFold(rs[i:], suffix) { + return string(rs[i : i+len([]rune(suffix))]), i + len([]rune(suffix)), true + } + } + // 2. [^\r\n\p{L}\p{N}]?\p{L}+ + { + start := i + if !isLetter(rs[start]) && rs[start] != '\r' && rs[start] != '\n' && !isNumber(rs[start]) { + start++ + } + if start < len(rs) && isLetter(rs[start]) { + end := start + for end < len(rs) && isLetter(rs[end]) { + end++ + } + return string(rs[i:end]), end, true + } + } + // 3. \p{N} - exactly one digit + if isNumber(rs[i]) { + return string(rs[i]), i + 1, true + } + // 4. " ?[^\s\p{L}\p{N}]+[\r\n]*" + { + start := i + if rs[start] == ' ' { + start++ + } + if start < len(rs) && !isSpace(rs[start]) && !isLetter(rs[start]) && !isNumber(rs[start]) { + end := start + for end < len(rs) && !isSpace(rs[end]) && !isLetter(rs[end]) && !isNumber(rs[end]) { + end++ + } + for end < len(rs) && (rs[end] == '\r' || rs[end] == '\n') { + end++ + } + return string(rs[i:end]), end, true + } + } + // 5. \s*[\r\n]+ - whitespace ENDING in a newline. Greedy \s* then >=1 + // newline: the match therefore stops at the LAST newline of the run, leaving + // any trailing spaces for the next alternative. That is why "a \n b" splits + // as [" \n", " b"] and not as [" \n ", "b"]. + { + end := i + for end < len(rs) && isSpace(rs[end]) { + end++ + } + lastNewline := -1 + for k := i; k < end; k++ { + if rs[k] == '\r' || rs[k] == '\n' { + lastNewline = k + } + } + if lastNewline >= 0 { + return string(rs[i : lastNewline+1]), lastNewline + 1, true + } + } + // 6/7. Whitespace runs without a newline. + // + // "\s+(?!\S)" comes first and, being backtracking, matches the whole run + // MINUS its last character whenever something non-space follows: the + // lookahead then sees the remaining space and succeeds. The leftover space is + // picked up by the next alternative, which is how "a b" becomes + // ["a", " ", " b"] - with " b" costing one token instead of two. Fall back to + // the full run when "!"(\S) cannot hold, i.e. at the end of the text. + { + end := i + for end < len(rs) && isSpace(rs[end]) { + end++ + } + if end > i { + if end-i >= 2 && end < len(rs) && !isSpace(rs[end]) { + return string(rs[i : end-1]), end - 1, true + } + return string(rs[i:end]), end, true + } + } + return "", i, false +} + +func hasRunePrefixFold(rs []rune, suffix string) bool { + s := []rune(suffix) + if len(rs) < len(s) { + return false + } + for i, r := range s { + if unicode.ToLower(rs[i]) != unicode.ToLower(r) { + return false + } + } + return true +} + +func hasNewline(rs []rune) bool { + for _, r := range rs { + if r == '\r' || r == '\n' { + return true + } + } + return false +} + +func isLetter(r rune) bool { return unicode.IsLetter(r) } +func isNumber(r rune) bool { return unicode.IsNumber(r) } +func isSpace(r rune) bool { return unicode.IsSpace(r) } + +// byteToRune is the GPT-2 byte -> unicode alphabet used by byte-level BPE. +var byteToRune = buildByteToRune() + +// runeToByte is its inverse, needed to report the original byte when a byte-level +// symbol has no token of its own. +var runeToByte = buildRuneToByte() + +func buildRuneToByte() map[rune]byte { + m := make(map[rune]byte, len(byteToRune)) + for b, r := range byteToRune { + m[r] = byte(b) + } + return m +} + +func buildByteToRune() [256]rune { + var bs, cs []int + for i := '!'; i <= '~'; i++ { + bs = append(bs, int(i)) + } + for i := 0xA1; i <= 0xAC; i++ { + bs = append(bs, int(i)) + } + for i := 0xAE; i <= 0xFF; i++ { + bs = append(bs, int(i)) + } + cs = append(cs, bs...) + extra := 0 + for b := 0; b < 256; b++ { + if containsInt(bs, b) { + continue + } + bs = append(bs, b) + cs = append(cs, 256+extra) + extra++ + } + var table [256]rune + for i := range bs { + table[bs[i]] = rune(cs[i]) + } + return table +} + +func containsInt(xs []int, v int) bool { + for _, x := range xs { + if x == v { + return true + } + } + return false +} + +// --------------------------------------------------------------------------- +// Counters +// --------------------------------------------------------------------------- + +type bpeCounter struct { + model *bpeModel + path string + counterID string +} + +func (c *bpeCounter) ID() string { return c.counterID } + +// IDs reports the token ids, for the oracle test's id-level comparison. +func (c *bpeCounter) IDs(text string) []int32 { return c.model.tokenIDs(text) } + +// Names reports the token text of every token, which the oracle test compares +// directly with the served tokenizer's tokens. +func (c *bpeCounter) Names(text string) []string { return c.model.tokenNames(text) } + +// SourcePath is the tokenizer.json this counter loaded. +func (c *bpeCounter) SourcePath() string { return c.path } + +func (c *bpeCounter) Count(text string) int { return c.model.countTokens(text) } + +func (c *bpeCounter) TrimToLimit(text string, limit int) string { + return c.model.trimToLimit(text, limit) +} + +func (c *bpeCounter) Available() bool { return c.model != nil } + +// Pinned digests of the shipped tokenizer.json assets. +var expectedBPEHashes = map[string]string{ + "Qwen/Qwen3-Embedding-0.6B/tokenizer.json": "e6592f4d8d1678e963da9188090ac3eee916ed6b", + "intfloat/e5-mistral-7b-instruct/tokenizer.json": "92cb22e16e9be5e5d24851a89b3f912cbfcda26b", +} + +const ( + llamaBPEAssetPinKey = "intfloat/e5-mistral-7b-instruct/tokenizer.json" + qwenBPEAssetPinKey = "Qwen/Qwen3-Embedding-0.6B/tokenizer.json" +) + +var llamaBPEAssetNames = []string{ + "ragflow_deps/huggingface.co/intfloat/e5-mistral-7b-instruct/tokenizer.json", + "llama_bpe_tokenizer.json", +} + +var qwenBPEAssetNames = []string{ + "ragflow_deps/huggingface.co/Qwen/Qwen3-Embedding-0.6B/tokenizer.json", + "qwen_bpe_tokenizer.json", +} + +var ( + llamaBPEOnce struct { + sync.Once + counter Counter + err error + } + qwenBPEOnce struct { + sync.Once + counter Counter + err error + } +) + +// LoadLlamaBPECounter loads (once) the Llama/Mistral BPE counter. +func LoadLlamaBPECounter() (Counter, error) { + llamaBPEOnce.Do(func() { + llamaBPEOnce.counter, llamaBPEOnce.err = loadBPECounter(llamaBPEAssetNames, llamaBPEAssetPinKey, bpeFlavorLlama, CounterLlamaBPE) + }) + return llamaBPEOnce.counter, llamaBPEOnce.err +} + +// LoadQwenBPECounter loads (once) the Qwen byte-level BPE counter. +func LoadQwenBPECounter() (Counter, error) { + qwenBPEOnce.Do(func() { + qwenBPEOnce.counter, qwenBPEOnce.err = loadBPECounter(qwenBPEAssetNames, qwenBPEAssetPinKey, bpeFlavorQwen, CounterQwenBPE) + }) + return qwenBPEOnce.counter, qwenBPEOnce.err +} + +func loadBPECounter(names []string, pinKey string, flavor bpeFlavor, counterID string) (Counter, error) { + path, _, err := readTokenizerAsset(names, pinKey, expectedBPEHashes) + if err != nil { + return nil, err + } + model, err := loadBPEModel(path, flavor) + if err != nil { + return nil, err + } + return &bpeCounter{model: model, path: path, counterID: counterID}, nil +} + +func init() { + RegisterCounterLoader(CounterLlamaBPE, LoadLlamaBPECounter) + RegisterCounterLoader(CounterQwenBPE, LoadQwenBPECounter) +} diff --git a/internal/tokenizer/bpe_loader.go b/internal/tokenizer/bpe_loader.go index 80ad8eeb85..17eae01eae 100644 --- a/internal/tokenizer/bpe_loader.go +++ b/internal/tokenizer/bpe_loader.go @@ -35,6 +35,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "ragflow/internal/common" @@ -51,11 +52,11 @@ func init() { // 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) +// filename. The loader names its cached file after the URL // (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`. +// ragflow_deps/download_go_deps.py: `sha1sum cl100k_base.tiktoken`. var expectedBpeHashes = map[string]string{ "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken": "6494e42d5aad2bbb441ea9793af9e7db335c8d9c", } @@ -63,6 +64,26 @@ var expectedBpeHashes = map[string]string{ // localBpeLoader resolves tiktoken BPE tables from the local filesystem. type localBpeLoader struct{} +// cl100kTablePath records the table file the loader accepted, for diagnostics: the +// availability report (CounterStatuses) prints it, which is how an operator sees which +// copy of the table is in use. +var ( + cl100kTablePathMu sync.Mutex + cl100kTablePath string +) + +func recordCl100kTablePath(path string) { + cl100kTablePathMu.Lock() + cl100kTablePath = path + cl100kTablePathMu.Unlock() +} + +func cl100kTableSource() string { + cl100kTablePathMu.Lock() + defer cl100kTablePathMu.Unlock() + return cl100kTablePath +} + // LoadTiktokenBpe implements tiktoken.BpeLoader. // // bpeURL is the upstream table URL that tiktoken-go would otherwise download; @@ -96,6 +117,7 @@ func (localBpeLoader) LoadTiktokenBpe(bpeURL string) (map[string]int, error) { // name collision. Continuing to the next candidate would mask it. return nil, fmt.Errorf("BPE table %s is malformed: %w", candidate, err) } + recordCl100kTablePath(candidate) return ranks, nil } @@ -138,6 +160,12 @@ func bpeCandidatePaths(bpeURL string) []string { } } + // MODEL_ASSETS_DIR: the shared model-asset root (see common.ModelAssetCandidates). + // Listed before the working-directory walk because it is an explicit instruction. + for _, candidate := range common.ModelAssetCandidates(bundledName) { + add(candidate) + } + for _, root := range bpeSearchRoots() { // Same layout the Dockerfile creates: the table sits in the // installation root under its sha1 name. diff --git a/internal/tokenizer/bpe_merge_test.go b/internal/tokenizer/bpe_merge_test.go new file mode 100644 index 0000000000..3a7421683f --- /dev/null +++ b/internal/tokenizer/bpe_merge_test.go @@ -0,0 +1,44 @@ +package tokenizer + +import ( + "reflect" + "testing" +) + +// The BPE batch loop defers candidates that are cheaper than the round in progress +// (the batch definition finishes the round first). A deferred candidate must never +// be LOST when the heap drains: dropping one leaves the piece split into more +// symbols than the merge table allows, i.e. an over-count — the direction that makes +// a provider reject a request. +// +// The table below is deliberately NOT training-consistent (the rank-0 merge needs a +// symbol that only the rank-5 merge creates). A well-formed table cannot reach this +// state, which is exactly why it is easy to lose: the loop used to exit with the +// candidate still deferred and segment "abc" as two tokens. +func TestBpeMergeDoesNotLoseDeferredCandidatesWhenTheHeapDrains(t *testing.T) { + m := &bpeModel{merges: map[string]int{ + "a b": 5, + "ab c": 0, + }} + + got := m.merge([]string{"a", "b", "c"}) + want := []string{"abc"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("deferred merge was dropped: got %v, want %v", got, want) + } +} + +// Sanity counterpart: with a training-consistent table the deferred path is never +// taken, so the ordinary cheapest-pair-first behaviour is unchanged. +func TestBpeMergePicksTheCheapestPairFirst(t *testing.T) { + m := &bpeModel{merges: map[string]int{ + "a b": 0, + "ab c": 1, + }} + + got := m.merge([]string{"a", "b", "c"}) + want := []string{"abc"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("got %v, want %v", got, want) + } +} diff --git a/internal/tokenizer/counter_oracle_test.go b/internal/tokenizer/counter_oracle_test.go new file mode 100644 index 0000000000..4ad6d4c019 --- /dev/null +++ b/internal/tokenizer/counter_oracle_test.go @@ -0,0 +1,319 @@ +// +// 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 ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + "testing" +) + +// This file is the offline half of the verification plan in +// internal/tokenizer/embedding_token_limits.md: every hand-written counter must +// agree, sample by sample, with the tokenizer the model is actually served with. +// +// Fixtures come from scripts/gen_tokenizer_oracle.py, which tokenizes a fixed +// corpus with the HuggingFace `tokenizers` library reading the model's own +// tokenizer.json: +// +// mkdir -p /tmp/tokenizer_oracle +// for m in bge-m3 bge-large-en qwen3-embedding e5-mistral; do +// uv run scripts/gen_tokenizer_oracle.py --model $m --corpus --out /tmp/tokenizer_oracle/$m.json +// done +// TOKENIZER_ORACLE_DIR=/tmp/tokenizer_oracle ./build.sh --test -run TestCountersMatchOracle ./internal/tokenizer/ +// +// The test reads every fixture in the directory, maps fixture -> counter id via +// the fixture's own "counter" field, and fails on any count mismatch: a +// hand-written tokenizer that is "close" is not good enough, because an +// under-count lets an oversized input through and the provider answers 400. + +type oracleFixture struct { + Model string `json:"model"` + Counter string `json:"counter"` + Oracle string `json:"oracle"` + Reference string `json:"reference"` + Crosscheck string `json:"crosscheck"` + Samples []struct { + Label string `json:"label"` + Text string `json:"text"` + Count int `json:"count"` + Chars int `json:"chars"` + IDs []int32 `json:"ids"` + Tokens []string `json:"tokens"` + KnownApprox bool `json:"known_approx"` + } `json:"samples"` +} + +// idCounter is implemented by the counters that can report the exact token ids. The +// oracle test uses it to compare segmentations; the Count interface alone only +// supports comparing their lengths, which two different segmentations can share. +type idCounter interface { + IDs(text string) []int32 +} + +// nameCounter is implemented by the counters that can report the token texts. The +// test prefers them over ids because a token id is only meaningful inside one +// numbering: the HuggingFace conversion of bge-m3 reorders the SentencePiece +// specials and inserts /, so every regular piece is +1 relative to the +// .model file while the tokens themselves are identical. +type nameCounter interface { + Names(text string) []string +} + +func TestCountersMatchOracle(t *testing.T) { + dir := os.Getenv("TOKENIZER_ORACLE_DIR") + if dir == "" { + t.Skip("set TOKENIZER_ORACLE_DIR to a directory of fixtures from scripts/gen_tokenizer_oracle.py") + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("reading the oracle directory %s: %v", dir, err) + } + seen := 0 + for _, entry := range entries { + if entry.IsDir() || filepath.Ext(entry.Name()) != ".json" { + continue + } + path := filepath.Join(dir, entry.Name()) + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("reading the oracle fixture %s: %v", path, err) + } + var fx oracleFixture + if err := json.Unmarshal(raw, &fx); err != nil { + t.Fatalf("parsing the oracle fixture %s: %v", path, err) + } + seen++ + t.Run(fx.Counter, func(t *testing.T) { checkCounterAgainstOracle(t, fx) }) + } + if seen == 0 { + t.Fatalf("no fixtures found in %s", dir) + } +} + +func checkCounterAgainstOracle(t *testing.T, fx oracleFixture) { + counter, ok := CounterByID(fx.Counter) + if !ok { + // A fixture names the counter it was generated for, so a miss cannot be a + // reason to pass. Distinguish the two causes, because the fix differs: + // either this build does not have that counter at all, or it has it and the + // asset it needs is not on disk. + for _, status := range CounterStatuses() { + if status.ID == fx.Counter { + t.Fatalf("fixture %q needs counter %q, which is declared but unavailable: its asset is missing from ragflow_deps/huggingface.co (run `uv run ragflow_deps/download_go_deps.py`)", fx.Model, fx.Counter) + } + } + t.Fatalf("fixture %q names counter %q, which this build does not implement", fx.Model, fx.Counter) + } + ids, haveIDs := counter.(idCounter) + names, haveNames := counter.(nameCounter) + var ( + corpus string + corpusOracle int + mismatches int + idCompared int + nameCompared int + // Per sample, not per check: one sample is examined by both the count + // comparison and the token-text comparison, and the agreement line at the + // end must not subtract it twice. + approximated = map[string]bool{} + ) + for _, sample := range fx.Samples { + got := counter.Count(sample.Text) + if got != sample.Count { + // The one documented approximation (two or more combining marks on one + // base character) may move the count by a token or two; the fixture + // marks those samples explicitly. Everything else must be exact. + if sample.KnownApprox && abs(got-sample.Count) <= knownApproxTokenTolerance { + approximated[sample.Label] = true + } else { + mismatches++ + t.Errorf("%s: our count=%d, %s oracle=%d (chars=%d)", sample.Label, got, fx.Model, sample.Count, sample.Chars) + } + } + if sample.Label == corpusSampleLabel { + corpus, corpusOracle = sample.Text, sample.Count + } + if len(sample.IDs) != sample.Count { + mismatches++ + t.Errorf("%s: fixture is self-inconsistent: %d ids but count=%d", sample.Label, len(sample.IDs), sample.Count) + continue + } + // Our own count and our own segmentation must agree; if they ever do not, + // one of the two paths inside the counter has drifted from the other. + checkOwnAgreement := func(t *testing.T, label string, got int, n int) bool { + if n != got { + t.Helper() + mismatches++ + t.Errorf("%s: segmentation has %d tokens but count=%d inside our own counter", label, n, got) + return false + } + return true + } + if haveNames && len(sample.Tokens) > 0 { + ours := names.Names(sample.Text) + nameCompared++ + if !checkOwnAgreement(t, sample.Label, got, len(ours)) { + continue + } + // A base character with two or more combining marks is the one shape + // where our NFKC-based normalization provably differs from the served + // Precompiled charmap (see the fixture's known_approx flag, set by + // scripts/gen_tokenizer_oracle.py). The divergence is bounded here + // rather than excused: anything larger, or any divergence at all on a + // sample without that shape, still fails. + if sample.KnownApprox { + if delta := len(ours) - len(sample.Tokens); delta > knownApproxTokenTolerance || delta < -knownApproxTokenTolerance { + mismatches++ + t.Errorf("%s: known-approximation sample differs by %d tokens (tolerance %d): ours=%d oracle=%d", + sample.Label, delta, knownApproxTokenTolerance, len(ours), len(sample.Tokens)) + } else { + approximated[sample.Label] = true + } + continue + } + if at, ok := firstDivergence(ours, sample.Tokens); ok { + mismatches++ + t.Errorf("%s: tokens diverge at %d of %d\n ours %s\n oracle %s\n text %s", + sample.Label, at, len(sample.Tokens), describeWindow(ours, at), describeWindow(sample.Tokens, at), headRunes(sample.Text, 120)) + } + continue + } + if !haveIDs || len(sample.IDs) == 0 { + continue + } + ours := ids.IDs(sample.Text) + idCompared++ + if !checkOwnAgreement(t, sample.Label, got, len(ours)) { + continue + } + if at, ok := firstDivergence(ours, sample.IDs); ok { + mismatches++ + t.Errorf("%s: token ids diverge at %d of %d\n ours %s\n oracle %s\n text %s", + sample.Label, at, len(sample.IDs), describeWindow(ours, at), describeWindow(sample.IDs, at), headRunes(sample.Text, 120)) + } + } + mode := "counts" + switch { + case nameCompared > 0: + mode = "token texts" + case idCompared > 0: + mode = "id sequences" + } + t.Logf("agreement: %d/%d samples compared by %s (oracle: %s)", len(fx.Samples)-mismatches-len(approximated), len(fx.Samples), mode, fx.Oracle) + if len(approximated) > 0 { + t.Logf("within the documented approximation on %d/%d samples (combining marks in a shape the served charmap normalizes differently - see the known_approx flag in scripts/gen_tokenizer_oracle.py - tolerance %d tokens)", + len(approximated), len(fx.Samples), knownApproxTokenTolerance) + } + if fx.Crosscheck != "" { + t.Logf("second reference: %s", fx.Crosscheck) + } + + // The corpus, pinned as a test. The document passed in with --corpus is the one + // the incident failed on, and its whole content is sampled - so what can be + // pinned here is the window: the model's tokenizer sees MORE than the + // 8192-token limit, which is why this document has to take the exact-count path + // rather than a coarser one. The under-count that actually triggered 400/20015 + // lived in a 25k-character span of this document (cl100k below the limit while + // the model was above it); that shape belongs to the span, not to the document, + // and a whole-document sample cannot reproduce it - pass a corpus file cut to + // that span to pin it. The counters' agreement with their oracles on the whole + // document is what the loop above checks, token by token. + if fx.Counter != CounterXLMRSentence { + return + } + if corpus == "" { + t.Errorf("fixture has no %q sample: regenerate it with --corpus (scripts/gen_tokenizer_oracle.py)", corpusSampleLabel) + return + } + if corpusOracle <= 8192 { + t.Errorf("the corpus document counts %d tokens in %s, expected more than the 8192-token window", corpusOracle, fx.Model) + } + if cl100k := CountCL100K(); cl100k.Available() { + t.Logf("corpus document: cl100k=%d vs %s=%d", cl100k.Count(corpus), fx.Model, corpusOracle) + } +} + +// knownApproxTokenTolerance bounds the divergence allowed on the samples the fixture +// marks as the documented normalization approximation (a base character carrying two +// or more combining marks, where NFKC composes and the served charmap does not). +const knownApproxTokenTolerance = 2 + +// corpusSampleLabel is the fixture sample holding the whole document passed to +// scripts/gen_tokenizer_oracle.py --corpus. The incident check at the end of +// checkCounterAgainstOracle pins that sample, so a fixture generated without a +// corpus (the flag is required) cannot quietly drop the guarantee. +const corpusSampleLabel = "incident_full" + +func abs(v int) int { + if v < 0 { + return -v + } + return v +} + +// firstDivergence reports where two token sequences stop agreeing (including a pure +// length difference, which is a divergence at the shorter end). +func firstDivergence[T comparable](ours, oracle []T) (int, bool) { + for i := range ours { + if i >= len(oracle) || ours[i] != oracle[i] { + return i, true + } + } + if len(ours) != len(oracle) { + return len(ours), true + } + return 0, false +} + +// describeWindow renders a small window of tokens around a divergence; the marker +// points at the first token that differs. +func describeWindow[T any](values []T, at int) string { + lo, hi := at-3, at+4 + if lo < 0 { + lo = 0 + } + if hi > len(values) { + hi = len(values) + } + var b strings.Builder + b.WriteString("[") + for i := lo; i < hi; i++ { + if i > lo { + b.WriteString(" ") + } + if i == at { + b.WriteString("»") + } + fmt.Fprintf(&b, "%v", values[i]) + } + b.WriteString("]") + return b.String() +} + +// headRunes quotes the first n runes of text, for context in a failure message. +func headRunes(text string, n int) string { + runes := []rune(text) + if len(runes) <= n { + return strconv.Quote(string(runes)) + } + return strconv.Quote(string(runes[:n])) + "…" +} diff --git a/internal/tokenizer/counter_properties_test.go b/internal/tokenizer/counter_properties_test.go new file mode 100644 index 0000000000..4e0f093cec --- /dev/null +++ b/internal/tokenizer/counter_properties_test.go @@ -0,0 +1,141 @@ +// +// 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 + +// Counter property tests: the invariants every Counter must hold, checked on +// every counter whose asset is present in this environment. +// +// These are deliberately separate from TestCountersMatchOracle. The oracle test +// answers "is this the model's tokenizer?" and needs the 17 MB oracle fixtures; +// this file answers "is this a *well-behaved* counter?" and needs only the asset +// the counter itself loads. Both are required: a counter can match the oracle on +// 35 samples and still violate Count(TrimToLimit(t, n)) <= n for an input nobody +// sampled, and the ingest path relies on that inequality to stay under the +// provider's window. +// +// If a counter is skipped here, its asset is missing - not its correctness. + +import ( + "strings" + "testing" + "unicode/utf8" +) + +// propertyCorpus is one text per content class. The classes are the ones that +// have produced real bugs (see embedding_token_limits.md "one bug, one sample +// class"): prose, table rows, digits, CJK, emoji, base64-like, a long run of one +// character, and whitespace-heavy text. +func propertyCorpus() map[string]string { + return map[string]string{ + "prose": strings.Repeat("The quick brown fox jumps over the lazy dog. ", 60), + "table_rows": strings.Repeat("| 1976 | | 383/1 | 383/2 | 383/3 | 383/4 |\n", 40), + "digits": strings.Repeat("0123456789 ", 300), + "cjk": strings.Repeat("中文分词测试，用于对比 tokenizer。", 40), + "emoji": strings.Repeat("🚀🔥 embedding ✅ 测试 ", 40), + "base64_like": strings.Repeat("QWxhZGRpbjpvcGVuIHNlc2FtZQ", 100), + "one_char_run": strings.Repeat("a", 2000), + "whitespace": " padded text \t with \n\n edge whitespace ", + "accents": "e\u0301 a\u0300 o\u0308 café naïve résumé ", + "long_word": strings.Repeat("supercalifragilisticexpialidocious", 20), + } +} + +// propertyLimits are the budgets the ingest path realistically uses: the floor of +// the shrink ladder, small windows (the catalog has 512-token models), and a full +// 8192-token window minus the margin. +func propertyLimits() []int { return []int{1, 8, 32, 64, 512, 2048, 8028} } + +// TestCountersSatisfyTrimProperties holds every available counter to the same +// contract. Each property maps to something the ingest path depends on: +// +// prefix - L3 re-trims an already trimmed text, so trimming must be a +// prefix operation (never reorder, never re-add text) +// rune boundary - a trimmed chunk must remain valid UTF-8 +// count <= limit - the whole point: fit the model's window +// idempotence - the ladder and the migration path trim repeatedly +// monotone - a larger limit must never yield less text (the trim +// implementations binary-search prefixes, which assumes this) +func TestCountersSatisfyTrimProperties(t *testing.T) { + ids := []string{CounterCL100K, CounterXLMRSentence, CounterBERTWordPiece, CounterQwenBPE, CounterLlamaBPE} + corpus, limits := propertyCorpus(), propertyLimits() + ran := make([]string, 0, len(ids)) + for _, id := range ids { + counter, ok := CounterByID(id) + if !ok { + t.Logf("%s: skipped (asset not present in this environment)", id) + continue + } + ran = append(ran, id) + t.Run(id, func(t *testing.T) { + // An empty text is zero tokens for every family. This is not + // academic: the SPM counter returned 1 here before the dummy-prefix + // rule was fixed. + if got := counter.Count(""); got != 0 { + t.Errorf("Count(\"\") = %d, want 0", got) + } + if got := counter.TrimToLimit("anything at all", 0); got != "" { + t.Errorf("TrimToLimit(text, 0) = %q, want empty", got) + } + for name, text := range corpus { + prevCount, prevLen := -1, -1 + for _, limit := range limits { + trimmed := counter.TrimToLimit(text, limit) + if !strings.HasPrefix(text, trimmed) { + t.Fatalf("%s limit=%d: result is not a prefix of the input", name, limit) + } + if !utf8.ValidString(trimmed) || strings.ContainsRune(trimmed, utf8.RuneError) { + t.Fatalf("%s limit=%d: result is not valid UTF-8", name, limit) + } + count := counter.Count(trimmed) + if count > limit { + t.Fatalf("%s limit=%d: trimmed text still counts %d tokens", name, limit, count) + } + if again := counter.TrimToLimit(trimmed, limit); again != trimmed { + t.Fatalf("%s limit=%d: trimming is not idempotent (%d vs %d bytes)", name, limit, len(trimmed), len(again)) + } + if count < prevCount || len(trimmed) < prevLen { + t.Fatalf("%s limit=%d: a larger limit produced less text (%d tokens/len %d after %d/%d)", + name, limit, count, len(trimmed), prevCount, prevLen) + } + prevCount, prevLen = count, len(trimmed) + } + } + }) + } + t.Logf("counters exercised: %v", ran) +} + +// TestCountTrimFitsEveryLimit is the ingest-path-shaped version of the property +// above: take a chunk that is far too long, ask each counter to fit it into the +// same window the embedder would use, and assert the result fits. +func TestCountTrimFitsEveryLimit(t *testing.T) { + text := strings.Repeat("| 1976 | | 383/1 | 383/2 | 383/3 | 383/4 | 383/5 |\n", 500) + window := EmbeddingTokenLimit(8192) + for _, id := range []string{CounterCL100K, CounterXLMRSentence, CounterBERTWordPiece, CounterQwenBPE, CounterLlamaBPE} { + counter, ok := CounterByID(id) + if !ok { + continue + } + trimmed := counter.TrimToLimit(text, window) + if got := counter.Count(trimmed); got > window { + t.Errorf("%s: %d tokens after trimming to %d", id, got, window) + } + if len(trimmed) == 0 { + t.Errorf("%s: trimmed the whole chunk away", id) + } + } +} diff --git a/internal/tokenizer/embedding_token_limits.md b/internal/tokenizer/embedding_token_limits.md new file mode 100644 index 0000000000..9a0ec6c4d9 --- /dev/null +++ b/internal/tokenizer/embedding_token_limits.md @@ -0,0 +1,823 @@ +# Embedding input token limits + +How the ingest path decides how much text to hand an embedding model, and what it +does when the model disagrees. This is an implementation document: everything +described here is in the tree and verified, and every number quoted was measured. + +## Summary + +- **Invariant.** Text handed to an embedding API must be counted with the model's own + tokenizer, or with a calibrated upper bound of it; cl100k alone is never enough. And: + **one over-limit rejection must never fail a whole document.** +- **Layering.** L1 exact counters (registry + four family implementations) → L2 calibrated + fallback (learn the real/cl100k ratio from the provider's `usage`) → L3 shrink, then + isolate and retry, when a request is rejected. +- **Why it cannot be avoided.** Embedding providers *refuse* over-length input instead of + truncating it (bge-m3 rejects an input 0.55% over spec with `400`/`20015`, and one bad item + rejects the whole batch), the only token count they report comes back with a *successful* + response, and the tokenizer architectures disagree so much that cl100k under-counts bge-m3 by + up to 33% on our samples and the spread between estimators reaches 11x on ordinary text. + See [§Why this module has to exist](#why-this-module-has-to-exist). +- **How it is verified.** `scripts/gen_tokenizer_oracle.py` generates fixtures and + `TestCountersMatchOracle` compares every sample with the model's own tokenizer + (**4 families, the same 33 base samples plus `--fuzz`, all equal**); a second tier holds *every* counter to the same + contract (prefix, UTF-8 integrity, `Count(Trim) <= limit`, idempotence, monotonicity). +- **Normative basis.** Each counter is anchored to the model's own tokenizer file (SHA-1 + pinned, verified at load) and to the algorithm that file declares; **a family with no + oracle fixture is never tagged** — it falls back to L2 instead. See + [§How the implementation is kept honest](#how-the-implementation-is-kept-honest). +- **Result.** `78785.md`, which used to fail with `400`/`20015`, now ingests normally + (36 chunks, MySQL == ES) and **without touching L2/L3** — it takes the exact-count path. +- **Two authorities.** The `tokenizer` field in the model catalog `conf/all_models.json` + decides *which* counter is used; the model's own tokenizer file decides *how it counts* + (not the declarations inside that file's config). + +## Why this module has to exist + +Three measured properties of embedding services make "just send it and let the model deal with +it" impossible. All the numbers below are measured against the provider this deployment uses +(SILICONFLOW, `BAAI/bge-m3`, declared `max_tokens = 8192`). Together they are the reason for an +explicit pre-flight trim, per-model counters, and a calibrated fallback rather than one helper +function. + +### 1. Over-length input is refused, never truncated, so the caller must trim + +| Request | Result | +|---|---| +| 23,000 chars = **8,176** tokens (inside spec) | `200 OK`, `usage.prompt_tokens = 8176` | +| the incident span, 25,000 chars = **8,237** tokens (**0.55% over spec**) | `400`, `{"code":20015,"message":"The parameter is invalid. Please check again."}` | +| batch `["hello", ]` | `400` — the **whole batch**, not just the bad item | +| batch `["hello","world","ok"]` | `200 OK`, `usage.prompt_tokens = 10` | + +- **There is no silent-truncation option.** Nothing in the request says "cut it to my window", + and the service does not do it on its own initiative. Any caller that wants text embedded has + to cut it explicitly, before sending, itself. +- **The rejection does not diagnose itself.** An input 0.55% over the limit comes back as + `400 "The parameter is invalid. Please check again."` — indistinguishable, without measuring, + from a malformed request, a wrong model name or a missing field. That is what makes the incident + in §Background so hard to read: the same request shape succeeds for 8,397 documents and fails + for one, and the error text points at "parameters", not at length. +- **The count arrives too late to be used.** `usage.prompt_tokens` is returned *with a successful + response*; the request whose count you needed is precisely the one that fails and reports + nothing. A caller therefore cannot probe its way to an answer — it must estimate the count + *before* sending. And since one over-length item fails the entire batch, the estimate has to be + conservative; there is no "let this one through and skip it" mode. (Which is why L3 exists, and + why it isolates the failing input instead of retrying the batch as-is.) + +### 2. There is no single way to count tokens across models + +The catalog (`conf/all_models.json`) lists 215 embedding entries (152 distinct base models - the +rest are task and quantisation variants); the ones this deployment can count by name fall into +four tokenizer architectures, and the remainder are either hosted services that publish nothing at +all or architectures that are not implemented here (the split, regenerated from the catalog, is in +§Test strategy across all embedding models): + +| family | examples | scheme | why an estimate from another family does not transfer | +|---|---|---|---| +| BERT WordPiece | bge-*-en, e5-*, gte-base, jina-v2 | 30k word pieces, greedy longest-match, case folding; words longer than 100 chars collapse to a single `[UNK]` | a 671-char single word is **2** tokens here and **146** under bge-m3 (73x) | +| XLM-R SentencePiece | bge-m3, multilingual-e5, m3e, jina-v3 | 250k Unigram, per-language subwords, byte fallback | CJK and punctuation split far more finely than in a byte-level BPE | +| byte-level BPE | Qwen3-Embedding, gte-Qwen, e5-mistral | merges over a 256-byte alphabet, never `[UNK]`, `▁` for spaces | no unknown-token escape hatch: every byte is billable, including invisible ones | +| tiktoken BPE | text-embedding-3-*, ada-002 | ~100k byte-level BPE tuned on English and code | compresses digits and repeated markup far more than the multilingual ones | + +Two consequences: the provider cannot be asked for the answer (no embedding API in this +deployment exposes a pre-flight "how many tokens is this?" call, and the hosted models — +Gemini, Cohere, voyage, Titan, … — publish no tokenizer at all), so the **tokenizer file is the +only authority**; and it has to be *loaded and followed*, not guessed, because the file's own +declarations can be wrong (bge-m3's `.model` says `remove_extra_whitespaces = true` and the +tokenizer actually served does not honour it). + +### 3. The estimates differ substantially, in both directions and per content class + +Same text, each model's own tokenizer (`TOKENIZER_ORACLE_DIR` fixtures; cl100k measured locally): + +| sample | chars | cl100k | bge-m3 | Qwen3-Embedding | e5-mistral | spread | +|---|---|---|---|---|---|---| +| `prose` | 1,800 | 401 | 521 | 401 | 481 | 1.3x | +| `markdown_table` | 1,770 | 1,110 | 1,171 | 1,440 | 1,500 | 1.6x | +| `digits` | 1,100 | 500 | 401 | 1,100 | 1,101 | 2.7x | +| `punctuation` | 1,600 | 481 | 641 | 481 | 681 | 2.5x | +| `fullwidth_forms` | 360 | 521 | 141 | 360 | 881 | **11x** | +| `emoji_zwj` | 360 | 702 | 361 | 480 | 841 | **10.5x** | +| incident span (t[42000:67000]) | 25,000 | 8,143 | 8,237 | 9,782 | 10,259 | 1.4x | +| incident document (`78785.md`) | 78,616 | 28,286 | 27,416 | 34,320 | 35,678 | 1.4x | + +- The disagreement is not a rounding error, and **its sign changes with content**: cl100k is + 1.1% *below* bge-m3 on the incident span (the direction that silently defeats a safety margin) + and **52% above** it on `numbers_and_units`; over the 34 non-trivial samples it is below on + **14**, by up to **33%** (`punctuation`), and the spread between the cheapest and the most + expensive counter is a **median of 2.3x**, up to **11x**. +- It is **content-dependent**, so a global fudge factor cannot fix it either: cl100k and bge-m3 + differ by 1.1% on the incident span but by 3.7x on `fullwidth_forms`. A constant margin wide + enough for the worst class would discard most of the window for every other class. This is why + the margin is proportional and why L2 *calibrates* the ratio per model from measured `usage` + instead of assuming one. + +> Fidelity to a tokenizer is not the same as being a usable estimator for a different +> architecture, and `long_word` shows how far apart they can be: a 671-character single word is +> 2 WordPiece tokens (the >100-char rule collapses it to one `[UNK]`) and 146 bge-m3 tokens. Sizing +> a byte-level model with a WordPiece count would be wrong by two orders of magnitude, in the +> unsafe direction. + +In short: every send needs an explicit trim before it goes out (1); that trim can only be +computed accurately by the model's own tokenizer, which in practice means per-family +implementations (2); and for models whose tokenizer cannot be obtained or trusted, the only +honest fallback is a calibrated upper bound plus a shrink-and-retry path (3). + +## The invariant + +1. **Counting**: text handed to an embedding API is counted either with the + embedding model's own tokenizer or with a **calibrated upper bound** of it. + `cl100k_base` alone is never sufficient: it and a model's tokenizer disagree by + roughly ±2%, and the sign depends on the content. +2. **Failure**: an over-limit rejection shrinks the input and retries. It never + fails a document on its own, because the chunks of that document have already + been deleted by the time the embedding runs. + +Everything below exists to hold those two lines. + +## Pipeline + +``` +kb.embd_id ──► ModelProviderService.GetEmbeddingModel + └─ service/model_service.go model + APIConfig + declared max tokens + │ + ├─ EmbeddingModel.ResolveTokenizerID() catalog "tokenizer" id + └─ EmbeddingModel.ResolveMaxTokens() extra.max_tokens → context_length → 2048 + │ + ▼ + TokenizerComponent.Invoke internal/ingestion/component/tokenizer.go + · per chunk: trimForEmbedding(embedder, text) + └─ embedder implements EmbedderTrimmer ⇒ embedder.Trim(text) + · embeddings: embedder.Encode(ctx, batch) (batch = BatchSize(), default 16) + │ + ▼ + embedder (internal/ingestion/task/embedder.go) + · limiter() = Counter(registry) + ratio(L2 calibration) + margin + · Trim(text) = limiter.Trim(text, ResolveMaxTokens()) + · Encode: call → usage → TokenCount distribution → calibration + └─ on over-limit: overLimitLadder → encodeIsolating → floor truncation +``` + +## The counters (L1) + +`internal/tokenizer` owns a small registry. A `Counter` is: + +```go +type Counter interface { + ID() string // catalog id + Count(text string) int // tokens for this tokenizer + TrimToLimit(text string, limit int) string // longest prefix with Count <= limit + Available() bool // the asset loaded +} +``` + +`RegisterCounterLoader(id, load)` installs a lazy loader (run at most once per +process; a failed load is remembered). `ResolveCounter(id)` falls back to cl100k, +and `CounterExact(id)` reports whether the id resolved to a real counter — that is +what decides whether a calibrated ratio has to be applied. Callers that must not +degrade check `CounterExact` first and refuse: the ingest embedder does, so a model +that declares an unavailable counter fails instead of being calibrated silently. + +| id | family / models | asset (SHA-1 pinned) | implementation | +|---|---|---|---| +| `cl100k_base` | OpenAI `text-embedding-3-*`, `ada-002` | `ragflow_deps/cl100k_base.tiktoken` (pre-existing) | tiktoken-go + offline loader (`bpe_loader.go`) | +| `xlmr-spm` | `bge-m3`, `multilingual-e5-*`, `m3e`, `gte-multilingual`, `jina-embeddings-v3` | `…/BAAI/bge-m3/sentencepiece.bpe.model` `7e88c49f…` | `spm.go`: Unigram Viterbi, byte fallback, NFKC + whitespace + charmap-lite | +| `bert-wordpiece` | `bge-{large,base,small}-en-v1.5`, `bge-en-icl`, `e5-{base,large}-v2`, `gte-{base,large}`, `jina-embeddings-v2-*` | `…/BAAI/bge-large-en-v1.5/vocab.txt` `c3b41053…` | `wordpiece.go`: BertNormalizer + BertPreTokenizer + `##` longest-match | +| `qwen-bpe` | `Qwen3-Embedding-*` | `…/Qwen/Qwen3-Embedding-0.6B/tokenizer.json` `e6592f4d…` | `bpe.go`: hand-written Split scanner, GPT-2 byte alphabet, BPE | +| `llama-bpe` | `e5-mistral` (Mistral/Llama SentencePiece-BPE) | `…/intfloat/e5-mistral-7b-instruct/tokenizer.json` `92cb22e1…` | `bpe.go`: prepend/replace `▁`, whole-unit merge, byte fallback | + +The XLM-R family shares one vocabulary file, so one 5 MB asset covers five model +families; the BERT family shares `bert-base-uncased`'s vocab. + +## Choosing a limit (margin + L2) + +`EmbeddingTokenLimit(maxTokens)` = `maxTokens − max(32, 2%)`. The margin is a +ratio on purpose: the previous constant was 10 tokens, 0.12% of an 8192 window, +an order of magnitude smaller than the disagreement it was meant to absorb. + +`ResolveEmbeddingMaxTokens(declared, contextLength)` prefers the model's own +`max_tokens`, then the catalog's `context_length`, then `EmbeddingTokenLimitDefault` +(2048 — deliberately the small end, because overshooting a window is a rejected +request while undershooting only truncates). The catalog has 512-token embedding +models; a hard-coded 8192 would have overshot those by 16x. + +`Limiter` binds a counter, a calibration and a margin: + +```go +Limiter.Limit(maxTokens) = EmbeddingTokenLimit(floor(ResolveEmbeddingMaxTokens(...) / Ratio())) +Limiter.Trim(text, maxTokens) (trimmed string, tokens int) +``` + +- `NewExactLimiter(counter)` — the counter *is* the model's tokenizer, ratio 1. +- `NewCalibratedLimiter(counter, key, cal)` — ratio read from the calibration **on + every call**, so an over-limit rejection recorded by a retry tightens the budget + of the very next attempt, not of the next document. + +### Calibration + +`Calibration` learns `real / own` per `(provider instance, model)` — the key is +`embedder.quotaKey()` (endpoint + region + model + hashed API key), the same key the +rate-limit cooldown uses. It only ratchets **up**: + +| Observation | Inference | +|---|---| +| `ObserveUsage(key, own, real)` on a successful call | `ratio = max(ratio, real/own)` | +| `ObserveOverLimit(key, own, maxTokens)` on a rejection | `ratio = max(ratio, maxTokens/own × 1.01)` | + +`ObserveOverLimit` must be given **the largest single input's count, not the batch +total**: the window bounds each input, and a batch total is normally above the +window, which would imply a ratio below 1 and teach nothing. + +For a model with no exact counter the starting bound is +`DefaultUncountedRatioUpper = 1.05`, which costs ~5% of truncation headroom and is +what makes unknown and future models safe without any asset. + +## Recovering from an over-limit rejection (L3) + +`isOverLimitErr` classifies a 4xx whose message matches `20015` (SiliconFlow's code +for exactly this — it answers a generic "The parameter is invalid") or the usual +"too long / maximum context / context length / token limit / …" phrasings. A 5xx is +never treated as a size problem. + +`overLimitLadder(budget)` = `budget, 0.75b, 0.5b, 0.25b, 0.125b`, floored at +`overLimitFloorTokens = 64`. The first steps are gentle because an over-limit +rejection usually means the input is a little past the window; cutting a quarter of +the content for that would lose text for nothing. Trimming is a **prefix +operation**, so re-trimming an already trimmed text simply shortens it — the ladder +needs no access to the untrimmed input. + +If the whole batch still fails at the smallest budget, `encodeIsolating` embeds the +inputs one at a time (a batch is allowed to contain exactly one pathological input) +and returns an error only if an input fails at the floor, i.e. for reasons unrelated +to size. + +## Wiring, and the seams that keep it back-compatible + +- `component/tokenizer.go` prefers two **optional** interfaces, so every existing + test stub still satisfies `Embedder`: + `EmbedderTrimmer{ Trim(text) (string, int) }` and + `EmbedderMaxResolver{ ResolveMaxTokens() int }`. Without them the component falls + back to `truncateForEmbedding`, which now also uses the resolved window and the + proportional margin. +- `embedder.Encode` passes a `*common.ModelUsage` into the driver and reads + `InputTokens` back. `distributeTokenCount` spreads that request-level total over + the inputs by our own counts (remainder to the largest, so the sum equals the + provider's number); with no usage it falls back to our own count rather than to 0, + which is what it used to report. +- `EmbeddingData.TokenCount` documents that per-input semantics; the stub FIXME is + gone. `openai.go` now records usage like `siliconflow.go` already did. + +## Data: the model catalog + +`conf/all_models.json` gained a `tokenizer` field per model (next to the existing +`max_batch_size` / `context_length`). **36** entries are tagged (7 `xlmr-spm`, 14 +`bert-wordpiece`, 9 `qwen-bpe`, 6 `cl100k_base`). + +Only families whose counter was verified against the model's own tokenizer are +tagged. Untagged models are not a gap: they resolve to the calibrated cl100k +fallback, which errs toward more headroom. A wrong tag is the dangerous direction, +so guessing is deliberately avoided - §Test strategy across all embedding models +states the invariant this rests on and shows how much of the catalog each tag covers. + +Two escape hatches for operators (both win over the catalog): +`TOKENIZER_EMBEDDING_TOKENIZER` (pin an id) and `TOKENIZER_EMBEDDING_MAX_TOKENS` +(pin a window). A third one is about where the assets live rather than how they are used: +`MODEL_ASSETS_DIR` points at a directory holding the downloaded model assets in the +HuggingFace repo layout (see §Assets and the oracle). + +## Test strategy across all embedding models + +Verification here is per **architecture**, not per catalog entry - and the catalog is +mostly entries, not architectures. `scripts/tokenizer_coverage.py` reads +`conf/all_models.json` and prints the current picture (no network, no per-model +guessing): + +``` +embedding entries: 215 | distinct base models: 152 | tagged entries: 36 | calibrated entries: 179 +``` + +| tagged counter | entries | distinct base models | scheme | +|---|---|---|---| +| `bert-wordpiece` | 14 | 14 | BERT WordPiece | +| `qwen-bpe` | 9 | 6 | byte-level BPE (Qwen pre-tokenizer regex) | +| `xlmr-spm` | 7 | 7 | XLM-R Unigram SentencePiece | +| `cl100k_base` | 6 | 6 | tiktoken BPE (OpenAI table) | + +| untagged group (name heuristic) | entries | base models | what it would take | +|---|---|---|---| +| hosted services (cohere, voyage, gemini, titan, nvidia NIM, …) | 77 | 77 | nothing to verify against: no published tokenizer. L2/L3 only | +| other architectures (`jinaai/jina-embeddings-v5*`, `google/embeddinggemma-*`, …) | 95 | 35 | a new implementation + its fixture + an oracle | +| `bert-wordpiece` | 4 | 4 | byte-identity evidence against the family fixture; a different vocabulary (`BAAI/bge-large-zh-v1.5`, `netease-youdao/bce-embedding-base_v1`, …) needs its own artifact + fixture + oracle - see The tag invariant | +| `llama-bpe` | 3 | 3 | byte-identity evidence against the family fixture; a different vocabulary (`mistral/mistral-embed`, two `nvidia/llama-nemoretriever-*`) needs its own artifact + fixture + oracle - see The tag invariant | + +The groups are a *prioritisation* heuristic, not a claim: a model is tagged only after +its own artifact has been through the oracle (recipe below). + +### What one fixture per architecture does and does not prove + +| architecture | fixture artifact | what varies between models of this family | read from the artifact? | exercised by the corpus? | +|---|---|---|---|---| +| `bert-wordpiece` | `BAAI/bge-large-en-v1.5/vocab.txt` | `do_lower_case`, `strip_accents`, `handle_chinese_chars`, `max_input_chars_per_word`, the unk token | **no** - `wordpiece.go` hard-codes the family's shape (`lowercase=true`, `strip_accents` follows it, `max_input_chars_per_word=100`); a cased or differently-configured BERT needs its own family | **partly** - the fixture is `lowercase=true`; no cased model is in the corpus | +| `xlmr-spm` | `BAAI/bge-m3/sentencepiece.bpe.model` | the precompiled charmap, `add_dummy_prefix`, `remove_extra_whitespaces`, `escape_whitespaces`, byte fallback, whether a `Metaspace` pre-tokenizer splits at all | yes | whitespace and charmap classes yes; **byte fallback no** (bge-m3 has none) | +| `qwen-bpe` | `Qwen/Qwen3-Embedding-0.6B/tokenizer.json` | the pre-tokenizer regex, `byte_fallback`, unk presence, `NFC` | yes - the regex is copied verbatim from the artifact | one artifact only | +| `llama-bpe` | `intfloat/e5-mistral-7b-instruct/tokenizer.json` | `byte_fallback`, the `▁` prepend/replace normalizer | yes | one artifact only | +| `cl100k_base` | the shipped `cl100k_base.tiktoken` | nothing - one table | n/a | complete | + +So the honest statement of coverage is: **the algorithm is verified per family, the +vocabulary per fixture artifact.** A model that differs from the fixture only in +vocabulary size is *not* covered by the fixture, and that is the axis the table is for. + +### The tag invariant + +> A tag means "this model's tokenizer **is** the artifact the fixture verified", not +> "this model looks like that family". + +The counters load **one fixed artifact per family** (`bert-wordpiece` always reads +`bge-large-en-v1.5/vocab.txt`, `xlmr-spm` always reads `bge-m3/sentencepiece.bpe.model`). +Tagging a model whose vocabulary differs - a Chinese `bge-large-zh`, a Jina-v2 BERT with +its own vocab, a cased model - would count it with the wrong vocabulary and *look exact +while being wrong*, which is worse than the calibrated fallback. Two consequences worth +knowing: the loaders also search a bare `vocab.txt` in the working directory, so dropping +an unrelated `vocab.txt` there silently replaces the verified artifact; and auditing the +36 existing tags against their own artifacts is an open item (the runtime calibration in +§Calibration is what would surface a mistake, as a drift in `real/cl100k`). + +### Adding a model (recipe) + +1. **Identify the architecture** from the model's own artifact (`tokenizer.json`, or + `vocab.txt` / `sentencepiece.bpe.model`). If the vocabulary is byte-identical to a + family fixture, only step 4 is needed; if it differs, the model needs its own fixture. +2. **Add the artifact** to `ragflow_deps/download_go_deps.py`'s `TOKENIZER_ASSETS` list - + individual files, not `snapshot_download` of a multi-GB repo - pin its SHA-1 in the + loader map so a mismatch fails at load instead of counting with the wrong table, and add + it to the copy loop in `Dockerfile` / `Dockerfile_base` / `Dockerfile_go` so every runtime image ships it + (the `ragflow_deps` image is built from this script's output, so it needs a refresh too - + see §Assets and the oracle for the whole chain). +3. **Add the fixture line** to `ORACLES` in `scripts/gen_tokenizer_oracle.py`, generate it, + and run `TestCountersMatchOracle`. Passing means every sample agrees, or sits inside the + flagged bounded class (the sample count follows `--fuzz`, see §Assets and the oracle); + anything else is a bug in the counter, not in the model. +4. **Then tag it** in `conf/all_models.json`. +5. **Record it** - regenerate the coverage table with + `scripts/tokenizer_coverage.py --markdown`. If the new model exercises an axis that was + not exercised before, add the sample class to the corpus in the same change. + +Acceptance criteria for steps 3-5 are the ones in §Verification: oracle at token level, +property tier green, `gofmt` clean. + +### What this strategy deliberately does not attempt + +- **A fixture per catalog entry** (152 base models). The code encodes the *algorithm* and + the artifact encodes the *vocabulary*; per-entry fixtures would mean ~150 downloads and + no extra code coverage - the axes table above is what finds bugs, not the model count. +- **Claiming exactness for untagged models.** They are *calibrated* at runtime and + observed through the provider's `usage`; the doc says "calibrated", never "verified". +- **Following the catalog's `context_length` blindly** - it is honoured as declared, and + §Not covered keeps the caveat that nothing here validates what a provider accepts. + +### Next moves, cheapest first + +1. Tag the **7 base models** whose architecture is already implemented - it needs their + artifacts, one fixture each, no new code. +2. Close the **`lowercase=false` WordPiece axis** (one cased model fixture) - today that + flag is read but never exercised. +3. Exercise **`byte_fallback=true` on the SPM family** (a `.model` that ships the byte + pieces), which the current fixture cannot reach. +4. Only then a **new architecture**, chosen by entry count: the `jina-embeddings-v5*` / + `embeddinggemma` group is ~35 base models. + +## Assets and the oracle + +Assets are fetched, never committed: `ragflow_deps/download_go_deps.py` (the Go-side +downloader - `download_deps.py` is upstream and stays untouched) has a `TOKENIZER_ASSETS` +list that downloads **individual files** (not `snapshot_download` of a repo that also +carries multi-GB weights) into `ragflow_deps/huggingface.co//`, which +`.gitignore` already excludes (`huggingface.co/`). Loaders read them from disk with a +SHA-1 pin, walk up from the working and executable directories like `bpe_loader.go`, and +perform no network I/O. A missing asset fails loudly at every layer that can report one: the +runtime image does not build without it (the copy loop below exits non-zero), an embedder whose +model declares a tokenizer it cannot load refuses to count rather than substituting the +calibrated estimate, and the cl100k table - the one asset that still comes from +`download_deps.py`, which fetches it from the OpenAI blob - has always been a startup panic. + +**The whole chain, because a gap in it is invisible**: `download_go_deps.py` writes +`ragflow_deps/huggingface.co/…`; the `ragflow_deps` image is built from that directory +(`ragflow_deps/Dockerfile` copies `huggingface.co` to `/huggingface.co`); the root +`Dockerfile`, `Dockerfile_base` and `Dockerfile_go` bind-mount that image and copy the four *runtime* +tokenizer assets into `/ragflow/ragflow_deps/huggingface.co/…`, which is where the +counters look (the working directory is `/ragflow`). The `tokenizer.json` files that +exist only as cross-check oracles are deliberately **not** shipped. Two consequences +worth knowing: adding an asset means touching all three places (script, deps image, +runtime image) - and the copy loop **fails the build** when the deps image does not carry +it, so no image can be produced with a counter that would silently fall back - and an +embedder whose model declares a tokenizer it cannot load **fails the ingest** with +`embedding tokenizer "" is declared for but its asset is unavailable …: +refusing to count with the calibrated estimate`, instead of counting differently. + +The three places are held together by `ragflow_deps/test_tokenizer_assets.py`, which +fails when a `"runtime"` asset is missing from either Dockerfile, when an `"oracle"` +asset is shipped, or when a runtime asset has no 40-hex SHA-1 pin in +`internal/tokenizer/*.go`. It is there because the first version of this chain silently +skipped step three: the counters loaded fine in a developer checkout and quietly +degraded to the calibrated path in the image. + +### Where the assets are looked up, and how to move them + +By default the loaders walk the working directory and the executable's directory (and +their ancestors) for the shipped relative names. `MODEL_ASSETS_DIR` overrides that: it +names a directory playing the role of `ragflow_deps/` - downloaded model assets in the +HuggingFace repo layout - and it is **consulted first**. Three spellings are accepted, so +an operator can point at whichever tree they have: + +| `MODEL_ASSETS_DIR` points at | candidate used | +|---|---| +| the `ragflow_deps`-like root | `/huggingface.co//` | +| a `huggingface.co`-like root | `//` | +| a flat directory | `/` | + +The variable is deliberately **not** embedding-specific: DeepDoc's weights use the same +layout (`huggingface.co/InfiniFlow/deepdoc`) and `resolveDeepDocModelDir` honours it too, +so one mount can serve every downloaded model asset. `ragflow_deps/download_go_deps.py` and +the three runtime Dockerfiles keep writing the same tree, so nothing changes for a default +deployment. The layout and the precedence are pinned by tests: +`internal/common/model_assets_test.go` for the candidate list, and +`internal/tokenizer/asset_dir_test.go` for the end-to-end case (a child process with an +unrelated working directory loads the SPM counter from the configured tree). + +Two things make the state of the assets observable: + +- each counter reports the file it loaded (`SourcePath`), so logs and failures can name it; +- the ingestor logs an **availability report once at startup** - + `embedding tokenizer counters {available: [...]}` plus, when something is missing, a + warning naming the unavailable ones and pointing at `ragflow_deps/download_go_deps.py` / + `MODEL_ASSETS_DIR`. The report is diagnostics rather than the guard: startup is not fatal + (untagged models count with the calibrated estimate by design, and they keep working), but + ingesting a model that declares a missing asset fails outright - so this report is what + tells an operator which asset to restore. + +The oracle is the point of the whole test story: a hand-written tokenizer that is +"close" is not good enough, because an **under-count** lets an oversized input +through and the provider answers 400. + +``` +mkdir -p /tmp/tokenizer_oracle +for m in bge-m3 bge-large-en qwen3-embedding e5-mistral cl100k; do + uv run scripts/gen_tokenizer_oracle.py --model $m --corpus /path/to/78785.md \ + --out /tmp/tokenizer_oracle/$m.json --fuzz 300 +done +TOKENIZER_ORACLE_DIR=/tmp/tokenizer_oracle ./build.sh --test -run TestCountersMatchOracle ./internal/tokenizer/ +``` + +`--corpus` is required: the document's **whole** content becomes one sample (`incident_full`), +so the corpus is whatever the caller wants checked - the incident document `78785.md` (see +§Background) is the one that exposed the bug. + +Current state: **339 samples per counter (`--fuzz 300`: 32 hand-written samples + the whole +`--corpus` document + 300 fuzzed + 6 fixed extremes), all five counters passing** - four of them +compared token by token, `cl100k_base` by token id. The fixture records the reference's +**token ids and token texts**; comparing texts is what makes the comparison meaningful, +because two different segmentations can share a token count (the incident document had +exactly one such spot, 18,182 tokens in). + +### What each counter is checked against + +| counter | reference implementation | what is compared | +|---|---|---| +| `cl100k_base` | OpenAI's own `tiktoken` (Python 0.12.0) over the table shipped in `ragflow_deps` | token **ids** (the numbering is shared - both read the same table) | +| `xlmr-spm` | the HuggingFace conversion of `BAAI/bge-m3` (the files SILICONFLOW serves), via `tokenizers` 0.22.2 reading the model's `tokenizer.json` | token **texts** - this family's ids are renumbered by the conversion (below) | +| `bert-wordpiece` | the same library reading `BAAI/bge-large-en-v1.5/tokenizer.json`; the algorithm itself is Google's `WordPiece` / `BasicTokenizer` | token texts (ids match as well) | +| `qwen-bpe`, `llama-bpe` | the same library reading each model's `tokenizer.json` | token texts (ids match as well) | + +Two things this table deliberately does **not** claim: + +- **Google's `sentencepiece` is a second opinion, not the authority.** It is installed + (`.venv`, 0.2.2); run over the same corpus it leaves 253/341 samples identical and comes + out *lower* than the served conversion on 82 of them, every one of them whitespace-shaped + (the `.model`'s `remove_extra_whitespaces`, which the canonical implementation honours and + the served conversion does not). Following it there would **under-count**, i.e. fail in the + dangerous direction - which is the measured reason the served conversion is the reference. + Method, the remaining classes and the limits of the check are in §The three-way cross-check + below. +- **Token ids are not portable across artifacts.** The bge-m3 conversion reorders the + SentencePiece specials and inserts two tokens, so the same piece is `id` in the + `.model` and `id+1` in the served vocabulary (`,,,…` becomes + `,,,,…`). Comparing ids across that boundary would report a fake + failure on every sample; comparing the *texts* isolates real segmentation + differences. (HuggingFace also renders an unknown token as the source text it could + not encode, so the fixture normalizes those to ``.) + +### The three-way cross-check (what it is, and how to read it) + +For `xlmr-spm` three parties are compared: **two readers of the same `.model` file**, plus +the conversion that is actually served. + +| party | what it reads | how its ids are numbered | +|---|---|---| +| **ours** (`spm.go`) | `sentencepiece.bpe.model` | `.model` piece index | +| **the served conversion** (HF `tokenizers`) | `tokenizer.json` built from that `.model` | *renumbered*: same pieces, specials reordered, ``/`` inserted ⇒ regular pieces are `+1` | +| **Google `sentencepiece`** (canonical C++) | the same `sentencepiece.bpe.model` | `.model` piece index - i.e. **the same numbering we use** (verified: `▁3`=137, `43`=11547, ``=0, 250,000 pieces) | + +How it runs, and how to read it: + +```bash +uv run scripts/gen_tokenizer_oracle.py --model bge-m3 --corpus /path/to/78785.md \ + --out /tmp/tokenizer_oracle/bge-m3.json --fuzz 300 --spm-crosscheck +``` + +The flag adds `spm_count` per sample and a `crosscheck` note to the fixture, and prints the +number of samples where the two readers disagree. **The Go test never reads that field**, so +the cross-check cannot change a verdict - it is evidence, not an assertion. + +Measured result on the 341-sample corpus the cross-check was run on (the corpus document +is sampled whole now, so a re-run is two samples short of that): + +| class | samples | cause | what it means | +|---|---|---|---| +| both readers agree, count **and** piece sequence | **245/341** | - | the canonical implementation independently reproduces the segmentation; since our counter is token-identical to the served reader on this corpus, our ids equal Google's there too | +| count equal, piece sequence differs | 8/341 | combining marks (NFKC composes, the charmap does not) | the same hazard class as the `known_approx` flag | +| Google **lower** by 1 | 80/341 | `.model` declares `remove_extra_whitespaces = true`; the canonical reader honours it (trims leading/trailing whitespace), the served conversion does not | following Google here would **under-count** - the dangerous direction | +| Google lower by 2 | 2/341 | same | same | +| Google **higher** by 1 | 6/341 | combining-mark / charmap shapes | same class as the `known_approx` flag | + +**Decision rule.** The served conversion decides. Google's implementation is an *independent +check of our modelling*: if our output matched neither reader, that would be a bug in `spm.go`. +A disagreement *between the two readers* is never by itself a reason to move `spm.go` towards +`.model` semantics - the table above is where such a change has to be argued, with numbers. + +**Limits of the check**, so it is not over-read: + +- It compares **counts** (plus the piece strings of the two readings). Comparing ids directly + still needs two differences modelled first: the served conversion's renumbering, and the fact + that the served pipeline runs one Viterbi *per pre-token* while SentencePiece runs one over + the whole text. +- It covers `xlmr-spm` only - it is the one family with a `.model` file, and therefore the only + one with a canonical implementation to check against. +- It is a second opinion, not an extra gate: only `TestCountersMatchOracle` and the property + tier fail a build. + +### The fuzz corpus + +`--fuzz N` appends `N` deterministic random strings - seeded, with the texts stored in +the fixture so the Go side never has to reproduce the RNG - drawn from 17 alphabets: +ASCII words, CJK, Greek, Cyrillic, Hebrew, Arabic, emoji (ZWJ sequences and skin-tone +modifiers included), combining marks, zero-width characters, control characters, +full-width and compatibility forms. It also appends boundary cases (`a`×99, `a`×100, +`a`×101 for WordPiece's 100-character rule) and a printable-ASCII sweep. + +The two halves catch different things: the hand-written samples found the first five bugs +in the table below, and the fuzz corpus found every one of the five at the end of it - on +its first run, after the hand-written corpus had been passing. + +### What the verification stack has caught + +Every one of these looked correct in review and passed a prose-only corpus. The +"caught by" column is the reason the tiers in the next section are not optional: +the oracle catches *modelling* mistakes, the property test catches *contract* +mistakes, and neither substitutes for the other. + +| Bug | Symptom | Caught by | +|---|---|---| +| SPM followed `.model`'s `remove_extra_whitespaces` instead of the served HF conversion | one token short on 15/23 samples | oracle, `leading_trailing_space` | +| BERT `isPunctuation` included Unicode **Symbols** | every emoji cost an extra token | oracle, `emoji` | +| BERT `clean_text` turned control/format characters into spaces instead of deleting them | "a\u200bb" split into two words | oracle, `control_chars` / `zero_width_and_nbsp` | +| BERT `strip_accents` short-circuited when the input was already NFD | decomposed accents counted twice (the marks were the thing to remove) | oracle, `decomposed_accents` | +| SPM ignored the precompiled charmap's delete/space rules | control and zero-width characters were billed as bytes | oracle, `control_chars` | +| BPE merged **one occurrence** of the lowest-rank pair per evaluation, and the heap broke rank ties arbitrarily | 32 `a`s counted as 9 instead of 7 | oracle, `long_word` | +| `cl100k`'s trim cut the token list without snapping to a rune boundary | a trimmed emoji chunk was invalid UTF-8 (and would have gone into the index that way) | **property test**, `emoji` at `limit=1` | +| BERT WordPiece kept a word's already-matched prefix and then one `[UNK]` for the rest | Google's rule is that if **any** position of a word fails to match, the whole word becomes a single `[UNK]`; the old behaviour over-counted such words | **token-level oracle**, fuzz corpus | +| SPM ran one Viterbi over the whole text | the served pipeline runs one **per Metaspace pre-token**; on the incident document that is count-neutral (27,416 tokens either way) but a different segmentation at byte 46,673 | **token-level oracle**, `incident_full` | +| SPM billed a bare `▁` for text the charmap deletes entirely | `"\x7f\x01\x7f"` counted 1 token where the model sees 0 | **token-level oracle**, fuzz corpus | +| SPM emitted one token per unmatched character | the served model fuses a whole run into **one** unknown token (`🧑🧑`, five NULs, each one token) | **token-level oracle**, fuzz corpus | +| SPM accumulated Viterbi scores in float32 | ties between equal-scoring segmentations broke the other way (a run of 99 `a`s) | **token-level oracle**, fuzz corpus | + +The BPE one is worth repeating for anyone touching `bpe.go`: the canonical rule is +"take the lowest-rank adjacent pair and merge **every** non-overlapping occurrence +of it", and when several occurrences share a rank the leftmost one goes first. Both +halves matter, and both are pinned by tests. + +## How the implementation is kept honest + +Two questions decide whether a counter may be trusted: **what is the normative +reference**, and **what has been proven against it**. Both are answered per counter, +because "it looked right in review" is how all twelve bugs above got in. + +### 1. The normative reference is the model's own artifact, nothing else + +Each counter is derived from a file that ships in the repository and is pinned by +SHA-1. The loader refuses a digest mismatch, so a test can never silently run +against a different vocabulary: + +| counter | authoritative artifact (pinned) | format / algorithm reference | behavioural reference | +|---|---|---|---| +| `cl100k_base` | `ragflow_deps/cl100k_base.tiktoken` (`expectedBpeHashes`, `bpe_loader.go`) | tiktoken's table format (` ` per line) | the table itself, plus the Python anchors in `TestNumTokensFromString_MatchesPythonAnchors` | +| `xlmr-spm` | `BAAI/bge-m3/sentencepiece.bpe.model` `7e88c49f…` | `sentencepiece/model.proto` wire format (decoded by the minimal protobuf reader in `spm.go`): `pieces`, `trainer_spec.model_type = UNIGRAM`, `normalizer_spec`; Unigram Viterbi; `BYTE` pieces for fallback | the HuggingFace conversion of the same repo (`tokenizer.json`) — it decides whitespace and pre-tokenization behaviour, and it **wins over the `.model` flags** when they disagree | +| `bert-wordpiece` | `BAAI/bge-large-en-v1.5/vocab.txt` `c3b41053…` | HuggingFace `BertNormalizer` / `BertPreTokenizer` / `WordPiece`, whose configuration is quoted from the model's `tokenizer.json` in `wordpiece.go`'s header; BERT's own `_is_punctuation` (category **P**, not S) | the same `tokenizer.json` | +| `qwen-bpe` | `Qwen/Qwen3-Embedding-0.6B/tokenizer.json` `e6592f4d…` | the artifact's `pre_tokenizer` regex, copied verbatim into `bpe.go`, plus GPT-2's `bytes_to_unicode` alphabet and canonical BPE merge order | the artifact's own `vocab` + `merges` (`NFC`, `byte_fallback=false`) | +| `llama-bpe` | `intfloat/e5-mistral-7b-instruct/tokenizer.json` `92cb22e1…` | the artifact's `normalizer` (`Prepend("▁")` + `Replace(" "→"▁")`) and `byte_fallback=true` | the same artifact | + +Two consequences of taking the artifact as the spec, both learned the hard way: + +- **A declaration inside the artifact is not evidence.** bge-m3's `.model` declares + `remove_extra_whitespaces = true`; the tokenizer providers actually serve does not + honour it. When the file and the served conversion disagree, the served conversion + is what decides whether a request fits the window, so it is the reference we + follow — and the oracle is what tells us they disagree. +- **What cannot be reproduced exactly is enumerated, not glossed over.** The + `nmt_nfkc` charmap is a 232 KB Darts trie; `spm.go` implements the subset that + NFKC does not already cover, and each rule in it was derived by **probing the + model one character at a time** (delete `U+0001-0008`, `U+000B`, `U+000C`, + `U+000E-001F`, `U+007F-0084`, `U+0086-009F`; `U+200B-200D`/`U+FEFF` to space; + keep `U+0000`, `U+0085`, `U+2060`) rather than copied from a description. The + same applies to the models with no public tokenizer at all (see §4). + +### 2. What is proven, tier by tier + +| Tier | Test | What it would catch | +|---|---|---| +| Artifact integrity | SHA-1 pins (`expectedSPMHashes`, `expectedWordPieceHashes`, `expectedBPEHashes`), plus `model_assets_test.go` / `asset_dir_test.go` for `MODEL_ASSETS_DIR` | a corrupt download, a swapped vocabulary, or an asset tree the process cannot see — the load fails instead of counting with the wrong table, and the startup report says which counters are usable | +| Model fidelity | `TestCountersMatchOracle` (fixtures from `scripts/gen_tokenizer_oracle.py`, `--fuzz 300`) | "our tokenizer is not the model's tokenizer" - per token, not per sample count, so a segmentation that happens to have the right length still fails | +| Counter contract | `TestCountersSatisfyTrimProperties`, `TestCountTrimFitsEveryLimit` | a counter that does not return a prefix, splits a rune, exceeds the limit it was given, is not idempotent, or shrinks when given a larger limit | +| Layer maths | `embedlimit_test.go` (margin, window resolution, calibration ratchet + over-limit inference, limiter, byte-level fallback) | an off-by-one in the margin, a ratio that should have ratcheted up, a fallback that reports invented counts | +| Ingest behaviour | `embedder_limit_test.go` (shrink ladder, per-input isolation, `20015` classification, token distribution, window honoured) | a document that fails instead of shrinking; token counts that do not sum to the provider's | +| Runtime oracle | `Calibration.ObserveUsage` / `ObserveOverLimit` | drift in a counter nobody can test offline — verification keeps running in production | +| End-to-end gate | re-ingest the incident document, compare MySQL vs ES, `covsweep` = 0 unserved words | everything above being green while content is still lost | + +The oracle assertion is **equality**, deliberately not a tolerance: "within 2%" +would have hidden every row of the bug table, because the whole failure mode is a +2% disagreement in the unsafe direction (too few tokens counted ⇒ the provider +rejects what we sent). + +### 3. One bug, one sample class + +A bug is not fixed until the corpus can fail on it. The expectation is regenerated +from the model (`gen_tokenizer_oracle.py`), never typed by hand, so "expected" always +means "what the model does": + +| Sample class | Exists because of | +|---|---| +| `emoji` (incl. ZWJ) | BERT treating Symbols as punctuation; the `cl100k` UTF-8 cut | +| `control_chars` | BERT `clean_text` and SPM charmap delete rules | +| `zero_width_and_nbsp` | BERT Cc/Cf deletion | +| `decomposed_accents`, `vietnamese` | BERT `strip_accents` short-circuit | +| `long_word` (a long run of one character) | BPE merge order and rank tie-breaking | +| `leading_trailing_space`, `tabs_and_crlf` | SPM whitespace semantics | +| `fullwidth_forms`, `compatibility_chars`, `cjk_compat`, `math_and_currency` | the charmap approximation (NFKC vs `nmt_nfkc`) | +| `prose`, `markdown_table`, `issue_numbers`, `digits`, `base64_like`, `incident_*` | the real corpus, including the span that caused this work | +| the fuzz corpus (`--fuzz N`, 17 alphabets) | everything the hand-written samples missed: the pre-token boundary, the run-of-unknowns rule, whitespace-only-after-deletion, the float32 tie, and WordPiece's whole-word `[UNK]` | + +### 4. Rules that keep it from rotting + +1. **No counter without a fixture.** A family whose tokenizer cannot be checked + against the model is not tagged in the catalog; it falls back to the calibrated + path. Guessing a tag is the unsafe direction. +2. **Nothing is hand-typed.** Fixtures and expected counts come from the model's + tokenizer; assets come from the downloaders (`download_go_deps.py` for the tokenizer + files, `download_deps.py` for the cl100k table) with a SHA-1 pin. +3. **Equality, not resemblance** (see above), and the dangerous direction + (under-count) is the one the tests are shaped around. +4. **Every approximation is named** in §Known approximations together with the + backstop that covers it — an unnamed approximation is a bug report waiting to + happen. +5. **Verification continues at runtime.** 77 of the 215 catalog entries (name heuristic; + §Test strategy across all embedding models) are hosted services with no downloadable + tokenizer, so the provider's own `usage.prompt_tokens` is the only reference that exists + for them. That is why `TokenCount` was implemented and why calibration is not optional. +6. **The gate is end-to-end.** Unit and oracle tests can all be green while a + document still loses text; the acceptance check is the incident document plus + `covsweep` reporting zero unserved words. + +### 5. What is not provable, and what covers it instead + +| Gap | Coverage | +|---|---| +| Hosted models with no tokenizer (Gemini, Cohere, voyage, Titan, …) | L2 calibration from `usage` + L3 shrink-and-retry | +| `nmt_nfkc` charmap outside the probed subset | L2/L3; the gap is enumerated in §Known approximations | +| A provider that reports no `usage` | falls back to our own count; L3 still guards the limit | +| Content classes nobody has sampled yet | the property tier holds for arbitrary input, and a mis-count shows up as a rejected request that L3 absorbs rather than as a failed document | +| Google's `sentencepiece` (canonical implementation of the `.model`) is cross-checked, but it disagrees with the served conversion on 88/341 samples | enumerated and explained in §The three-way cross-check (whitespace trimming, 80 × `-1`); the served conversion remains the reference because the other reading under-counts in the dangerous direction | +| Two shapes of combining marks, where NFKC and the served charmap provably differ | flagged per sample and held to ≤ 2 tokens (see §Known approximations); every other sample is compared exactly | +| A finite corpus, however adversarial: 339 samples per counter is evidence, not proof | the corpus is seeded and regenerated from the model on demand (`--fuzz N`), so coverage is a dial rather than a constant, and the runtime oracle covers the rest | + +The BPE one is worth repeating for anyone touching `bpe.go`: the canonical rule is +"take the lowest-rank adjacent pair and merge **every** non-overlapping occurrence +of it", and when several occurrences share a rank the leftmost one goes first. Both +halves matter, and both are pinned by tests. + +## Performance + +- `spm.go` / `wordpiece.go` count in O(n) over the text (trie walk + Viterbi); + `trimToLimit` binary-searches rune prefixes, so O(n log n). +- `bpe.go` keeps a min-heap over adjacent pairs keyed by `(rank, position)` and + drains it per rank: O(n log n), plus a per-piece memo (pieces repeat far more than + they differ). The first two implementations were O(n²) or wrong; the current one + counts the 78 KB incident document in **0.09 s**, against **181 s** for the + rescanning version. +- The whole oracle suite - five counters, 341 samples each, including the 78 KB + incident document - runs in about 1.2 s. + +## Known approximations (deliberate) + +- `xlmr-spm` implements the **subset** of the model's 232 KB precompiled `nmt_nfkc` + charmap that NFKC does not already cover, restricted to the rules verified + character by character against the model (delete `U+0001-0008`, `U+000B`, + `U+000C`, `U+000E-001F`, `U+007F-0084`, `U+0086-009F`; `U+200B-200D`, `U+FEFF` and + the escaped space `U+2581` to space; keep `U+0000`, `U+0085`, `U+2060`). +- **Combining marks are where that subset is measurably short**, and the oracle test + says exactly which samples are affected instead of hiding them. Our normalization + starts from NFKC, which composes and reorders; the served charmap does neither, and + it drops marks that cannot attach to a letter (measured: `②`+`U+0327` normalizes to + `2`, a mark orphaned by a deleted control character disappears, `a`+`U+0301`+`U+0327` + keeps three characters). Two shapes are therefore flagged per sample by + `gen_tokenizer_oracle.py` (`has_combining_run`, `has_dangling_mark`): two or more + marks on one base, and a mark that does not follow a letter. Those samples are held + to a **bounded** difference (≤ 2 tokens) rather than to equality; every other sample + is still compared exactly, which is 156/339 for `xlmr-spm` and 339/339 for + `cl100k_base` in the current fixture. Without `--fuzz` no sample is flagged at all - + the hazard needs combining marks, which real prose rarely has. +- Unusual Unicode outside that set can still be off by a token or two; L2 and L3 are + the backstop. +- Untagged model families (e.g. `mistral-embed`, `gte-Qwen`) use the calibrated + fallback on purpose. +- Truncation is silent by design here: the counter's job is to fit the window. Not + sending an oversized chunk at all is the chunker's job, and oversized chunks still + exist in the index (up to 104k characters) — see "Not covered". + +## Not covered + +- **Chunker hard cap.** Tables and other delimiter-free blocks are still emitted as + chunks far above `chunk_token_size`. Today that costs an extra shrink-and-retry; + it should be fixed in the chunker so the content is complete rather than truncated. +- **Other providers' usage.** Only the OpenAI-compatible and SiliconFlow embedding + paths report `usage` today; other drivers leave `TokenCount` at 0 and the + calibration simply gets no observation from them. +- Windows of models with declared limits of 32768/131072 are honoured as declared; + nothing here validates that a provider accepts that much. + +## Verification + +```bash +# unit tests for the layer (margin, calibration, limiter, L3 ladder, distribution) +# plus the counter contract properties for every counter whose asset is present +./build.sh --test ./internal/tokenizer/ ./internal/ingestion/component/ \ + ./internal/ingestion/task/ ./internal/entity/models/ + +# counter fidelity against the models themselves (339 samples x 5 counters, token level) +mkdir -p /tmp/tokenizer_oracle +for m in bge-m3 bge-large-en qwen3-embedding e5-mistral cl100k; do + uv run scripts/gen_tokenizer_oracle.py --model $m --corpus /path/to/78785.md \ + --out /tmp/tokenizer_oracle/$m.json --fuzz 300 +done +TOKENIZER_ORACLE_DIR=/tmp/tokenizer_oracle \ + ./build.sh --test -run TestCountersMatchOracle ./internal/tokenizer/ + +# optional second opinion for xlmr-spm: Google's sentencepiece on the .model file. +# Expect disagreement on whitespace-shaped samples - see +# "The three-way cross-check (what it is, and how to read it)" above. +uv run scripts/gen_tokenizer_oracle.py --model bge-m3 --corpus /path/to/78785.md \ + --out /tmp/tokenizer_oracle/bge-m3.json --fuzz 300 --spm-crosscheck + +# build, then end-to-end: re-ingest a document and check MySQL vs ES +bash build.sh --go +bash scripts/restart_ingestor.sh + +# how much of the catalog the tags cover (regenerate the tables in the strategy section) +python3 scripts/tokenizer_coverage.py --markdown + +# asset plumbing: download list vs runtime images vs SHA-1 pins +python3 -m pytest ragflow_deps/test_tokenizer_assets.py -q +``` + +Expected output today: all packages `ok`; the oracle test logs +`agreement: 341/341 samples compared by token texts` for `xlmr-spm`, `bert-wordpiece`, +`qwen-bpe` and `llama-bpe`, `341/341 samples compared by id sequences` for +`cl100k_base`, plus a line naming how many samples were held only to the bounded +combining-mark tolerance, and `counters exercised: [cl100k_base xlmr-spm +bert-wordpiece qwen-bpe llama-bpe]` for the property test. A counter that is skipped +there is missing an asset, not exempt from the contract. + +Gate status at the time of writing (2026-09-16), so nobody reads a plan as a result: + +| Gate | Status | +|---|---| +| Unit + property + oracle tiers | **observed green** | +| Incident document re-ingested, MySQL `chunk_num` == ES `_count` | **observed** (36 == 36) | +| `covsweep` reporting zero unserved words under the strict criterion | **not yet run** (the earlier sweep was interrupted); command: `covsweep --docs /tmp/cov_docs.tsv --shards 8 --workers 16 --out outputs/cov_after_20260916` | + +## Background: the incident + +`78785.md` (77 KB of Wikipedia tables) embedded fine until the parser's separator +fix changed how the issue-number tables are rendered. It then failed +deterministically with `400 … {"code":20015,"message":"The parameter is invalid"}` +at the Tokenizer's content-embedding step, and because one bad input failed the +whole batch, the document was left with **0 chunks** — 34 lost. + +Measured, on the exact span that failed: + +| | tokens | +|---|---| +| `cl100k_base` (what the old guard counted with) | 8,143 → below the 8,182 guard, so **truncation never fired** | +| `bge-m3`'s own tokenizer | **8,237** → above the 8192 window → provider rejects | + +The same document now ingests in 2 s: `run=3`, `chunk_num=36`, +`token_num=43581`, ES `_count=36`, and the shrink-and-retry path did **not** have to +fire — the primary path is exact. diff --git a/internal/tokenizer/embedlimit.go b/internal/tokenizer/embedlimit.go new file mode 100644 index 0000000000..e4b669dddb --- /dev/null +++ b/internal/tokenizer/embedlimit.go @@ -0,0 +1,604 @@ +// +// 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 + +// Embedding input limits: counters, the margin, and calibration. +// +// The invariant this file exists to enforce: +// +// Text handed to an embedding API must be counted with the embedding model's +// OWN tokenizer, or with a CALIBRATED UPPER BOUND of it. Counting with +// cl100k_base alone is never sufficient. +// +// Why it matters, concretely: cl100k_base (tiktoken) and a model's own tokenizer +// disagree by roughly +/-2% and the sign depends on the content. On repetitive +// numeric tables cl100k compresses far harder than XLM-R SentencePiece, so a +// 25,000-character chunk that cl100k scores at 8,143 tokens (just under the +// 8,182-token guard) can be 8,250 tokens in the model's own tokenizer. The +// provider then rejects the whole request ("400 ... code 20015 The parameter is +// invalid" on SiliconFlow) and, before this change, the whole document failed +// with it. See internal/tokenizer/embedding_token_limits.md. + +import ( + "fmt" + "math" + "sort" + "strings" + "sync" +) + +// Counter is one model tokenizer. Implementations must be exact for their own +// tokenizer: TrimToLimit returns a prefix whose Count is <= limit. +type Counter interface { + // ID is the stable identifier shared with the model catalog + // (conf/all_models.json "tokenizer"). + ID() string + // Count is the number of tokens text occupies for this tokenizer. + Count(text string) int + // TrimToLimit returns the longest prefix of text with Count(prefix) <= limit. + // A limit <= 0 yields "". + TrimToLimit(text string, limit int) string + // Available reports whether the tokenizer's data loaded. An unavailable + // counter must not be trusted: an encoder that failed to load silently + // reports 0 tokens for everything (see bpe_loader.go). + Available() bool +} + +// Counter ids. These are the values accepted in the model catalog's "tokenizer" +// field; anything else (or a missing asset) resolves to the calibrated fallback. +// +// Only ids that actually have a loader are declared: a phantom id would silently +// become a calibrated cl100k count, which is safe but surprising. Adding a family +// means adding its constant together with its loader (see embedding_token_limits.md). +const ( + CounterCL100K = "cl100k_base" + CounterXLMRSentence = "xlmr-spm" + CounterBERTWordPiece = "bert-wordpiece" + CounterQwenBPE = "qwen-bpe" + CounterLlamaBPE = "llama-bpe" +) + +// --------------------------------------------------------------------------- +// Registry +// --------------------------------------------------------------------------- + +var counterRegistry = struct { + sync.RWMutex + byID map[string]Counter + ctor map[string]func() (Counter, error) +}{byID: make(map[string]Counter), ctor: make(map[string]func() (Counter, error))} + +// RegisterCounter installs a ready counter under id, replacing any previous one. +func RegisterCounter(id string, c Counter) { + if id == "" || c == nil { + return + } + counterRegistry.Lock() + defer counterRegistry.Unlock() + counterRegistry.byID[id] = c +} + +// RegisterCounterLoader installs a lazy loader for id. The loader runs at most +// once per process; a loader that fails (asset missing) is remembered as +// failed so a broken asset does not re-read the disk for every document. +func RegisterCounterLoader(id string, load func() (Counter, error)) { + if id == "" || load == nil { + return + } + counterRegistry.Lock() + defer counterRegistry.Unlock() + if _, ok := counterRegistry.byID[id]; ok { + return + } + counterRegistry.ctor[id] = load +} + +// CounterByID returns the counter for id when it is registered and usable. +func CounterByID(id string) (Counter, bool) { + if id == "" { + return nil, false + } + counterRegistry.RLock() + c, ok := counterRegistry.byID[id] + ctor := counterRegistry.ctor[id] + counterRegistry.RUnlock() + if ok && c != nil { + return c, c.Available() + } + if ctor != nil { + loaded, err := ctor() + counterRegistry.Lock() + delete(counterRegistry.ctor, id) + if err == nil && loaded != nil { + counterRegistry.byID[id] = loaded + } + counterRegistry.Unlock() + if err == nil && loaded != nil { + return loaded, loaded.Available() + } + } + return nil, false +} + +// CounterStatus is one line of the availability report a process can log at startup. +type CounterStatus struct { + ID string + Available bool + Source string // the file that was loaded, when the counter knows it +} + +// sourcePather is implemented by the counters that know which file they loaded. +type sourcePather interface{ SourcePath() string } + +// CounterStatuses reports every counter the model catalog can reference, so a process can +// say once, at startup, whether the exact counters are usable in this deployment. +// +// This is the earliest sign of a missing asset: the ingest path fails loudly on the first +// document whose model declares an unavailable counter (it refuses to substitute the +// calibrated cl100k estimate, which under-counts some tokenizers - see +// embedding_token_limits.md), so this report names the asset to restore before that happens. +func CounterStatuses() []CounterStatus { + ids := []string{CounterCL100K, CounterXLMRSentence, CounterBERTWordPiece, CounterQwenBPE, CounterLlamaBPE} + out := make([]CounterStatus, 0, len(ids)) + for _, id := range ids { + counter, ok := CounterByID(id) + if !ok || counter == nil { + out = append(out, CounterStatus{ID: id}) + continue + } + status := CounterStatus{ID: id, Available: true} + if pathAware, ok := counter.(sourcePather); ok { + status.Source = pathAware.SourcePath() + } + out = append(out, status) + } + return out +} + +// ResolveCounter maps a catalog tokenizer id onto a counter. Unknown ids, empty +// ids and unloadable assets all resolve to cl100k_base, which the caller compensates +// for with a calibrated ratio (see Limiter). Callers that must NOT degrade - an embedder +// whose model declares a tokenizer - check CounterExact first and refuse instead +// (internal/ingestion/task/embedder.go). +func ResolveCounter(id string) Counter { + if c, ok := CounterByID(id); ok { + return c + } + if c, ok := CounterByID(CounterCL100K); ok { + return c + } + return unavailableCounter{id: id} +} + +// CounterExact reports whether id names a counter that is loaded and therefore +// safe to use without a calibrated ratio. +func CounterExact(id string) bool { + _, ok := CounterByID(id) + return ok +} + +// --------------------------------------------------------------------------- +// The declared limit and the margin +// --------------------------------------------------------------------------- + +const ( + // EmbeddingMarginRatio is the safety margin kept below the model's declared + // input limit. It is a *ratio*, not a constant: the previous 10-token + // margin was 0.12% of an 8192-token limit, an order of magnitude smaller + // than the disagreement between two tokenizers. + EmbeddingMarginRatio = 0.02 + // EmbeddingMarginFloor keeps the ratio from collapsing for small models. + EmbeddingMarginFloor = 32 + // EmbeddingTokenLimitDefault is used when a model declares no limit at all. + // The provider catalog's context_length is the preferred source; this is + // the last resort and is deliberately the smallest common embedding window + // rather than 8192, because overshooting a model's window is a hard 400 + // while undershooting only truncates. + EmbeddingTokenLimitDefault = 2048 + // DefaultUncountedRatioUpper is the starting upper bound for real/own when + // the model's tokenizer is not available. The L3 shrink-and-retry path + // catches whatever this misses; being slightly too conservative here only + // truncates a few percent more text. + DefaultUncountedRatioUpper = 1.05 +) + +// EmbeddingTokenLimit returns how many tokens of text may be sent to a model +// whose declared input limit is maxTokens. +func EmbeddingTokenLimit(maxTokens int) int { + if maxTokens <= 0 { + return 0 + } + margin := int(math.Ceil(float64(maxTokens) * EmbeddingMarginRatio)) + if margin < EmbeddingMarginFloor { + margin = EmbeddingMarginFloor + } + if margin >= maxTokens { + // Tiny windows: keep at least one token instead of going negative. + if maxTokens <= 1 { + return maxTokens + } + return maxTokens / 2 + } + return maxTokens - margin +} + +// ResolveEmbeddingMaxTokens picks the declared input limit for a model, in the +// order the rest of the system already uses elsewhere: an explicit model value +// first, then the provider catalog's context_length, then the default. It exists +// because defaulting to a hard-coded 8192 overshoots the window of every model +// with a smaller one (the catalog contains 512-token embedding models), and an +// overshoot is a rejected request rather than a truncated one. +func ResolveEmbeddingMaxTokens(declared, contextLength int) int { + for _, candidate := range []int{declared, contextLength} { + if candidate > 0 { + return candidate + } + } + return EmbeddingTokenLimitDefault +} + +// --------------------------------------------------------------------------- +// Calibration (L2) +// --------------------------------------------------------------------------- + +// calibrationEntry is what we have learned about one (provider instance, model) +// pair's tokenizer ratio: real_tokens / own_tokens. Only ever ratchets up. +type calibrationEntry struct { + ratio float64 + samples int + limitRejects int +} + +// calibrationKey identifies a tokenizer's owner: it must include the endpoint and +// the model, because "bge-m3" tokenizes differently from "bge-m3 @ another +// provider" only in what the provider accepts, not in the tokenizer — but the +// *limit* differs per deployment, so the key stays per provider instance. +type calibrationKey string + +// Calibration learns the true/own token ratio per model from provider usage. +type Calibration struct { + mu sync.RWMutex + entries map[calibrationKey]*calibrationEntry + // defaultRatio is the starting upper bound before any observation. + defaultRatio float64 +} + +var defaultCalibration = NewCalibration(DefaultUncountedRatioUpper) + +// NewCalibration creates an empty calibration. A non-positive default falls back +// to DefaultUncountedRatioUpper. +func NewCalibration(defaultRatio float64) *Calibration { + if defaultRatio <= 0 { + defaultRatio = DefaultUncountedRatioUpper + } + return &Calibration{entries: make(map[calibrationKey]*calibrationEntry), defaultRatio: defaultRatio} +} + +// DefaultCalibration is the process-wide calibration used by the ingest path. +func DefaultCalibration() *Calibration { return defaultCalibration } + +// ObserveUsage records a successful call: the provider reported realCount tokens +// for text our counter scored at ownCount. +func (c *Calibration) ObserveUsage(key string, ownCount, realCount int) { + if c == nil || ownCount <= 0 || realCount <= 0 { + return + } + ratio := float64(realCount) / float64(ownCount) + c.mu.Lock() + defer c.mu.Unlock() + e := c.entryLocked(key) + e.samples++ + if ratio > e.ratio { + e.ratio = ratio + } +} + +// ObserveOverLimit records that the provider rejected a request for exceeding +// the model's input window while our own counter scored **the offending input** +// at ownCount tokens against a declared window of maxTokens. That observation is +// enough to raise the ratio bound: the true count is above maxTokens, so the true +// ratio exceeds maxTokens/ownCount. This is how a rejection teaches the next +// attempt to be more conservative instead of repeating the same failure. +// +// ownCount must be a SINGLE input's count, not a batch total: the window bounds +// each input individually, and a batch total is normally above the window, which +// would imply a ratio below 1 and teach nothing. +func (c *Calibration) ObserveOverLimit(key string, ownCount, maxTokens int) { + if c == nil || ownCount <= 0 || maxTokens <= 0 { + return + } + // A little headroom on top of the implied bound: the observed count is a + // lower bound on the true count, and we want the next attempt to pass. + ratio := float64(maxTokens) / float64(ownCount) * 1.01 + c.mu.Lock() + defer c.mu.Unlock() + e := c.entryLocked(key) + e.limitRejects++ + if ratio > e.ratio { + e.ratio = ratio + } +} + +// RatioUpper returns the upper bound of real/own for key, never below 1. +func (c *Calibration) RatioUpper(key string) float64 { + if c == nil { + return DefaultUncountedRatioUpper + } + c.mu.RLock() + e := c.entries[calibrationKey(key)] + c.mu.RUnlock() + if e == nil || e.ratio <= 0 { + return c.defaultRatio + } + if e.ratio < 1 { + return 1 + } + return e.ratio +} + +// Stats reports what has been learned, for logging and tests. +func (c *Calibration) Stats(key string) (ratio float64, samples, limitRejects int, ok bool) { + if c == nil { + return 0, 0, 0, false + } + c.mu.RLock() + defer c.mu.RUnlock() + e := c.entries[calibrationKey(key)] + if e == nil { + return c.defaultRatio, 0, 0, false + } + // The ratio is computed inline rather than through RatioUpper: that helper + // takes the read lock again, and a writer waiting between the two RLock calls + // blocks the second one (Go's RWMutex starves new readers once a writer is + // waiting) while it waits for the first lock to be released - a deadlock. + ratio = c.defaultRatio + if e.ratio > 0 { + ratio = e.ratio + } + if ratio < 1 { + ratio = 1 + } + return ratio, e.samples, e.limitRejects, true +} + +// Reset drops what has been learned for key (used when a model or its tokenizer +// configuration changes). +func (c *Calibration) Reset(key string) { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + delete(c.entries, calibrationKey(key)) +} + +// Keys lists the observed keys, sorted, for diagnostics. +func (c *Calibration) Keys() []string { + if c == nil { + return nil + } + c.mu.RLock() + defer c.mu.RUnlock() + out := make([]string, 0, len(c.entries)) + for k := range c.entries { + out = append(out, string(k)) + } + sort.Strings(out) + return out +} + +func (c *Calibration) entryLocked(key string) *calibrationEntry { + e := c.entries[calibrationKey(key)] + if e == nil { + // Seed from the configured upper bound, never below 1. The calibration + // only ratchets UP, so seeding a fresh entry at 1 would drop the safety + // margin the moment any observation arrives, and a later input with a + // different token distribution could then be under-counted and rejected. + e = &calibrationEntry{ratio: math.Max(c.defaultRatio, 1)} + c.entries[calibrationKey(key)] = e + } + return e +} + +// --------------------------------------------------------------------------- +// Limiter: counter + calibration + margin, in one place +// --------------------------------------------------------------------------- + +// Limiter turns "the model accepts maxTokens" into "send at most this much text". +type Limiter struct { + counter Counter + // ratio is the upper bound of real/own. It is 1 for an exact counter and is + // only consulted when cal is nil. + ratio float64 + // cal, when set, is read on every call rather than captured, so an + // over-limit rejection observed by a retry immediately tightens the budget + // of the next attempt instead of the next document. + cal *Calibration + key string +} + +// NewExactLimiter is for counters that ARE the model's tokenizer. +func NewExactLimiter(counter Counter) Limiter { + return Limiter{counter: counter, ratio: 1} +} + +// NewCalibratedLimiter is for everything else, including counters that are only +// an approximation: the ratio comes from Calibration. +func NewCalibratedLimiter(counter Counter, key string, cal *Calibration) Limiter { + if cal == nil { + cal = defaultCalibration + } + return Limiter{counter: counter, cal: cal, key: key} +} + +// LimiterFor resolves the counter for a catalog tokenizer id and binds the +// right ratio: exact when the model's own tokenizer is loaded, calibrated +// otherwise. +func LimiterFor(tokenizerID, calibrationKey string, cal *Calibration) Limiter { + if CounterExact(tokenizerID) { + return NewExactLimiter(ResolveCounter(tokenizerID)) + } + return NewCalibratedLimiter(ResolveCounter(tokenizerID), calibrationKey, cal) +} + +// Counter returns the underlying counter. +func (l Limiter) Counter() Counter { return l.counter } + +// Ratio returns the upper bound of real/own in use. +func (l Limiter) Ratio() float64 { + if l.cal != nil { + r := l.cal.RatioUpper(l.key) + if r < 1 { + return 1 + } + return r + } + if l.ratio < 1 { + return 1 + } + return l.ratio +} + +// Calibration exposes the calibration a calibrated limiter reads from, so the +// retry path can record an over-limit rejection against the same key. +func (l Limiter) Calibration() (*Calibration, string) { return l.cal, l.key } + +// Limit is the number of tokens of *model* tokens we may send. +func (l Limiter) Limit(maxTokens int) int { + declared := ResolveEmbeddingMaxTokens(maxTokens, 0) + if declared <= 0 { + return 0 + } + budget := int(math.Floor(float64(declared) / l.Ratio())) + if budget < 1 { + budget = 1 + } + return EmbeddingTokenLimit(budget) +} + +// Trim cuts text down to the limiter's budget for maxTokens and reports both the +// trimmed text and its count in the limiter's own counter. +func (l Limiter) Trim(text string, maxTokens int) (string, int) { + limit := l.Limit(maxTokens) + if l.counter == nil || !l.counter.Available() { + // No usable counter: fall back to a byte-level bound that cannot be too + // generous for any BPE/SPM tokenizer we know of (>= 1 token per 4 + // bytes is the practical floor for text). This path only runs when the + // cl100k table itself is missing, which the loader already screams + // about. + return trimByBytes(text, limit), limit + } + trimmed := l.counter.TrimToLimit(text, limit) + return trimmed, l.counter.Count(trimmed) +} + +// trimByBytes is the last-resort bound used when no counter is available. +func trimByBytes(text string, limit int) string { + if limit <= 0 { + return "" + } + // One byte per token is the only bound that holds for every tokenizer (a + // token consumes at least one byte), so `limit` bytes cannot exceed `limit` + // tokens. A larger byte budget assumes more bytes per token than a worst-case + // input provides, which would let the provider reject the input for exceeding + // its window - exactly the failure this fallback exists to avoid. + maxBytes := limit + if len(text) <= maxBytes { + return text + } + cut := maxBytes + for cut > 0 && !utf8Start(text[cut]) { + cut-- + } + return text[:cut] +} + +// utf8Start reports whether b is not a UTF-8 continuation byte. +func utf8Start(b byte) bool { return b&0xC0 != 0x80 } + +// --------------------------------------------------------------------------- +// Counters we can build without extra assets +// --------------------------------------------------------------------------- + +type cl100kCounter struct{ id string } + +// CountCL100K is the cl100k_base counter, backed by the BPE table shipped in +// ragflow_deps (see bpe_loader.go). +func CountCL100K() Counter { return cl100kCounter{id: CounterCL100K} } + +func (c cl100kCounter) ID() string { return c.id } + +func (c cl100kCounter) Count(text string) int { return NumTokensFromString(text) } + +// IDs reports the token ids cl100k_base assigns to text, for the oracle test's +// id-level comparison. +func (c cl100kCounter) IDs(text string) []int32 { + enc, err := getCL100KEncoder() + if err != nil || enc == nil { + return nil + } + tokens := enc.Encode(text, nil, nil) + ids := make([]int32, len(tokens)) + for i, tok := range tokens { + ids[i] = int32(tok) + } + return ids +} + +func (c cl100kCounter) TrimToLimit(text string, limit int) string { + return TrimContentToTokenLimit(text, limit) +} + +func (c cl100kCounter) Available() bool { + enc, err := getCL100KEncoder() + return err == nil && enc != nil +} + +// SourcePath is the table file the loader accepted, when it got that far. +func (c cl100kCounter) SourcePath() string { return cl100kTableSource() } + +// unavailableCounter is returned when nothing else could be built. It reports +// zero tokens, so callers must check Available() before trusting a count; the +// Limiter does. +type unavailableCounter struct{ id string } + +func (c unavailableCounter) ID() string { + if c.id == "" { + return "unavailable" + } + return c.id +} + +func (unavailableCounter) Count(string) int { return 0 } + +func (unavailableCounter) TrimToLimit(text string, limit int) string { + return trimByBytes(text, limit) +} + +func (unavailableCounter) Available() bool { return false } + +func init() { + RegisterCounter(CounterCL100K, CountCL100K()) +} + +// DescribeCounter renders a counter for logs. +func DescribeCounter(c Counter) string { + if c == nil { + return "counter=" + } + return fmt.Sprintf("counter=%s(available=%t)", strings.TrimSpace(c.ID()), c.Available()) +} diff --git a/internal/tokenizer/embedlimit_test.go b/internal/tokenizer/embedlimit_test.go new file mode 100644 index 0000000000..be5f576213 --- /dev/null +++ b/internal/tokenizer/embedlimit_test.go @@ -0,0 +1,272 @@ +// +// 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 ( + "strings" + "testing" +) + +// fakeCounter counts one unit per N runes, so tests can exercise the limit maths +// without depending on a shipped asset. +type fakeCounter struct { + id string + perRune int + available bool +} + +func (f fakeCounter) ID() string { return f.id } + +func (f fakeCounter) Count(s string) int { + if f.perRune <= 0 { + return 0 + } + return len([]rune(s)) * f.perRune +} + +func (f fakeCounter) TrimToLimit(s string, limit int) string { + if limit <= 0 { + return "" + } + maxRunes := limit / f.perRune + r := []rune(s) + if len(r) <= maxRunes { + return s + } + return string(r[:maxRunes]) +} + +func (f fakeCounter) Available() bool { return f.available } + +func TestEmbeddingTokenLimitMargin(t *testing.T) { + cases := []struct { + maxTokens int + want int + }{ + // 2% of 8192 = 163.84 -> 164, well above the 32-token floor. + {8192, 8192 - 164}, + // 2% of 512 = 10.24 -> the 32-token floor wins. + {512, 512 - 32}, + // Degenerate windows keep a usable budget instead of going negative. + {32, 16}, + {2, 1}, + {0, 0}, + } + for _, c := range cases { + if got := EmbeddingTokenLimit(c.maxTokens); got != c.want { + t.Errorf("EmbeddingTokenLimit(%d) = %d, want %d", c.maxTokens, got, c.want) + } + } +} + +func TestResolveEmbeddingMaxTokens(t *testing.T) { + cases := []struct { + declared, contextLength, want int + }{ + {8192, 32768, 8192}, // the explicit model value wins + {0, 512, 512}, // catalog context_length, not a hard-coded 8192 + {0, 0, EmbeddingTokenLimitDefault}, + } + for _, c := range cases { + if got := ResolveEmbeddingMaxTokens(c.declared, c.contextLength); got != c.want { + t.Errorf("ResolveEmbeddingMaxTokens(%d,%d) = %d, want %d", c.declared, c.contextLength, got, c.want) + } + } +} + +func TestCalibrationRatchet(t *testing.T) { + cal := NewCalibration(DefaultUncountedRatioUpper) + key := "siliconflow|bge-m3" + if got := cal.RatioUpper(key); got != DefaultUncountedRatioUpper { + t.Fatalf("unobserved ratio = %v, want %v", got, DefaultUncountedRatioUpper) + } + // A successful call whose real count is lower than ours must not lower the + // bound: it stays at the configured default. The calibration only ratchets + // up, so one under-counting observation cannot replace the margin with an + // estimate of its own. + cal.ObserveUsage(key, 8182, 8027) + if got := cal.RatioUpper(key); got != DefaultUncountedRatioUpper { + t.Fatalf("ratio after an under-counting observation = %v, want the default %v", got, DefaultUncountedRatioUpper) + } + // A real count above ours ratchets the bound up. + cal.ObserveUsage(key, 1000, 1030) + if got := cal.RatioUpper(key); got < 1.03 { + t.Fatalf("ratio after ObserveUsage = %v, want >= 1.03", got) + } + ratio, samples, rejects, ok := cal.Stats(key) + if !ok || samples != 2 || rejects != 0 || ratio < 1.03 { + t.Fatalf("Stats = (%v,%d,%d,%t), want ratio >= 1.03, samples 2, rejects 0, ok", ratio, samples, rejects, ok) + } + cal.Reset(key) + if got := cal.RatioUpper(key); got != DefaultUncountedRatioUpper { + t.Fatalf("ratio after Reset = %v, want the default", got) + } +} + +// TestCalibrationOverLimitInfersRatio is the 78785.md case: our counter scored a +// chunk at 8,143 tokens against an 8192-token window and the provider rejected it +// for being over the limit. The rejection itself proves the true count is above +// 8192, so the ratio bound must rise above 1 even though no usage was returned. +func TestCalibrationOverLimitInfersRatio(t *testing.T) { + cal := NewCalibration(1.0) + key := "siliconflow|bge-m3" + cal.ObserveOverLimit(key, 8143, 8192) + ratio := cal.RatioUpper(key) + want := 8192.0 / 8143.0 * 1.01 + if ratio < want-1e-9 { + t.Fatalf("ratio after an over-limit rejection = %v, want >= %v", ratio, want) + } + if _, _, rejects, _ := cal.Stats(key); rejects != 1 { + t.Fatalf("limitRejects = %d, want 1", rejects) + } +} + +func TestLimiterLimitUsesRatioAndMargin(t *testing.T) { + exact := NewExactLimiter(fakeCounter{id: "exact", perRune: 1, available: true}) + if got, want := exact.Limit(8192), EmbeddingTokenLimit(8192); got != want { + t.Fatalf("exact Limit(8192) = %d, want %d", got, want) + } + + cal := NewCalibration(1.0) + key := "p|m" + calibrated := NewCalibratedLimiter(fakeCounter{id: "approx", perRune: 1, available: true}, key, cal) + if got := calibrated.Limit(8192); got != EmbeddingTokenLimit(8192) { + t.Fatalf("ratio-1 calibrated Limit(8192) = %d, want %d", got, EmbeddingTokenLimit(8192)) + } + + // The limiter reads the calibration live: an over-limit rejection recorded + // while it is in flight must tighten the very next Limit() call. + cal.ObserveOverLimit(key, 8143, 8192) + tightened := calibrated.Limit(8192) + if tightened >= EmbeddingTokenLimit(8192) { + t.Fatalf("Limit after an over-limit observation = %d, want less than %d", tightened, EmbeddingTokenLimit(8192)) + } + budget := int(8192 / calibrated.Ratio()) + if tightened != EmbeddingTokenLimit(budget) { + t.Fatalf("Limit = %d, want EmbeddingTokenLimit(%d) = %d", tightened, budget, EmbeddingTokenLimit(budget)) + } +} + +func TestLimiterTrimRespectsLimit(t *testing.T) { + counter := fakeCounter{id: "fake", perRune: 2, available: true} + limiter := NewExactLimiter(counter) + text := strings.Repeat("a", 10000) + trimmed, tokens := limiter.Trim(text, 8192) + limit := limiter.Limit(8192) + if tokens > limit { + t.Fatalf("Trim returned %d tokens, limit is %d", tokens, limit) + } + if got := counter.Count(trimmed); got > limit { + t.Fatalf("counter reports %d tokens for the trimmed text, limit is %d", got, limit) + } + if len(trimmed) == len(text) { + t.Fatalf("expected the text to be cut, got the whole %d-byte input", len(text)) + } +} + +// TestLimiterWithoutUsableCounter covers the case where the cl100k table is +// missing: the byte-level bound must still cut the text and must never report a +// token count that was invented from a dead encoder. +func TestLimiterWithoutUsableCounter(t *testing.T) { + limiter := NewExactLimiter(fakeCounter{id: "dead", perRune: 1, available: false}) + text := strings.Repeat("x", 1_000_000) + trimmed, tokens := limiter.Trim(text, 8192) + limit := limiter.Limit(8192) + // The fallback bounds BYTES by the token limit: one byte per token is the + // only budget that cannot exceed it. + if len(trimmed) > limit { + t.Fatalf("byte-level fallback kept %d bytes, want <= %d", len(trimmed), limit) + } + if tokens != limit { + t.Fatalf("fallback token count = %d, want the limit %d", tokens, limit) + } + if strings.ContainsRune(trimmed, '\uFFFD') { + t.Fatal("byte-level fallback split a multi-byte rune") + } +} + +func TestTrimByBytesKeepsRunesIntact(t *testing.T) { + text := strings.Repeat("中", 100) // 300 bytes + got := trimByBytes(text, 10) // 10-byte bound + if len(got) > 10 { + t.Fatalf("trimByBytes kept %d bytes, want <= 10", len(got)) + } + if !strings.HasPrefix(text, got) { + t.Fatal("trimByBytes did not return a prefix") + } + for _, r := range got { + if r == '\uFFFD' { + t.Fatal("trimByBytes cut a rune in half") + } + } +} + +func TestResolveCounterFallsBackForUnknownID(t *testing.T) { + // Loading cl100k here is fine: failfast_test.go re-execs itself in a fresh + // process precisely so that tiktoken-go's process-global encoding cache + // cannot make that test order-dependent. + defer resetCL100KEncoderForTest() + for _, id := range []string{"", "not-a-tokenizer", CounterXLMRSentence} { + c := ResolveCounter(id) + if c == nil { + t.Fatalf("ResolveCounter(%q) returned nil", id) + } + // Whether cl100k's table is present in this environment or not, the + // resolved counter must be safe to call and must trim without panicking. + text := strings.Repeat("word ", 5000) + trimmed := c.TrimToLimit(text, 100) + if !strings.HasPrefix(text, trimmed) { + t.Fatalf("ResolveCounter(%q) did not return a prefix", id) + } + if strings.ContainsRune(trimmed, '\uFFFD') { + t.Fatalf("ResolveCounter(%q) split a rune while trimming", id) + } + if c.Available() { + if got := c.Count(trimmed); got > 100 { + t.Fatalf("ResolveCounter(%q) kept %d tokens for a 100-token limit", id, got) + } + continue + } + if len(trimmed) > 100*4 { + t.Fatalf("ResolveCounter(%q) kept %d bytes for a 100-token limit with no counter available", id, len(trimmed)) + } + } +} + +// TestCL100KTrimNeverExceedsLimit is the property every counter must hold: the +// trimmed text fits the limit according to the counter that produced it. +func TestCL100KTrimNeverExceedsLimit(t *testing.T) { + defer resetCL100KEncoderForTest() + counter := CountCL100K() + if !counter.Available() { + t.Skip("cl100k table not present in this environment") + } + samples := []string{ + strings.Repeat("hello world ", 500), + strings.Repeat("| 1976 | | 383/1 | 383/2 | 383/3 |\n", 200), + strings.Repeat("中", 2000), + strings.Repeat("QWxhZGRpbjpvcGVuIHNlc2FtZQ", 200), + } + for _, limit := range []int{1, 50, 512, 8028} { + for i, s := range samples { + trimmed := counter.TrimToLimit(s, limit) + if got := counter.Count(trimmed); got > limit { + t.Errorf("sample %d limit %d: trimmed text counts %d tokens", i, limit, got) + } + } + } +} diff --git a/internal/tokenizer/failfast_test.go b/internal/tokenizer/failfast_test.go index ac386299e0..4bb8abaa58 100644 --- a/internal/tokenizer/failfast_test.go +++ b/internal/tokenizer/failfast_test.go @@ -3,7 +3,7 @@ // // 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 the License at +// You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // @@ -17,13 +17,18 @@ package tokenizer import ( + "bytes" "io" "os" + "os/exec" "path/filepath" "strings" "testing" ) +// failFastChildEnv marks the re-executed child that owns the real assertions. +const failFastChildEnv = "TOKENIZER_FAILFAST_CHILD" + // TestInitCL100KEncoder_FailFast pins the contract that InitCL100KEncoder fails // fast on a missing cl100k_base table, and — the part the regression actually // guards — that a PRESENT table yields a working encoder (NumTokensFromString > 0) @@ -43,6 +48,26 @@ import ( // that the table "present" later copies under /tmp is not discoverable by it via // searchRoots() walking up to the shared /tmp ancestor. func TestInitCL100KEncoder_FailFast(t *testing.T) { + // This test's premise is a process with an EMPTY tiktoken cache. tiktoken-go + // memoises loaded encodings in a package-global map with no reset API + // (encodingMap in its encoding.go), so as soon as any sibling test has loaded + // cl100k successfully, "no table present must fail" can no longer be observed + // here - resetCL100KEncoderForTest() only clears RAGFlow's own wrapper. A test + // whose outcome depends on execution order is not a test, so re-exec the test + // binary and run the real checks in a fresh process instead. + if os.Getenv(failFastChildEnv) == "" { + child := exec.Command(os.Args[0], "-test.run=^TestInitCL100KEncoder_FailFast$", "-test.v") + child.Env = append(os.Environ(), failFastChildEnv+"=1") + out, err := child.CombinedOutput() + if err != nil { + t.Fatalf("isolated child run failed (%v); the fail-fast contract is only observable in a fresh process\n%s", err, out) + } + if !bytes.Contains(out, []byte("--- PASS")) { + t.Fatalf("isolated child run reported no PASS:\n%s", out) + } + return + } + // fail-fast: an empty scoped dir must surface a hard error, not a silent 0. t.Run("absent", func(t *testing.T) { resetCL100KEncoderForTest() diff --git a/internal/tokenizer/spm.go b/internal/tokenizer/spm.go new file mode 100644 index 0000000000..ca712052b9 --- /dev/null +++ b/internal/tokenizer/spm.go @@ -0,0 +1,956 @@ +// +// 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 + +// SentencePiece (Unigram) counter: the XLM-R family. +// +// BAAI/bge-m3, intfloat/multilingual-e5-*, moka-ai/m3e-*, Alibaba-NLP/gte- +// multilingual-base and jina-embeddings-v3 all tokenize with the same +// XLM-RoBERTa SentencePiece model (Unigram, 250k pieces, byte_fallback), which +// is the vocabulary shipped in ragflow_deps - so one asset and one implementation +// cover the whole family. See internal/tokenizer/embedding_token_limits.md. +// +// What is faithful and what is approximate: +// +// - pieces, scores, byte-fallback pieces and the whitespace flags are read +// straight from the .model file, so segmentation of the normalized text is +// the model's own; +// - normalization applies NFKC, which is the bulk of the model's `nmt_nfkc` +// normalizer. The asset additionally carries a precompiled Darts charmap +// (~232 KB of tables) that SPM applies before NFKC; this implementation does +// not, so text containing unusual Unicode can normalize slightly +// differently. Plain prose, tables, digits and CJK - i.e. essentially all +// ingested content - are unaffected, and the embedder also verifies its +// counting against the provider's reported usage (Calibration.ObserveUsage) +// and shrinks-and-retries if a request is nonetheless rejected. +// +// The counter therefore reports Available() as soon as the asset loads, and the +// ingest path treats it as the model's tokenizer while keeping that verification +// running. + +import ( + "crypto/sha1" + "encoding/binary" + "errors" + "fmt" + "math" + "os" + "path/filepath" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" + + "ragflow/internal/common" +) + +// SentencePiece piece types (model.proto SentencePiece.Type). +const ( + spmTypeNormal = 1 + spmTypeUnknown = 2 + spmTypeControl = 3 + spmTypeUser = 4 + spmTypeUnused = 5 + spmTypeByte = 6 +) + +// sentencePieceModel is a SentencePiece Unigram model, parsed from the .model +// protobuf. Only the fields the counter needs are decoded. +type sentencePieceModel struct { + pieces []spmPiece + trie *spmTrie + unkScore float32 + unkID int32 // index of the piece, -1 when the model has none + byteByValue map[byte]int32 // "<0xXX>" piece indexes + byteFallback bool + + addDummyPrefix bool + removeExtraWhitespaces bool + escapeWhitespaces bool + + normalizerName string + hasCharmap bool +} + +type spmPiece struct { + text string + score float32 + typ int32 +} + +// spmTrie is a byte trie over the vocabulary. It returns every piece that starts +// at a position, which is what Unigram's Viterbi needs (any matching length can +// be the optimal one). +type spmTrie struct { + nodes []spmTrieNode +} + +type spmTrieNode struct { + next map[byte]int32 + piece int32 + score float32 +} + +func newSPMTrie() *spmTrie { + t := &spmTrie{} + t.nodes = append(t.nodes, spmTrieNode{piece: -1}) + return t +} + +func (t *spmTrie) insert(piece string, idx int32, score float32) { + node := int32(0) + for i := 0; i < len(piece); i++ { + if t.nodes[node].next == nil { + t.nodes[node].next = make(map[byte]int32, 2) + } + next, ok := t.nodes[node].next[piece[i]] + if !ok { + t.nodes = append(t.nodes, spmTrieNode{piece: -1}) + next = int32(len(t.nodes) - 1) + t.nodes[node].next[piece[i]] = next + } + node = next + } + t.nodes[node].piece = idx + t.nodes[node].score = score +} + +// parseSentencePieceModel decodes what the counter needs from a .model file. +func parseSentencePieceModel(raw []byte) (*sentencePieceModel, error) { + m := &sentencePieceModel{ + unkID: -1, + byteByValue: make(map[byte]int32, 256), + // model.proto declares these `optional bool ... [default = true]`; proto + // semantics therefore treat an absent field as true. + addDummyPrefix: true, + removeExtraWhitespaces: true, + escapeWhitespaces: true, + } + modelType := int64(0) + + r := &pbReader{b: raw} + for { + field, wire, err := r.readKey() + if err != nil { + if errors.Is(err, errPBEOF) { + break + } + return nil, fmt.Errorf("sentencepiece: %w", err) + } + switch { + case field == 1 && wire == pbWireBytes: // ModelProto.pieces + pieceBytes, err := r.readBytes() + if err != nil { + return nil, fmt.Errorf("sentencepiece: piece: %w", err) + } + piece, err := parseSPMPiece(pieceBytes) + if err != nil { + return nil, err + } + m.pieces = append(m.pieces, piece) + case field == 2 && wire == pbWireBytes: // ModelProto.trainer_spec + spec, err := r.readBytes() + if err != nil { + return nil, fmt.Errorf("sentencepiece: trainer_spec: %w", err) + } + // TrainerSpec.model_type = 3 (1 = UNIGRAM). + sr := &pbReader{b: spec} + for { + f, w, err := sr.readKey() + if err != nil { + break + } + if f == 3 && w == pbWireVarint { + v, err := sr.readVarint() + if err != nil { + break + } + modelType = int64(v) + continue + } + if err := sr.skip(w); err != nil { + break + } + } + case field == 3 && wire == pbWireBytes: // ModelProto.normalizer_spec + spec, err := r.readBytes() + if err != nil { + return nil, fmt.Errorf("sentencepiece: normalizer_spec: %w", err) + } + if err := m.parseNormalizerSpec(spec); err != nil { + return nil, err + } + default: + if err := r.skip(wire); err != nil { + return nil, fmt.Errorf("sentencepiece: %w", err) + } + } + } + + if len(m.pieces) == 0 { + return nil, errors.New("sentencepiece: model has no pieces") + } + if modelType != 1 { + // 1 = UNIGRAM. A BPE model in the same protobuf needs a different + // encoder; refusing is better than counting with the wrong algorithm. + return nil, fmt.Errorf("sentencepiece: model_type %d is not Unigram", modelType) + } + + // The unknown score is the floor SPM uses for pieces it cannot match. + m.unkScore = math.MaxFloat32 + for i, p := range m.pieces { + if p.typ == spmTypeUnknown { + m.unkScore = p.score + m.unkID = int32(i) + } + } + if m.unkScore == math.MaxFloat32 { + minScore := m.pieces[0].score + for _, p := range m.pieces { + if p.score < minScore { + minScore = p.score + } + } + m.unkScore = minScore - 10 + } + + m.trie = newSPMTrie() + for i := range m.pieces { + p := &m.pieces[i] + if p.typ == spmTypeControl || p.typ == spmTypeUnknown || p.typ == spmTypeUnused { + continue + } + m.trie.insert(p.text, int32(i), p.score) + if p.typ == spmTypeByte && len(p.text) == 6 && strings.HasPrefix(p.text, "<0x") { + var v int + if _, err := fmt.Sscanf(p.text, "<0x%02X>", &v); err == nil && v >= 0 && v < 256 { + m.byteByValue[byte(v)] = int32(i) + } + } + } + m.byteFallback = len(m.byteByValue) == 256 + return m, nil +} + +func parseSPMPiece(raw []byte) (spmPiece, error) { + piece := spmPiece{typ: spmTypeNormal} + r := &pbReader{b: raw} + for { + field, wire, err := r.readKey() + if err != nil { + break + } + switch { + case field == 1 && wire == pbWireBytes: // SentencePiece.piece + b, err := r.readBytes() + if err != nil { + return piece, fmt.Errorf("sentencepiece: piece text: %w", err) + } + piece.text = string(b) + case field == 2 && wire == pbWireFixed32: // SentencePiece.score (float) + b, err := r.readFixed32() + if err != nil { + return piece, fmt.Errorf("sentencepiece: piece score: %w", err) + } + piece.score = math.Float32frombits(b) + case field == 3 && wire == pbWireVarint: // SentencePiece.type + v, err := r.readVarint() + if err != nil { + return piece, fmt.Errorf("sentencepiece: piece type: %w", err) + } + piece.typ = int32(v) + default: + if err := r.skip(wire); err != nil { + return piece, fmt.Errorf("sentencepiece: piece: %w", err) + } + } + } + return piece, nil +} + +func (m *sentencePieceModel) parseNormalizerSpec(raw []byte) error { + r := &pbReader{b: raw} + for { + field, wire, err := r.readKey() + if err != nil { + break + } + switch { + case field == 1 && wire == pbWireBytes: // name + b, err := r.readBytes() + if err != nil { + return fmt.Errorf("sentencepiece: normalizer name: %w", err) + } + m.normalizerName = string(b) + case field == 2 && wire == pbWireBytes: // precompiled_charsmap + b, err := r.readBytes() + if err != nil { + return fmt.Errorf("sentencepiece: charmap: %w", err) + } + m.hasCharmap = len(b) > 0 + case (field == 3 || field == 4 || field == 5) && wire == pbWireVarint: + v, err := r.readVarint() + if err != nil { + return fmt.Errorf("sentencepiece: normalizer flag: %w", err) + } + on := v != 0 + switch field { + case 3: + m.addDummyPrefix = on + case 4: + m.removeExtraWhitespaces = on + case 5: + m.escapeWhitespaces = on + } + default: + if err := r.skip(wire); err != nil { + return fmt.Errorf("sentencepiece: normalizer_spec: %w", err) + } + } + } + return nil +} + +// normalize mirrors the normalization the models are actually served with. +// +// Verified against the HuggingFace conversion of this exact vocabulary +// (scripts/gen_tokenizer_oracle.py), because that is what providers run: +// +// normalizer: Sequence(Precompiled(charmap), Replace(Regex(" {2,}"), " ")) +// pre_tokenizer: Metaspace(replacement="▁", prepend_scheme=always, split=true) +// +// Three details that are easy to get wrong, all pinned by the oracle test: +// +// - runs of two or more SPACES collapse to one, but whitespace is never +// trimmed: "hello world " keeps its trailing ▁ and costs one token more than +// a trimmed version. +// - any other whitespace run (\n, \t, \r\n) becomes a SINGLE ▁. +// - a leading ▁ is always prepended, except for empty input and when the text +// already starts with whitespace (the leading run fills that slot). +// +// DELIBERATELY not driven by the .model's flags. That file declares +// `remove_extra_whitespaces = true` (and would strip the trailing whitespace), +// but the tokenizer providers actually serve is the HuggingFace conversion, which +// drops that flag and keeps the Metaspace behavior above. Following the .model +// made this counter systematically one token short - 15 of 23 oracle samples - +// and the provider, not the file, decides what fits in the window. +// +// Normalization applies NFKC, which is the bulk of the model's `nmt_nfkc` +// normalizer; the precompiled charmap is not applied (see the file header). +func (m *sentencePieceModel) normalize(text string) string { + if text == "" { + return "" + } + s := collapseSpaces(normalizerLite(norm.NFKC.String(text))) + if s == "" { + // Everything the charmap deletes (e.g. a text of nothing but control + // characters) leaves nothing for the dummy prefix to attach to. Emitting a + // bare "▁" here would bill one token for text the model sees as empty. + return "" + } + const space = '\u2581' // ▁ + var b strings.Builder + b.Grow(len(s) + 2) + b.WriteRune(space) + prevSpace := false + for i, r := range s { + if unicode.IsSpace(r) { + // A leading whitespace run is absorbed by the dummy prefix. + if prevSpace || i == 0 { + prevSpace = true + continue + } + prevSpace = true + b.WriteRune(space) + continue + } + prevSpace = false + b.WriteRune(r) + } + return b.String() +} + +// normalizerLite applies the part of the model's precompiled `nmt_nfkc` charmap +// that NFKC does not cover, so the counter does not spend tokens on characters the +// served tokenizer deletes or turns into whitespace. +// +// The rules below are the ones scripts/gen_tokenizer_oracle.py exercises, each +// verified individually against the model (probe -> what bge-m3 does): +// +// U+0001..U+0008, U+000B, U+000C, U+000E..U+001F deleted +// U+007F..U+0084, U+0086..U+009F deleted (0x85 is NOT deleted) +// U+200B..U+200D, U+FEFF become a space +// U+2581 (the escaped space itself) becomes a space +// U+0000, U+0085, U+2060 kept, each as its own token +// +// This is deliberately a subset: the real charmap is a 232 KB Darts trie with +// hundreds of rules. What is not implemented here is unusual by construction - the +// probe above is what keeps this honest, and the ingest path additionally checks +// its counting against the provider's reported usage (Calibration.ObserveUsage) +// and shrinks-and-retries if a request is rejected anyway. +func normalizerLite(s string) string { + if !strings.ContainsFunc(s, func(r rune) bool { return isCharmapSpecial(r) }) { + return s + } + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + switch { + case isDeletedByCharmap(r): + continue + case isCharmapSpace(r): + b.WriteRune(' ') + default: + b.WriteRune(r) + } + } + return b.String() +} + +func isCharmapSpecial(r rune) bool { + return isDeletedByCharmap(r) || isCharmapSpace(r) +} + +// isCharmapSpace reports the characters nmt_nfkc maps to a plain space rather than +// deleting them. U+2581 is the escaped space itself: the served tokenizer turns a +// literal "▁" in the input into a space (verified: normalize_str("▁a▁▁b") = " a b"), +// and Metaspace turns it back into one "▁" per run. +func isCharmapSpace(r rune) bool { + switch r { + case '\u200B', '\u200C', '\u200D', '\uFEFF', '\u2581': + return true + default: + return false + } +} + +func isDeletedByCharmap(r rune) bool { + switch { + case r >= 0x0001 && r <= 0x0008: + return true + case r == 0x000B || r == 0x000C: + return true + case r >= 0x000E && r <= 0x001F: + return true + case r >= 0x007F && r <= 0x0084: + return true + case r >= 0x0086 && r <= 0x009F: + return true + default: + return false + } +} + +// collapseSpaces collapses runs of two or more spaces into one, mirroring the +// tokenizer's Replace(Regex(" {2,}"), " ") normalizer step. Other whitespace is +// left untouched. +func collapseSpaces(s string) string { + if !strings.Contains(s, " ") { + return s + } + var b strings.Builder + b.Grow(len(s)) + prevSpace := false + for _, r := range s { + if r == ' ' { + if prevSpace { + continue + } + prevSpace = true + b.WriteRune(r) + continue + } + prevSpace = false + b.WriteRune(r) + } + return b.String() +} + +// countTokens is the Unigram Viterbi count over the normalized text: the number +// of pieces in the highest-scoring segmentation. +func (m *sentencePieceModel) countTokens(text string) int { + s := m.normalize(text) + if s == "" { + return 0 + } + total := 0 + for _, chunk := range splitPreTokens(s) { + n, _ := m.viterbi(chunk, false) + total += n + } + return total +} + +// preTokenSpace is the escaped-space marker the model's normalizer emits and its +// Metaspace pre-tokenizer splits on. +const preTokenSpace = '\u2581' // ▁ + +// hasPieceAt reports whether any vocabulary piece starts at byte offset i, which is +// what decides whether the unknown-token run ends there. +func (m *sentencePieceModel) hasPieceAt(s string, i int) bool { + node := int32(0) + for j := i; j < len(s); j++ { + next, ok := m.trie.nodes[node].next[s[j]] + if !ok { + return false + } + node = next + if m.trie.nodes[node].piece >= 0 { + return true + } + } + return false +} + +// splitPreTokens splits the normalized text into the chunks the served tokenizer +// actually feeds to its Unigram model. +// +// This matters, and not only for cosmetics: bge-m3's tokenizer.json declares +// `Metaspace(add_prefix_space=true)` with the default `split=true`, so the model +// runs one Viterbi *per whitespace-delimited chunk*, not one over the whole text. +// A global Viterbi is free to pick a different (globally better-scoring) +// segmentation - the incident document has such a spot around byte 46,673, where +// the whole-document optimum splits "343/1" as "▁34"+"3" while the served +// per-chunk optimum splits it as "▁3"+"43" - so counting globally can disagree with +// the provider on both the pieces and, in principle, their number. Our normalized +// string is already the concatenation of the pre-tokens (every "▁" starts one), so +// splitting on "▁" reproduces the served pre-tokenization exactly. +func splitPreTokens(s string) []string { + if s == "" { + return nil + } + chunks := make([]string, 0, strings.Count(s, string(preTokenSpace))+1) + start := 0 + for i, r := range s { + if r == preTokenSpace && i > start { + chunks = append(chunks, s[start:i]) + start = i + } + } + return append(chunks, s[start:]) +} + +// tokenIDs returns the exact piece ids the model's encoder emits, one id per token +// (byte fallback contributes one id per byte, exactly as the model does). It backs +// the oracle test's id-level comparison: two different segmentations can share a +// token count, so "same count" is weaker evidence than "same tokens", and a +// count-neutral segmentation bug would otherwise stay invisible. +// +// The ids are the .model file's piece indexes. That is not always the numbering the +// served tokenizer uses - the HuggingFace conversion of bge-m3 reorders the specials +// and inserts /, which shifts every regular piece by one - so the oracle +// test compares piece *texts* (tokenNames) for this family and keeps the ids for +// local reasoning. +func (m *sentencePieceModel) tokenIDs(text string) []int32 { + s := m.normalize(text) + if s == "" { + return nil + } + var ids []int32 + for _, chunk := range splitPreTokens(s) { + _, chunkIDs := m.viterbi(chunk, true) + ids = append(ids, chunkIDs...) + } + return ids +} + +// tokenNames returns the piece text of every token, which is numbering-independent +// and therefore directly comparable with the served tokenizer's own tokens. +func (m *sentencePieceModel) tokenNames(text string) []string { + ids := m.tokenIDs(text) + if len(ids) == 0 { + return nil + } + names := make([]string, 0, len(ids)) + for _, id := range ids { + if id >= 0 && int(id) < len(m.pieces) { + names = append(names, m.pieces[id].text) + continue + } + names = append(names, "") + } + return names +} + +// Sentinels for the Viterbi backpointers. A real entry holds a piece index, which +// is never negative, so negative values can mark the two fallbacks. +const ( + spmViterbiUnk int32 = -1 // the piece + spmViterbiByteRun int32 = -2 // one <0xXX> piece per byte of the rune +) + +// viterbi runs Unigram's Viterbi over s, returning the best token count and, when +// wantIDs is set, the chosen piece ids. +// +// The id bookkeeping is skipped on the count path: Count runs once per chunk on the +// ingest hot path, so the backpointer arrays are only allocated when a test asks +// for them. Both paths share this one function so they cannot drift apart. +func (m *sentencePieceModel) viterbi(s string, wantIDs bool) (int, []int32) { + n := len(s) + // spmNegInf stands in for -Inf: piece scores are small negative numbers, so + // the most negative float64 is unreachable by any real path. + // + // The accumulator is float64 even though the piece scores are float32: summing + // a few hundred float32s loses enough precision to flip ties between + // equal-scoring segmentations, and the reference implementation resolves those + // ties the other way. The incident-free case that exposed this is a run of 99 + // "a"s, where "▁a"+"a"+19×"aaaaa"+"aab" and "▁a"+2×"aaaaa"+"a"+... score the + // same; counting is unaffected, but a token-level comparison sees it. + const spmNegInf float64 = -math.MaxFloat64 + best := make([]float64, n+1) + count := make([]int, n+1) + var prev, pieceAt []int32 + if wantIDs { + prev = make([]int32, n+1) + pieceAt = make([]int32, n+1) + } + for i := range best { + best[i] = spmNegInf + } + best[0] = 0 + for i := 0; i < n; i++ { + if best[i] == spmNegInf { + continue + } + matched := false + node := int32(0) + for j := i; j < n; j++ { + next, ok := m.trie.nodes[node].next[s[j]] + if !ok { + break + } + node = next + if piece := m.trie.nodes[node].piece; piece >= 0 { + matched = true + cand := best[i] + float64(m.pieces[piece].score) + if cand > best[j+1] { + best[j+1] = cand + count[j+1] = count[i] + 1 + if wantIDs { + prev[j+1], pieceAt[j+1] = int32(i), piece + } + } + } + } + if matched { + continue + } + // No piece starts here. The served tokenizer consumes the whole RUN of + // characters that have no piece and bills it as ONE unknown token - not one + // per character: "🧑🧑" is a single token, five NULs are a single token, and + // a run only breaks at a character that does have a piece or at a + // pre-token boundary. Counting per character over-counts exactly this shape + // (it is what the fuzz corpus found). + runEnd := i + for runEnd < n { + _, size := utf8.DecodeRuneInString(s[runEnd:]) + if size <= 0 || m.hasPieceAt(s, runEnd) { + break + } + runEnd += size + } + if runEnd == i { + // Unreachable: the loop above always advances by at least one rune. + return count[i], nil + } + if m.byteFallback { + score := float64(0) + for k := i; k < runEnd; k++ { + if idx, ok := m.byteByValue[s[k]]; ok { + score += float64(m.pieces[idx].score) + } else { + score += float64(m.unkScore) + } + } + if cand := best[i] + score; cand > best[runEnd] { + best[runEnd] = cand + count[runEnd] = count[i] + (runEnd - i) // one <0xXX> per byte + if wantIDs { + prev[runEnd], pieceAt[runEnd] = int32(i), spmViterbiByteRun + } + } + continue + } + if cand := best[i] + float64(m.unkScore); cand > best[runEnd] { + best[runEnd] = cand + count[runEnd] = count[i] + 1 + if wantIDs { + prev[runEnd], pieceAt[runEnd] = int32(i), spmViterbiUnk + } + } + } + if best[n] == spmNegInf { + // Should not happen (the fallback always advances), but never claim 0 + // tokens for non-empty text: an under-count is the dangerous direction. + return n, nil + } + if !wantIDs { + return count[n], nil + } + ids := make([]int32, 0, count[n]) + for pos := n; pos > 0; { + switch pieceAt[pos] { + case spmViterbiUnk: + ids = append(ids, m.unkID) + case spmViterbiByteRun: + for k := int(prev[pos]); k < pos; k++ { + if idx, ok := m.byteByValue[s[k]]; ok { + ids = append(ids, idx) + } else { + ids = append(ids, m.unkID) + } + } + default: + ids = append(ids, pieceAt[pos]) + } + pos = int(prev[pos]) + } + for i, j := 0, len(ids)-1; i < j; i, j = i+1, j-1 { + ids[i], ids[j] = ids[j], ids[i] + } + return count[n], ids +} + +// trimToLimit returns the longest rune prefix of text with at most limit tokens. +func (m *sentencePieceModel) trimToLimit(text string, limit int) string { + if limit <= 0 { + return "" + } + if m.countTokens(text) <= limit { + return text + } + runes := []rune(text) + lo, hi := 0, len(runes) + for lo < hi { + mid := (lo + hi + 1) / 2 + if m.countTokens(string(runes[:mid])) <= limit { + lo = mid + } else { + hi = mid - 1 + } + } + return string(runes[:lo]) +} + +// sentencePieceCounter adapts a parsed model to the Counter interface. +type sentencePieceCounter struct { + model *sentencePieceModel + path string +} + +func (c *sentencePieceCounter) ID() string { return CounterXLMRSentence } + +// IDs reports the token ids this counter would emit. It is used by the oracle test +// to compare segmentations, not just their lengths; production code only needs +// Count and TrimToLimit. +func (c *sentencePieceCounter) IDs(text string) []int32 { return c.model.tokenIDs(text) } + +// Names reports the piece text of every token; the oracle test prefers this over +// ids because the served tokenizer renumbers this model's specials. +func (c *sentencePieceCounter) Names(text string) []string { return c.model.tokenNames(text) } + +// SourcePath is the `.model` file this counter loaded. +func (c *sentencePieceCounter) SourcePath() string { return c.path } + +func (c *sentencePieceCounter) Count(text string) int { return c.model.countTokens(text) } + +func (c *sentencePieceCounter) TrimToLimit(text string, limit int) string { + return c.model.trimToLimit(text, limit) +} + +func (c *sentencePieceCounter) Available() bool { return c.model != nil } + +// expectedSPMHashes pins the on-disk SentencePiece asset. A mismatch means a +// corrupt or tampered file: refuse to load it rather than counting with it. +var expectedSPMHashes = map[string]string{ + "BAAI/bge-m3/sentencepiece.bpe.model": "7e88c49faff6c6c136fdf4a3402d0cb534c6ab10", +} + +// xlmrSPMAssetNames are the file names the loader accepts, in priority order. +var xlmrSPMAssetNames = []string{ + "ragflow_deps/huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model", + "sentencepiece.bpe.model", + "xlmr_spm.model", +} + +var sentencePieceOnce struct { + sync.Once + counter Counter + err error +} + +// LoadXLMRSentencePieceCounter loads (once) the XLM-R SentencePiece counter. A +// missing asset is not fatal: callers fall back to a calibrated cl100k count. +func LoadXLMRSentencePieceCounter() (Counter, error) { + sentencePieceOnce.Do(func() { + sentencePieceOnce.counter, sentencePieceOnce.err = loadXLMRSentencePieceCounter() + }) + return sentencePieceOnce.counter, sentencePieceOnce.err +} + +func loadXLMRSentencePieceCounter() (Counter, error) { + path, raw, err := readTokenizerAsset(xlmrSPMAssetNames, spmAssetPinKey, expectedSPMHashes) + if err != nil { + return nil, err + } + model, err := parseSentencePieceModel(raw) + if err != nil { + return nil, fmt.Errorf("sentencepiece asset %s: %w", path, err) + } + return &sentencePieceCounter{model: model, path: path}, nil +} + +// spmAssetPinKey names the pinned file in expectedSPMHashes; the pin applies +// wherever the loader happened to find the same file. +const spmAssetPinKey = "BAAI/bge-m3/sentencepiece.bpe.model" + +// readTokenizerAsset finds the first readable candidate and verifies its digest +// when one is pinned. It mirrors bpe_loader.go: disk only, no network I/O, and a +// corrupt file fails loudly instead of being skipped. +func readTokenizerAsset(names []string, pinKey string, pins map[string]string) (string, []byte, error) { + var tried []string + // MODEL_ASSETS_DIR wins over wherever the process happens to be running: it is how + // a deployment points at a mounted asset tree (see common.ModelAssetCandidates). + for _, name := range names { + tried = append(tried, common.ModelAssetCandidates(name)...) + } + for _, root := range searchRoots() { + for _, name := range names { + tried = append(tried, filepath.Join(root, name)) + } + } + for _, candidate := range tried { + raw, err := os.ReadFile(candidate) + if err != nil { + if os.IsNotExist(err) { + continue + } + return "", nil, fmt.Errorf("reading tokenizer asset %s: %w", candidate, err) + } + if want, ok := pins[pinKey]; ok { + if got := sha1Hex(raw); got != want { + return "", nil, fmt.Errorf("tokenizer asset %s digest mismatch (got %s, want %s); refusing to load a corrupt or tampered file", candidate, got, want) + } + } + return candidate, raw, nil + } + return "", nil, fmt.Errorf("no tokenizer asset found; run `uv run ragflow_deps/download_go_deps.py` or set %s to a directory holding them; tried: %s", + common.EnvModelAssetsDir, strings.Join(tried, ", ")) +} + +func sha1Hex(raw []byte) string { + sum := sha1.Sum(raw) + return fmt.Sprintf("%x", sum) +} + +func init() { + RegisterCounterLoader(CounterXLMRSentence, LoadXLMRSentencePieceCounter) +} + +// --------------------------------------------------------------------------- +// Minimal protobuf reader (sentencepiece model.proto, wire format only) +// --------------------------------------------------------------------------- + +const ( + pbWireVarint = 0 + pbWireFixed64 = 1 + pbWireBytes = 2 + pbWireFixed32 = 5 +) + +var errPBEOF = errors.New("protobuf: unexpected end of input") + +type pbReader struct { + b []byte + i int +} + +func (r *pbReader) readKey() (int, int, error) { + v, err := r.readVarint() + if err != nil { + return 0, 0, err + } + return int(v >> 3), int(v & 7), nil +} + +func (r *pbReader) readVarint() (uint64, error) { + var ( + value uint64 + shift uint + ) + for { + if r.i >= len(r.b) { + return 0, errPBEOF + } + c := r.b[r.i] + r.i++ + value |= uint64(c&0x7F) << shift + if c&0x80 == 0 { + return value, nil + } + shift += 7 + if shift > 63 { + return 0, errors.New("protobuf: varint overflow") + } + } +} + +func (r *pbReader) readBytes() ([]byte, error) { + length, err := r.readVarint() + if err != nil { + return nil, err + } + if uint64(r.i)+length > uint64(len(r.b)) { + return nil, errPBEOF + } + out := r.b[r.i : r.i+int(length)] + r.i += int(length) + return out, nil +} + +func (r *pbReader) readFixed32() (uint32, error) { + if r.i+4 > len(r.b) { + return 0, errPBEOF + } + v := binary.LittleEndian.Uint32(r.b[r.i : r.i+4]) + r.i += 4 + return v, nil +} + +func (r *pbReader) skip(wire int) error { + switch wire { + case pbWireVarint: + _, err := r.readVarint() + return err + case pbWireFixed64: + if r.i+8 > len(r.b) { + return errPBEOF + } + r.i += 8 + return nil + case pbWireBytes: + _, err := r.readBytes() + return err + case pbWireFixed32: + _, err := r.readFixed32() + return err + default: + return fmt.Errorf("protobuf: unsupported wire type %d", wire) + } +} diff --git a/internal/tokenizer/tokenizer.go b/internal/tokenizer/tokenizer.go index 095bdca655..98446812ae 100644 --- a/internal/tokenizer/tokenizer.go +++ b/internal/tokenizer/tokenizer.go @@ -24,6 +24,7 @@ import ( "sync" "sync/atomic" "time" + "unicode/utf8" "github.com/pkoukk/tiktoken-go" "go.uber.org/zap" @@ -627,6 +628,13 @@ func NumTokensFromString(s string) int { // Like Python's trim_content, an unavailable encoder PANICS: the previous // byte-length fallback silently produced a differently-truncated string instead // of surfacing the missing table. See NumTokensFromString. +// The decoded prefix is snapped back to a rune boundary before it is returned. +// Cutting the token list is not enough on its own: a single emoji is four UTF-8 +// bytes and may span tokens, so a token-boundary cut can end mid-character and +// hand the caller - and the index - invalid UTF-8. (TestCountersSatisfyTrimProperties +// caught exactly that on the emoji corpus at limit=1.) Dropping the partial rune +// can only shrink the text, but the token count is re-checked anyway so the +// "at most limit tokens" contract holds unconditionally. func TrimContentToTokenLimit(s string, limit int) string { if limit < 0 { limit = 0 @@ -642,7 +650,35 @@ func TrimContentToTokenLimit(s string, limit int) string { if len(tokens) <= limit { return s } - return enc.Decode(tokens[:limit]) + trimmed := clampToValidUTF8(enc.Decode(tokens[:limit])) + for i := 0; i < utf8.UTFMax && len(trimmed) > 0 && len(enc.Encode(trimmed, nil, nil)) > limit; i++ { + trimmed = dropLastRune(trimmed) + } + return trimmed +} + +// clampToValidUTF8 returns the longest prefix of s that is valid UTF-8. Only the +// tail can be broken after a token-boundary cut, so this is at most three bytes +// of backtracking rather than a scan for an interior invalid byte. +func clampToValidUTF8(s string) string { + if utf8.ValidString(s) { + return s + } + for cut := 1; cut <= utf8.UTFMax && cut <= len(s); cut++ { + if prefix := s[:len(s)-cut]; utf8.ValidString(prefix) { + return prefix + } + } + return "" +} + +// dropLastRune removes the final rune of s, tolerating a trailing partial rune. +func dropLastRune(s string) string { + _, size := utf8.DecodeLastRuneInString(s) + if size <= 0 { + return "" + } + return s[:len(s)-size] } // EncodeCL100KTokens returns the cl100k_base token ids of s, from the same diff --git a/internal/tokenizer/usage.go b/internal/tokenizer/usage.go index ec41d973b7..cfd89632cf 100644 --- a/internal/tokenizer/usage.go +++ b/internal/tokenizer/usage.go @@ -34,8 +34,28 @@ import ( "context" "encoding/json" "sync" + "time" ) +// CallUsage is ONE LLM call's token split inside a run. The aggregate (RunUsage) +// answers "what did this question cost"; this answers "where did the cost go" — +// which ReAct iteration, repair turn or auditor pass burned the tokens, and how +// the input/output mix moved as the context grew. +type CallUsage struct { + // Seq is the 1-based call index within the run. + Seq int `json:"seq"` + // AtSeconds is the call's start relative to the run sink. Rounded to + // milliseconds: enough to order calls, not enough to bloat the artefact. + AtSeconds float64 `json:"at_seconds"` + // Model is the model the call went to, when the caller knows it (the + // explorer, the auditor and the synthesis fallback can use different ones). + Model string `json:"model,omitempty"` + // PromptTokens / CompletionTokens are the call's input and output split. + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + // Context key types — unexported to prevent direct external access. type runUsageKeyType struct{} type runAttrsKeyType struct{} @@ -51,11 +71,22 @@ type RunUsage struct { CompletionTokens int TotalTokens int Calls int + // startedAt anchors CallUsage.AtSeconds; calls keeps one record per LLM + // call so a per-question cost can be broken down turn by turn. + startedAt time.Time + calls []CallUsage } // Add atomically adds a single LLM call's token counts to the sink. // Safe to call concurrently from multiple goroutines. func (u *RunUsage) Add(prompt, completion, total int) { + u.AddFor("", prompt, completion, total) +} + +// AddFor is Add plus the model that served the call, and records the call's own +// token split. Callers that know the model (the chat-model wrapper does) get a +// per-turn breakdown for free; the rest still get a record with an empty model. +func (u *RunUsage) AddFor(model string, prompt, completion, total int) { if u == nil { return } @@ -71,6 +102,33 @@ func (u *RunUsage) Add(prompt, completion, total int) { u.TotalTokens += total } u.Calls++ + at := 0.0 + if !u.startedAt.IsZero() { + at = time.Since(u.startedAt).Seconds() + } + u.calls = append(u.calls, CallUsage{ + Seq: u.Calls, + AtSeconds: float64(int(at*1000)) / 1000, + Model: model, + PromptTokens: prompt, + CompletionTokens: completion, + TotalTokens: total, + }) +} + +// CallSnapshot returns a copy of the per-call records, oldest first. +func (u *RunUsage) CallSnapshot() []CallUsage { + if u == nil { + return nil + } + u.mu.Lock() + defer u.mu.Unlock() + if len(u.calls) == 0 { + return nil + } + out := make([]CallUsage, len(u.calls)) + copy(out, u.calls) + return out } // Snapshot returns a copy of the current cumulative counts. @@ -93,7 +151,7 @@ type RunAttrs struct { // WithRunUsage installs a fresh RunUsage sink on ctx. Should be called // once at the start of a canvas turn. func WithRunUsage(ctx context.Context) context.Context { - return context.WithValue(ctx, runUsageKeyType{}, &RunUsage{}) + return context.WithValue(ctx, runUsageKeyType{}, &RunUsage{startedAt: time.Now()}) } // GetRunUsage retrieves the per-run token usage sink from ctx. @@ -129,11 +187,18 @@ func GetRunAttrs(ctx context.Context) *RunAttrs { // active run sink on ctx. Safe to call from anywhere; when no run sink // is installed it is a no-op. func RecordRunTokenUsage(ctx context.Context, promptTokens, completionTokens, totalTokens int) { + RecordRunTokenUsageFor(ctx, "", promptTokens, completionTokens, totalTokens) +} + +// RecordRunTokenUsageFor is RecordRunTokenUsage plus the model that served the +// call, so the per-call breakdown can attribute tokens to the explorer, the +// auditor or the synthesis fallback. +func RecordRunTokenUsageFor(ctx context.Context, model string, promptTokens, completionTokens, totalTokens int) { sink := GetRunUsage(ctx) if sink == nil { return } - sink.Add(promptTokens, completionTokens, totalTokens) + sink.AddFor(model, promptTokens, completionTokens, totalTokens) } // UsageFromMap extracts a token usage split from a raw API response map. diff --git a/internal/tokenizer/usage_test.go b/internal/tokenizer/usage_test.go new file mode 100644 index 0000000000..cc9e6d2414 --- /dev/null +++ b/internal/tokenizer/usage_test.go @@ -0,0 +1,54 @@ +// +// 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 ( + "context" + "testing" + "time" +) + +// TestRunUsageCallSnapshot pins the per-LLM-call breakdown: the aggregate and +// the per-call records must agree, and each record must carry its own input / +// output split, the model that served it, and a monotonic start offset. +func TestRunUsageCallSnapshot(t *testing.T) { + ctx := WithRunUsage(context.Background()) + RecordRunTokenUsageFor(ctx, "MiniMax-M3", 1200, 80, 1280) + time.Sleep(2 * time.Millisecond) + RecordRunTokenUsage(ctx, 3000, 150, 3150) // no model known: still recorded + + sink := GetRunUsage(ctx) + pt, ct, tt, calls := sink.Snapshot() + if pt != 4200 || ct != 230 || tt != 4430 || calls != 2 { + t.Fatalf("aggregate = %d/%d/%d calls=%d", pt, ct, tt, calls) + } + turns := sink.CallSnapshot() + if len(turns) != 2 { + t.Fatalf("expected 2 per-call records, got %d", len(turns)) + } + if turns[0].Seq != 1 || turns[0].PromptTokens != 1200 || turns[0].CompletionTokens != 80 || turns[0].Model != "MiniMax-M3" { + t.Fatalf("first call wrong: %+v", turns[0]) + } + if turns[1].Seq != 2 || turns[1].Model != "" || turns[1].TotalTokens != 3150 { + t.Fatalf("second call wrong: %+v", turns[1]) + } + if turns[1].AtSeconds < turns[0].AtSeconds { + t.Fatalf("call offsets must be non-decreasing: %v then %v", turns[0].AtSeconds, turns[1].AtSeconds) + } + // Without a sink installed, recording is a no-op (not a panic). + RecordRunTokenUsageFor(context.Background(), "m", 1, 1, 2) +} diff --git a/internal/tokenizer/wordpiece.go b/internal/tokenizer/wordpiece.go new file mode 100644 index 0000000000..d645a1bbbc --- /dev/null +++ b/internal/tokenizer/wordpiece.go @@ -0,0 +1,436 @@ +// +// 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 + +// BERT WordPiece counter: the bge-*-en / e5-* / gte-base / jina-v2 family. +// +// These models are BERT-derived, so they tokenize with a 30k WordPiece vocab +// (`vocab.txt`) plus BERT's own text processing. Everything here mirrors the +// served tokenizer - the HuggingFace conversion of the same vocab, verified by +// scripts/gen_tokenizer_oracle.py: +// +// normalizer: BertNormalizer(clean_text=true, handle_chinese_chars=true, +// strip_accents=null (follows lowercase=true), lowercase=true) +// pre_tokenizer: BertPreTokenizer (whitespace, punctuation and CJK chars split) +// model: WordPiece(continuing_subword_prefix="##", max_input_chars_per_word=100) +// +// Note the asymmetry with the XLM-R counter: this family lowercases and strips +// accents, XLM-R does neither. Counting "HELLO" as one token for both would be +// wrong for one of them. + +import ( + "errors" + "fmt" + "path/filepath" + "strings" + "sync" + "unicode" + + "golang.org/x/text/unicode/norm" +) + +// wordPieceMaxCharsPerWord mirrors the model's max_input_chars_per_word: longer +// words become a single unk token instead of being split. +const wordPieceMaxCharsPerWord = 100 + +// wordPieceModel is a WordPiece vocabulary plus BERT's text processing flags. +type wordPieceModel struct { + vocab map[string]int32 + unkToken string + continuing string + lowercase bool + stripAccents bool + cleanText bool + handleChinese bool + maxCharsPerWord int + + // names is the id -> piece map, built on first use: only the oracle test asks + // for token texts, and production would not want a second copy of the vocab. + nameOnce sync.Once + names map[int32]string +} + +// bgLargeEnVocabPin is the digest of the shipped vocab.txt. +var expectedWordPieceHashes = map[string]string{ + "BAAI/bge-large-en-v1.5/vocab.txt": "c3b4105373feaa5b0b2c332321c1e592d6f39658", +} + +const wordPieceAssetPinKey = "BAAI/bge-large-en-v1.5/vocab.txt" + +var bertWordPieceAssetNames = []string{ + "ragflow_deps/huggingface.co/BAAI/bge-large-en-v1.5/vocab.txt", + "vocab.txt", + "bert_wordpiece_vocab.txt", +} + +// parseWordPieceVocab reads one token per line, index order. +func parseWordPieceVocab(raw []byte) (map[string]int32, error) { + lines := strings.Split(strings.ReplaceAll(string(raw), "\r\n", "\n"), "\n") + vocab := make(map[string]int32, len(lines)) + for i, line := range lines { + if line == "" && i == len(lines)-1 { + continue + } + token := strings.TrimSuffix(line, "\r") + if _, dup := vocab[token]; dup { + continue + } + vocab[token] = int32(i) + } + if len(vocab) == 0 { + return nil, errors.New("wordpiece: empty vocabulary") + } + return vocab, nil +} + +// normalize applies BertNormalizer's steps in the order HF applies them. +func (m *wordPieceModel) normalize(text string) string { + s := text + if m.cleanText { + s = cleanBertText(s) + } + if m.handleChinese { + s = padChineseChars(s) + } + if m.lowercase { + s = strings.ToLower(s) + } + if m.stripAccents { + s = stripAccents(s) + } + return s +} + +// cleanBertText removes control characters and normalizes whitespace, mirroring +// BertNormalizer.clean_text. +// +// The distinction that matters: whitespace is *replaced* by a space (so it still +// separates words), while control characters are *deleted*. Format characters +// (category Cf: zero-width space/joiner, word joiner, BOM) are deleted too - the +// oracle shows "a\u200bb" tokenizing as a single "ab". Treating either class as a +// space splits words the model keeps together, and leaves the counter over an +// entire token per occurrence. +func cleanBertText(s string) string { + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + switch { + case r == 0 || r == 0xFFFD: + continue + case unicode.IsSpace(r): // includes \t \n \r and NBSP-like spaces + b.WriteRune(' ') + case unicode.IsControl(r) || unicode.In(r, unicode.Cf): + continue + default: + b.WriteRune(r) + } + } + return b.String() +} + +// padChineseChars puts a space around every CJK character, mirroring +// BertNormalizer.handle_chinese_chars. +func padChineseChars(s string) string { + var b strings.Builder + b.Grow(len(s) + 8) + prevSpace := true + for _, r := range s { + if isChineseChar(r) { + if !prevSpace { + b.WriteRune(' ') + } + b.WriteRune(r) + b.WriteRune(' ') + prevSpace = true + continue + } + if r == ' ' { + prevSpace = true + } else { + prevSpace = false + } + b.WriteRune(r) + } + return b.String() +} + +// isChineseChar reports whether r is in a CJK block, using the same ranges BERT +// uses (CJK Unified Ideographs and extensions, and the CJK symbols block). +func isChineseChar(r rune) bool { + return (r >= 0x4E00 && r <= 0x9FFF) || + (r >= 0x3400 && r <= 0x4DBF) || + (r >= 0x20000 && r <= 0x2A6DF) || + (r >= 0x2A700 && r <= 0x2B73F) || + (r >= 0x2B740 && r <= 0x2B81F) || + (r >= 0x2B820 && r <= 0x2CEAF) || + (r >= 0xF900 && r <= 0xFAFF) || + (r >= 0x2F800 && r <= 0x2FA1F) +} + +// stripAccents removes combining marks, mirroring BertNormalizer.strip_accents. +// +// NFD first, then drop the marks. (An earlier version short-circuited when the +// input was already NFD, which is exactly the case that still has marks to +// remove: "e\u0301" is NFD, and the counter kept the accent as an extra token.) +func stripAccents(s string) string { + var b strings.Builder + b.Grow(len(s)) + for _, r := range norm.NFD.String(s) { + if unicode.Is(unicode.Mn, r) { + continue + } + b.WriteRune(r) + } + return b.String() +} + +// preTokenize splits text the way BertPreTokenizer does: on whitespace, around +// punctuation, and around CJK characters. +func preTokenizeBert(s string) []string { + var ( + out []string + cur strings.Builder + flush = func() { + if cur.Len() > 0 { + out = append(out, cur.String()) + cur.Reset() + } + } + ) + for _, r := range s { + switch { + case unicode.IsSpace(r): + flush() + case isPunctuation(r): + flush() + out = append(out, string(r)) + case isChineseChar(r): + flush() + out = append(out, string(r)) + default: + cur.WriteRune(r) + } + } + flush() + return out +} + +// isPunctuation mirrors BERT's _is_punctuation: the ASCII ranges between letters +// and digits, plus Unicode category **P** — and only P. Symbols (category S, which +// includes every emoji) are NOT punctuation for BERT, so "🚀🔥" stays inside a +// word instead of becoming one token per character. Verified by the oracle: the +// emoji sample is what caught an IsSymbol() added here by mistake. +func isPunctuation(r rune) bool { + if (r >= 33 && r <= 47) || (r >= 58 && r <= 64) || (r >= 91 && r <= 96) || (r >= 123 && r <= 126) { + return true + } + return unicode.IsPunct(r) +} + +// countTokens is the WordPiece greedy longest-match count, including the single +// unk token a word collapses to when nothing matches. +func (m *wordPieceModel) countTokens(text string) int { + n, _ := m.encode(text, false) + return n +} + +// tokenIDs returns the exact token ids, for the oracle test's id-level comparison +// (two segmentations can share a count; only the ids prove they are the same). +func (m *wordPieceModel) tokenIDs(text string) []int32 { + _, ids := m.encode(text, true) + return ids +} + +// tokenNames returns the piece of every token ("##"-prefixed for continuations, +// the model's unk token for unknowns), which is what the served tokenizer reports. +func (m *wordPieceModel) tokenNames(text string) []string { + n, ids := m.encode(text, true) + if n == 0 { + return nil + } + m.nameOnce.Do(func() { + m.names = make(map[int32]string, len(m.vocab)) + for piece, id := range m.vocab { + m.names[id] = piece + } + }) + names := make([]string, 0, len(ids)) + for _, id := range ids { + names = append(names, m.names[id]) + } + return names +} + +// encode runs BERT's greedy longest-match word splitting, returning the token count +// and - when wantIDs is set - the token ids. Both live in one function so the count +// path and the id path cannot drift apart. +func (m *wordPieceModel) encode(text string, wantIDs bool) (int, []int32) { + s := m.normalize(text) + if strings.TrimSpace(s) == "" { + return 0, nil + } + total := 0 + var ids []int32 + unkID := m.vocab[m.unkToken] + appendID := func(id int32) { + total++ + if wantIDs { + ids = append(ids, id) + } + } + // One word at a time. If ANY position of a word cannot be matched, the whole + // word collapses to a single [UNK] and the pieces already matched for it are + // dropped - that is what BERT's reference implementation does + // (`if is_bad: split_tokens.append(self.unk_token)`), and what the HuggingFace + // WordPiece model does (`is_bad` -> one unk). Emitting "matched prefix + one + // unk for the rest" instead over-counts such words, which is what the + // id-level oracle found on the fuzz corpus. + var pieces []int32 + for _, word := range preTokenizeBert(s) { + if word == "" { + continue + } + if len([]rune(word)) > m.maxCharsPerWord { + appendID(unkID) // the whole word becomes unk + continue + } + pieces = pieces[:0] + start := 0 + runes := []rune(word) + bad := false + for start < len(runes) { + end := len(runes) + matched := false + for end > start { + piece := string(runes[start:end]) + if start > 0 { + piece = m.continuing + piece + } + if id, ok := m.vocab[piece]; ok { + pieces = append(pieces, id) + start = end + matched = true + break + } + end-- + } + if !matched { + bad = true + break + } + } + if bad { + appendID(unkID) + continue + } + for _, id := range pieces { + appendID(id) + } + } + return total, ids +} + +// trimToLimit returns the longest rune prefix with at most limit tokens. +func (m *wordPieceModel) trimToLimit(text string, limit int) string { + if limit <= 0 { + return "" + } + if m.countTokens(text) <= limit { + return text + } + runes := []rune(text) + lo, hi := 0, len(runes) + for lo < hi { + mid := (lo + hi + 1) / 2 + if m.countTokens(string(runes[:mid])) <= limit { + lo = mid + } else { + hi = mid - 1 + } + } + return string(runes[:lo]) +} + +type wordPieceCounter struct { + model *wordPieceModel + path string +} + +func (c *wordPieceCounter) ID() string { return CounterBERTWordPiece } + +// IDs reports the token ids, for the oracle test's id-level comparison. +func (c *wordPieceCounter) IDs(text string) []int32 { return c.model.tokenIDs(text) } + +// Names reports the piece of every token, which the oracle test compares directly +// with the served tokenizer's tokens. +func (c *wordPieceCounter) Names(text string) []string { return c.model.tokenNames(text) } + +// SourcePath is the vocab file this counter loaded. +func (c *wordPieceCounter) SourcePath() string { return c.path } + +func (c *wordPieceCounter) Count(text string) int { return c.model.countTokens(text) } + +func (c *wordPieceCounter) TrimToLimit(text string, limit int) string { + return c.model.trimToLimit(text, limit) +} + +func (c *wordPieceCounter) Available() bool { return c.model != nil } + +var wordPieceOnce struct { + sync.Once + counter Counter + err error +} + +// LoadBERTWordPieceCounter loads (once) the BERT WordPiece counter from the +// shipped vocab.txt. A missing asset is not fatal: callers fall back to a +// calibrated cl100k count. +func LoadBERTWordPieceCounter() (Counter, error) { + wordPieceOnce.Do(func() { + wordPieceOnce.counter, wordPieceOnce.err = loadBERTWordPieceCounter() + }) + return wordPieceOnce.counter, wordPieceOnce.err +} + +func loadBERTWordPieceCounter() (Counter, error) { + path, raw, err := readTokenizerAsset(bertWordPieceAssetNames, wordPieceAssetPinKey, expectedWordPieceHashes) + if err != nil { + return nil, err + } + vocab, err := parseWordPieceVocab(raw) + if err != nil { + return nil, fmt.Errorf("wordpiece asset %s: %w", path, err) + } + model := &wordPieceModel{ + vocab: vocab, + unkToken: "[UNK]", + continuing: "##", + lowercase: true, + stripAccents: true, + cleanText: true, + handleChinese: true, + maxCharsPerWord: wordPieceMaxCharsPerWord, + } + if _, ok := vocab[model.unkToken]; !ok { + return nil, fmt.Errorf("wordpiece asset %s has no %s token", filepath.Base(path), model.unkToken) + } + return &wordPieceCounter{model: model, path: path}, nil +} + +func init() { + RegisterCounterLoader(CounterBERTWordPiece, LoadBERTWordPieceCounter) +} diff --git a/ragflow_deps/download_deps.py b/ragflow_deps/download_deps.py old mode 100644 new mode 100755 diff --git a/ragflow_deps/download_go_deps.py b/ragflow_deps/download_go_deps.py old mode 100644 new mode 100755 index 92a56f07eb..d0dbd46d67 --- a/ragflow_deps/download_go_deps.py +++ b/ragflow_deps/download_go_deps.py @@ -4,7 +4,6 @@ # /// script # requires-python = ">=3.10" # dependencies = [ -# "nltk", # "huggingface-hub", # "requests", # ] @@ -43,6 +42,14 @@ # to be set: # # bash build.sh --test-native # or: cd internal/deepdoc/native && bash run.sh +# +# Embedding tokenizer assets: the Go token counters in internal/tokenizer load a +# tokenizer file per family (XLM-R SentencePiece, BERT WordPiece, two byte-level BPE +# families) from `huggingface.co//` at the top of ragflow_deps/ - the path +# the loaders search and the one ragflow_deps/Dockerfile ships. This script fetches +# them too, so one invocation gives a Go checkout everything it counts with. Without +# them every model that declares a tokenizer silently counts with the calibrated +# cl100k estimate instead (see internal/tokenizer/embedding_token_limits.md). import argparse import os @@ -90,6 +97,36 @@ def _ort_normalized_dir(version): DEEPDOC_REPO = "InfiniFlow/deepdoc" DEEPDOC_MODEL_FILES = ["det.ort", "layout.ort", "tsr.ort", "rec.ort", "ocr.res"] +# Embedding tokenizer assets for the Go counters (internal/tokenizer). Each entry is +# (repo, file, kind): +# +# "runtime" - the counters load it in production, so the runtime images must ship it +# (the copy loop in Dockerfile / Dockerfile_base) and the loader must pin +# its SHA-1; +# "oracle" - only scripts/gen_tokenizer_oracle.py needs it, to regenerate the test +# fixtures, so it is deliberately NOT shipped to the image. +# +# ragflow_deps/test_tokenizer_assets.py holds those three places together, because a +# drift here is invisible at runtime: the counter simply becomes unavailable and ingest +# counts with the calibrated fallback. +# +# Fetched per file, not by snapshot: these repos also carry multi-GB weights we do not +# want. A missing asset is NOT fatal (unlike the DeepDoc weights above, without which +# the server cannot run at all). +TOKENIZER_ASSETS = [ + # XLM-R SentencePiece (Unigram) - the BAAI bge / multilingual-e5 / m3e / + # gte-multilingual / jina-v3 family shares this vocabulary. + ("BAAI/bge-m3", "sentencepiece.bpe.model", "runtime"), + # BERT WordPiece vocab for the bge-*-en / e5-* / gte-base family. + ("BAAI/bge-large-en-v1.5", "vocab.txt", "runtime"), + # Byte-level BPE families (Qwen3-Embedding, Mistral/Llama embeddings). + ("Qwen/Qwen3-Embedding-0.6B", "tokenizer.json", "runtime"), + ("intfloat/e5-mistral-7b-instruct", "tokenizer.json", "runtime"), + # Cross-check oracles only: the HF tokenizer.json files the same families convert to. + ("BAAI/bge-m3", "tokenizer.json", "oracle"), + ("BAAI/bge-large-en-v1.5", "tokenizer.json", "oracle"), +] + def prune_stale_onnxruntime(static_lib_dir, version): """Remove ONNX Runtime version dirs under static_lib that do NOT match @@ -293,6 +330,62 @@ def download_go_models(use_china_mirrors=False): print(" No MODEL_DIR env needed: the Go backend auto-discovers this directory.") +def download_tokenizer_assets(use_china_mirrors=False): + """Download the embedding tokenizer assets the Go counters load. + + Written to `ragflow_deps/huggingface.co//`: the path internal/tokenizer + searches and the one ragflow_deps/Dockerfile builds its image from. Fetched per file + (these repos also carry multi-GB weights we do not want), routed via hf-mirror.com + when --china-mirrors is set. + + A missing asset is fatal, like every other asset this script provisions. The counters + would otherwise fall back to the calibrated cl100k estimate, which under-counts some + tokenizers - and an under-count is what makes a provider answer 400. The runtime + images refuse to build on exactly the same condition (the Dockerfile copy loops exit + 1), so provisioning has to fail here rather than hand them a broken tree. + """ + if use_china_mirrors: + os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + # Imported lazily so the module stays importable (and its unit tests stay + # huggingface-free) without the huggingface_hub dependency installed. + from huggingface_hub import hf_hub_download + + base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "huggingface.co") + missing = [] + for repo_id, filename, _kind in TOKENIZER_ASSETS: + target_dir = os.path.join(base, repo_id) + if os.path.isfile(os.path.join(target_dir, filename)): + print(f" ✓ {repo_id}/{filename} already present") + continue + os.makedirs(target_dir, exist_ok=True) + print(f"Downloading tokenizer asset {repo_id}/{filename}...") + try: + hf_hub_download(repo_id=repo_id, filename=filename, local_dir=target_dir) + except Exception as e: # noqa: BLE001 - collected and surfaced below + missing.append((repo_id, filename, e)) + + if missing: + for repo_id, filename, e in missing: + print(f" ERROR: could not fetch {repo_id}/{filename}: {e}", file=sys.stderr) + print( + "\n" + "The embedding tokenizer counters in internal/tokenizer load these files from\n" + f" {base}\n" + "Without one, the models that declare its counter fall back to the calibrated\n" + "cl100k estimate: safe, but less precise - it under-counts some tokenizers, and\n" + "an under-count is a request the provider rejects. The runtime images refuse to\n" + "build on the same condition (their copy loops exit 1), so this script fails\n" + "instead of handing a broken tree to the image build. To recover:\n" + " - re-run this script (a transient HF/network error usually clears);\n" + " - behind the GFW, re-run with --china-mirrors (routes via hf-mirror.com);\n" + " - or point MODEL_ASSETS_DIR at a directory holding the same layout.\n", + file=sys.stderr, + ) + sys.exit(1) + + print(" ✓ embedding tokenizer assets ready") + + if __name__ == "__main__": # Anchor CWD to this file's directory so all relative outputs # (huggingface.co/, nltk_data/, *.deb, *.jar, *.tar.gz, etc.) land @@ -377,3 +470,7 @@ if __name__ == "__main__": # Download the Go DeepDoc `.ort` weights so this script is a one-stop for Go # dev: native libs (above) + model files (below) from a single invocation. download_go_models(args.china_mirrors) + + # Same one-stop idea for the token counters: their tokenizer files, in the layout + # the loaders search and ragflow_deps/Dockerfile ships. + download_tokenizer_assets(args.china_mirrors) diff --git a/ragflow_deps/test_tokenizer_assets.py b/ragflow_deps/test_tokenizer_assets.py new file mode 100644 index 0000000000..94d5467ae5 --- /dev/null +++ b/ragflow_deps/test_tokenizer_assets.py @@ -0,0 +1,103 @@ +# Guard: the tokenizer assets production loads must actually be shipped and pinned. +# +# The Go counters in internal/tokenizer read their vocabulary from +# ragflow_deps/huggingface.co//. Three places have to agree: +# +# 1. ragflow_deps/download_go_deps.py - fetches the file, and says whether it is +# "runtime" (the counters load it) or "oracle" +# (only scripts/gen_tokenizer_oracle.py does). +# It is the Go-side downloader: everything the +# Go side reads comes from here, and +# ragflow_deps/download_deps.py (upstream) is +# deliberately left alone; +# 2. Dockerfile / Dockerfile_base / Dockerfile_go +# - copy the runtime ones into the image at the +# path the counters search; +# 3. internal/tokenizer/*.go - pin the SHA-1 of the runtime ones. +# +# When they drift, the layer that can see it fails loudly: the runtime image does not +# build without a runtime asset, and an embedder whose model declares a counter that +# cannot load refuses to ingest rather than substituting the calibrated cl100k count +# (cl100k under-counts XLM-R on some content, and an under-count is what makes a provider +# answer 400). Those failures are still late - a broken image gets built, or the first +# document fails - which is why the drift is checked here, cheaply, instead. +# See internal/tokenizer/embedding_token_limits.md. + +import ast +import re +from pathlib import Path + +HERE = Path(__file__).resolve().parent +ROOT = HERE.parent +RUNTIME_IMAGES = ("Dockerfile", "Dockerfile_base", "Dockerfile_go") +GO_SOURCES = ("spm.go", "wordpiece.go", "bpe.go") +KINDS = ("runtime", "oracle") +DOWNLOADER = "download_go_deps.py" +UPSTREAM_DOWNLOADER = "download_deps.py" + + +def load_tokenizer_assets(): + """Read TOKENIZER_ASSETS out of the Go-side downloader without importing it (that + module pulls in huggingface_hub, which a unit test run does not need).""" + source = (HERE / DOWNLOADER).read_text(encoding="utf-8") + for node in ast.parse(source).body: + if isinstance(node, ast.Assign) and any(getattr(target, "id", None) == "TOKENIZER_ASSETS" for target in node.targets): + return [tuple(ast.literal_eval(element)) for element in node.value.elts] + raise AssertionError(f"TOKENIZER_ASSETS not found in {DOWNLOADER}") + + +def image_text(name): + return (ROOT / name).read_text(encoding="utf-8") + + +def test_assets_are_well_formed(): + assets = load_tokenizer_assets() + assert assets, "tokenizer_assets is empty" + for repo, filename, kind in assets: + assert kind in KINDS, f"{repo}/{filename}: unknown kind {kind!r}" + assert "/" in repo and filename, f"{repo}/{filename}: malformed entry" + runtime = [a for a in assets if a[2] == "runtime"] + assert runtime, "no runtime asset: at least one counter would always be unavailable" + + +def test_every_runtime_asset_is_copied_into_every_runtime_image(): + assets = [a for a in load_tokenizer_assets() if a[2] == "runtime"] + for image in RUNTIME_IMAGES: + text = image_text(image) + for repo, filename, _kind in assets: + assert f"{repo}/{filename}" in text, ( + f"{image} does not copy {repo}/{filename}; the matching counter would be unavailable in that image and ingest would silently fall back to the calibrated count" + ) + + +def test_every_runtime_asset_is_pinned_by_the_loader(): + sources = {name: (ROOT / "internal" / "tokenizer" / name).read_text(encoding="utf-8") for name in GO_SOURCES} + for repo, filename, kind in load_tokenizer_assets(): + if kind != "runtime": + continue + key = f"{repo}/{filename}" + pinned = [match.group(1) for text in sources.values() for match in re.finditer(rf'"{re.escape(key)}":\s*"([0-9a-f]{{40}})"', text)] + assert pinned, f"{key} has no 40-hex SHA-1 pin in internal/tokenizer/*.go" + for digest in pinned: + assert len(set(digest)) > 4, f"{key}: pin {digest!r} does not look like a digest" + + +def test_oracle_only_assets_are_not_shipped(): + assets = [a for a in load_tokenizer_assets() if a[2] == "oracle"] + for image in RUNTIME_IMAGES: + text = image_text(image) + for repo, filename, _kind in assets: + assert f"{repo}/{filename}" not in text, f"{image} ships {repo}/{filename}, which only the fixture generator needs" + + +def test_assets_are_fetched_by_the_go_downloader_only(): + """The Go-side assets belong to the Go-side downloader. + + ragflow_deps/download_deps.py is an upstream file (it snapshots the Python side's + dependencies); adding the Go counters' tokenizer files there split the list across + two scripts and made the upstream file carry Go concerns. Keeping the fetch in + download_go_deps.py is what makes a Go checkout self-sufficient. + """ + upstream = (HERE / UPSTREAM_DOWNLOADER).read_text(encoding="utf-8") + for repo, filename, _kind in load_tokenizer_assets(): + assert filename not in upstream, f"{UPSTREAM_DOWNLOADER} mentions {repo}/{filename}; the embedding tokenizer assets are fetched by {DOWNLOADER} only" diff --git a/scripts/gen_tokenizer_oracle.py b/scripts/gen_tokenizer_oracle.py new file mode 100755 index 0000000000..564551a0c9 --- /dev/null +++ b/scripts/gen_tokenizer_oracle.py @@ -0,0 +1,414 @@ +#!/usr/bin/env python3 + +# PEP 723 metadata +# /// script +# requires-python = ">=3.13" +# dependencies = [ +# "tokenizers", +# "tiktoken", +# "sentencepiece", +# ] +# /// +"""Generate tokenizer oracle fixtures for internal/tokenizer's counters. + +The Go counters are hand-written reimplementations of four tokenizer +architectures (SentencePiece Unigram, BERT WordPiece, byte-level BPE, +SentencePiece-BPE) plus the tiktoken table. "Looks right" is not evidence, so each +one is checked against the implementation that actually defines the answer: + + * cl100k - OpenAI's own tiktoken package on the shipped table; + * xlmr-spm - the HuggingFace `tokenizers` conversion of the model that + SILICONFLOW serves (bge-m3's tokenizer.json), with the + canonical Google `sentencepiece` package cross-checked under + --spm-crosscheck; + * bert-wordpiece - the model's own tokenizer.json (HF `tokenizers`); + * qwen-bpe / + llama-bpe - the model's own tokenizer.json (HF `tokenizers`). + +The fixture stores the expected token **ids**, not just their number. Two different +segmentations can share a token count, so a count comparison cannot tell "same +tokens" from "different tokens, same length"; the Go test compares ids element by +element, and a mismatch is reported with the first divergence and its context. + +Usage: + for m in bge-m3 bge-large-en qwen3-embedding e5-mistral cl100k; do + uv run scripts/gen_tokenizer_oracle.py --model $m --corpus --out /tmp/tokenizer_oracle/$m.json + done + TOKENIZER_ORACLE_DIR=/tmp/tokenizer_oracle \\ + ./build.sh --test -run TestCountersMatchOracle ./internal/tokenizer/ + +`--corpus` is REQUIRED, and its argument is a path, never a built-in default: the +document belongs to the caller, not to this script. Its ENTIRE content becomes one +regression sample (`incident_full`), the document that exposed the under-count. A +window of that document is deliberately NOT sampled: the window would silently stop +covering the document the moment the document changes. + +Every dependency is declared inline above (PEP 723), so `uv run` provisions them: there +is no `pip install` step, and no import-time fallback - a dependency that is missing +fails the run instead of degrading it silently. + +The same holds for the assets it reads: the model's `tokenizer.json`, the cl100k table +and (under `--spm-crosscheck`) the `sentencepiece` `.model` are checked up front, and a +missing one fails the run - a fixture computed without its reference is worse than none. + +Add `--fuzz 2000` for a deterministic random Unicode corpus (seeded, and the texts +are stored in the fixture, so the Go side never has to reproduce the RNG). Adding +`--spm-crosscheck` also records what Google's SentencePiece implementation returns +for the same text, which is how the two authorities are compared. + +A mismatch means the Go counter and the reference disagree: a bug in the Go counter. +An under-count is the dangerous direction (it lets an oversized input through and +the provider answers 400). +""" + +import argparse +import json +import os +import random +import re +import sys +import unicodedata + +import sentencepiece as spm +import tiktoken +from tokenizers import Tokenizer + + +# A base character carrying two or more combining marks. This is the one shape where +# our NFKC-based normalization provably differs from the served Precompiled charmap: +# NFKC composes `a` + U+0327 + U+0301 into a precomposed letter, while the charmap +# keeps the base and the marks separate (measured: HF normalizes it to "a\u0327\u0301" +# and emits three tokens). The samples are marked so the Go test can bound the +# divergence instead of hiding it: exact token equality everywhere else, and at most +# a two-token difference here. +def has_combining_run(text: str) -> bool: + """True when two or more combining marks sit on the same base character.""" + run = 0 + for ch in text: + if unicodedata.combining(ch): + run += 1 + if run >= 2: + return True + else: + run = 0 + return False + + +def has_dangling_mark(text: str) -> bool: + """True when a combining mark does not follow a letter. + + The served Precompiled charmap drops such marks (measured: "3" + normalizes to "32", and a mark left over after a control character is deleted is + dropped with it), while our NFKC-based normalization keeps them as their own + token - one token more than the model sees. The samples are marked so the Go + test bounds the difference instead of hiding it. + """ + previous = "" + for ch in text: + if unicodedata.combining(ch): + if not previous or not previous[-1].isalpha(): + return True + else: + previous += ch + return False + + +def is_known_approximation(text: str) -> bool: + return has_combining_run(text) or has_dangling_mark(text) + + +# The subset of the model's Precompiled charmap that internal/tokenizer's spm.go +# implements, replicated here so the fixture can tell whether our normalization +# agrees with the served one ON THIS SAMPLE. When it does, the Go counter must match +# the tokenization exactly; when it does not, the sample is marked and the test +# bounds the difference instead of hiding it. +CHARMAP_DELETE_RANGES = ((0x01, 0x08), (0x0B, 0x0C), (0x0E, 0x1F), (0x7F, 0x84), (0x86, 0x9F)) +CHARMAP_SPACE = {"\u200b", "\u200c", "\u200d", "\ufeff", "\u2581"} + + +def our_normalize(text: str) -> str: + """Our SPM normalization: NFKC, charmap subset, space collapse, ▁ escaping.""" + out = [] + for ch in unicodedata.normalize("NFKC", text): + cp = ord(ch) + if any(lo <= cp <= hi for lo, hi in CHARMAP_DELETE_RANGES): + continue + out.append(" " if ch in CHARMAP_SPACE else ch) + s = re.sub(r" {2,}", " ", "".join(out)) + if not s: + return "" + res, prev_space = ["\u2581"], False + for i, ch in enumerate(s): + if ch.isspace(): + if prev_space or i == 0: + prev_space = True + continue + prev_space = True + res.append("\u2581") + continue + prev_space = False + res.append(ch) + return "".join(res) + + +def served_normalize(hf_norm: str) -> str: + """The same string as the served pipeline builds it: Metaspace turns each + whitespace run into one ▁ and adds the leading one.""" + if not hf_norm: + return "" + collapsed = re.sub(r"\s+", " ", hf_norm).strip() + return "\u2581" + collapsed.replace(" ", "\u2581") + + +def normalization_differs(text: str, hf_norm: str) -> bool: + """True when our normalization cannot be expected to match the served one.""" + return our_normalize(text) != served_normalize(hf_norm) + + +ORACLES = { + # label: (tokenizer.json path, what the Go counter is called) + "bge-m3": ("huggingface.co/BAAI/bge-m3/tokenizer.json", "xlmr-spm"), + "bge-large-en": ("huggingface.co/BAAI/bge-large-en-v1.5/tokenizer.json", "bert-wordpiece"), + "qwen3-embedding": ("huggingface.co/Qwen/Qwen3-Embedding-0.6B/tokenizer.json", "qwen-bpe"), + "e5-mistral": ("huggingface.co/intfloat/e5-mistral-7b-instruct/tokenizer.json", "llama-bpe"), + # OpenAI's table, counted by OpenAI's own implementation. + "cl100k": (None, "cl100k_base"), +} + +# The canonical SentencePiece model for the XLM-R family, used only by +# --spm-crosscheck: the .model file is what Google's implementation reads, while +# the served behaviour comes from tokenizer.json above. +SPM_MODEL = "huggingface.co/BAAI/bge-m3/sentencepiece.bpe.model" + +# The cl100k table as shipped in ragflow_deps (the OpenAI blob download_deps.py fetches). +# The Go counter loads this file; tiktoken reads its own cached copy of the same blob, so +# the check in main() is what keeps the check-out honest about what it ships. +CL100K_TABLE = "cl100k_base.tiktoken" + +# The label of the corpus sample, and the incident it reproduces: an English +# Wikipedia article about Argosy magazine with large volume/issue-number tables. +# SILICONFLOW rejected that document with 400/20015 while cl100k counted it as +# *under* the 8192-token window. The document itself arrives as --corpus - nothing +# about it is built into this script, because the corpus is the caller's choice. +INCIDENT_SAMPLE_LABEL = "incident_full" + + +def readable_file(path: str) -> str: + """argparse type for --corpus: fail loudly instead of sampling nothing.""" + if not os.path.isfile(path): + raise argparse.ArgumentTypeError(f"not a file: {path}") + return path + + +def samples(corpus_path: str) -> list[tuple[str, str]]: + out: list[tuple[str, str]] = [ + ("empty", ""), + ("space_only", " "), + ("one_word", "hello"), + ("prose", "The quick brown fox jumps over the lazy dog. " * 40), + ("prose_with_newlines", "First line.\nSecond line.\n\nThird line after a blank.\n" * 20), + ("markdown_table", "| 1940 | Dates: | 6,13,20,27 | 3,10,17,24 | 2,9,16,23,30 |\n" * 30), + ("issue_numbers", "| 1976 | | 383/1 | 383/2 | 383/3 | 383/4 | 383/5 |\n" * 60), + ("digits", "0123456789 " * 100), + ("cjk", "中文分词测试，这是用于对比 tokenizer 的样本。" * 30), + ("mixed_cjk_latin", "RAGFlow 的 ingestion 组件会调用 embedding 模型，例如 BAAI/bge-m3。" * 15), + ("emoji", "🚀🔥 embedding ✅ 测试 " * 30), + ("punctuation", "!!! ??? ... --- ;;; ::: ((( ))) [[[ ]]] " * 40), + ("leading_trailing_space", " padded text with edge whitespace "), + ("tabs_and_crlf", "a\tb\r\nc\t\td\r\n" * 30), + ("accents", "café naïve résumé coöperate Ångström " * 30), + ("unicode_symbols", "≤ ≥ ≈ ∞ ∑ √ π ™ © ® ‰ § ¶ " * 30), + ("long_word", "a" * 500 + " " + "supercalifragilisticexpialidocious" * 5), + ("base64_like", "QWxhZGRpbjpvcGVuIHNlc2FtZQ" * 200), + ("html_table", "
1976383/1
" * 30), + ("numbers_and_units", "296/1 to 296/4, 12:30, 1,234.56 USD, 99.9%, 2020-01-02 " * 25), + # Adversarial Unicode. These exist to check the one place where the SPM + # counter is an approximation: normalization applies NFKC, while the + # served tokenizer applies the model's precompiled nmt_nfkc charmap + # first. If a sample below disagrees with the oracle, the fix is to + # implement the charmap - not to loosen the test. + ("fullwidth_forms", "ＡＢＣＤ１２３４ ｈｅｌｌｏ ！？ " * 20), + ("compatibility_chars", "fi fl ①②③ Ⅻ ㍿ ㎏ ㎞ ¾ ⅓ ㎡ " * 20), + ("decomposed_accents", "e\u0301 a\u0300 o\u0308 u\u030a n\u0303 " * 30), + ("vietnamese", "Tiếng Việt có dấu: Đặng Hữu Phúc, Nguyễn Thị Hương " * 15), + ("rtl_text", "مرحبا بالعالم هذا نص عربي للاختبار " * 15), + ("hebrew", "שלום עולם זהו טקסט לבדיקה " * 15), + ("control_chars", "a\u0001b\u007fc\u200bd\u200ce\u200df " * 20), + ("zero_width_and_nbsp", "a\u00a0b\u2009c\u202fd\u2060e " * 20), + ("cjk_compat", "漢字の互換文字：髙 﨑 塚 邊 ／ 全角カナ アイウ " * 15), + ("emoji_zwj", "👨‍👩‍👧‍👦 👍🏽 🇨🇳 🧑‍💻 " * 20), + ("mixed_scripts", "Latin Ελληνικά Кириллица 中文 日本語 한국어 " * 15), + ("math_and_currency", "∀x∈ℝ: x² ≥ 0 → €100 ± 5% ∞ ≈ ℵ₀ " * 20), + ] + # The WHOLE document, not a window of it: a fixed span would silently stop + # covering the document the moment the document changes, and the point of the + # sample is that what the caller passed in is what gets checked. + with open(corpus_path, encoding="utf-8") as fh: + out.append((INCIDENT_SAMPLE_LABEL, fh.read())) + return out + + +# Alphabets the fuzz corpus draws from. Every one of them has produced, or is +# capable of producing, a divergence between two tokenizer implementations: +# combining marks (normalizer order), zero-width characters (deletable), controls +# (deletable vs spaced), fullwidth/compatibility forms (NFKC vs nmt_nfkc), emoji +# (byte fallback vs Symbols-as-punctuation), and whitespace runs (SPM dummy prefix +# and collapse rules). +FUZZ_ALPHABETS = [ + "abcdefghijklmnopqrstuvwxyz", + "ABCDEFGHIJKLMNOPQRSTUVWXYZ", + "0123456789", + " \t\n", + ".,;:!?()[]{}'\"-—/\\|<>@#$%^&*+=~`", + "中文日本語한국어", + "éàüñçøå", + "αβγδεζηθικλμ", + "абвгдежзийк", + "אבגדהוזחט", + "ابجدهوزحط", + "🙂🚀🔥👍🏽✅🧑‍💻", + "\u0301\u0308\u0327", + "\u200b\u200c\u200d\u2060\ufeff", + "\u0000\u0001\u0009\u007f", + "ＡＢＣ１２３！？　", + "fifl①②Ⅻ㎏㎡¾", +] + + +def fuzz_samples(count: int, seed: int) -> list[tuple[str, str]]: + """Deterministic random Unicode strings, plus a few adversarial extremes. + + The texts are stored in the fixture, so the Go side never has to reproduce this + RNG - it only has to agree with the reference on the text it is given. + """ + rng = random.Random(seed) + out: list[tuple[str, str]] = [] + for i in range(count): + parts = [] + for _ in range(rng.randint(1, 30)): + alphabet = rng.choice(FUZZ_ALPHABETS) + parts.append("".join(rng.choice(alphabet) for _ in range(rng.randint(1, 6)))) + out.append((f"fuzz_{i:04d}", "".join(parts))) + out.append(("fuzz_long_word_99", "a" * 99 + "b")) + out.append(("fuzz_long_word_100", "a" * 100 + "b")) + out.append(("fuzz_long_word_101", "a" * 101 + "b")) + out.append(("fuzz_repeat_run", rng.choice("abcxyz") * rng.randint(2, 200))) + out.append(("fuzz_ascii_all_printable", "".join(chr(c) for c in range(32, 127)))) + out.append(("fuzz_one_codepoint_each", "".join(sorted(FUZZ_ALPHABETS[12] + FUZZ_ALPHABETS[13])))) + return out + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model", choices=sorted(ORACLES), required=True) + parser.add_argument("--out", required=True, help="fixture path to write") + parser.add_argument("--deps", default="ragflow_deps", help="directory holding huggingface.co/") + parser.add_argument( + "--corpus", + required=True, + type=readable_file, + help="document whose ENTIRE content is added as the regression sample (required: the corpus is the caller's, not this script's)", + ) + parser.add_argument("--fuzz", type=int, default=0, help="number of random Unicode samples to append") + parser.add_argument("--seed", type=int, default=20260916, help="seed for --fuzz") + parser.add_argument( + "--spm-crosscheck", + action="store_true", + help="also record what Google's sentencepiece implementation returns (bge-m3 only)", + ) + parser.add_argument("--quiet", action="store_true", help="do not print the per-sample table") + args = parser.parse_args() + + rel, counter_id = ORACLES[args.model] + rows_src = samples(args.corpus) + (fuzz_samples(args.fuzz, args.seed) if args.fuzz else []) + + if rel is None: + # cl100k: OpenAI's own implementation on the table shipped in ragflow_deps. + # tiktoken reads its own cached copy of the blob, so a checkout without the + # shipped table would be checked against something else without a word - + # fail loudly instead, the way the Go loader does. + table = os.path.join(args.deps, CL100K_TABLE) + if not os.path.isfile(table): + print(f"missing cl100k table: {table}\nrun `uv run ragflow_deps/download_deps.py`", file=sys.stderr) + return 1 + encoder = tiktoken.get_encoding("cl100k_base") + rows = [{"label": label, "text": text, "ids": encoder.encode(text), "chars": len(text)} for label, text in rows_src] + oracle = f"tiktoken {getattr(tiktoken, '__version__', '?')} cl100k_base" + else: + tokenizer_path = os.path.join(args.deps, rel) + if not os.path.isfile(tokenizer_path): + print(f"missing oracle tokenizer: {tokenizer_path}\nrun `uv run ragflow_deps/download_go_deps.py`", file=sys.stderr) + return 1 + tokenizer = Tokenizer.from_file(tokenizer_path) + vocab = tokenizer.get_vocab() + unk_id = vocab.get("") + rows = [] + for label, text in rows_src: + encoding = tokenizer.encode(text, add_special_tokens=False) + # HuggingFace renders an unknown token as the *source text* it could not + # encode (Unigram models do: token "🧑" arrives with id 3 = ""), + # while a counter can only report the unknown token itself. The fixture + # stores the canonical token text so the comparison is about the + # segmentation, not about a rendering convention. + tokens = ["" if unk_id is not None and tid == unk_id else tok for tok, tid in zip(encoding.tokens, encoding.ids)] + # Both witnesses are kept: the token texts are numbering-independent and + # are what the Go test compares; the ids are kept for local reasoning + # (and are the primary witness for cl100k, whose numbering is shared). + rows.append( + { + "label": label, + "text": text, + "ids": encoding.ids, + "tokens": tokens, + "chars": len(text), + "known_approx": is_known_approximation(text), + } + ) + oracle = tokenizer_path + + # Counts are derived, never typed: the id list is the evidence, the count is + # its length. Keeping both in the fixture makes a self-inconsistent fixture + # impossible to overlook. + for row in rows: + row["count"] = len(row["ids"]) + + crosscheck_note = None + if args.spm_crosscheck: + model_file = os.path.join(args.deps, SPM_MODEL) + if not os.path.isfile(model_file): + # --spm-crosscheck was asked for, so the asset it reads has to be there: + # skipping quietly would turn "the second opinion agrees" into a claim + # nobody made. Assets are not optional in this script - they are checked. + print(f"missing sentencepiece model: {model_file}\nrun `uv run ragflow_deps/download_go_deps.py`", file=sys.stderr) + return 1 + processor = spm.SentencePieceProcessor(model_file=model_file) + divergent = 0 + for row in rows: + spm_ids = processor.encode(row["text"], out_type=int) + row["spm_count"] = len(spm_ids) + if len(spm_ids) != row["count"]: + divergent += 1 + crosscheck_note = f"Google sentencepiece {getattr(spm, '__version__', '?')} on {SPM_MODEL}: {divergent}/{len(rows)} samples differ from the served tokenizer.json" + print(f"cross-check: {crosscheck_note}") + + fixture = { + "model": args.model, + "counter": counter_id, + "oracle": oracle, + "reference": "https://github.com/huggingface/tokenizers" if rel else "https://github.com/openai/tiktoken", + "samples": rows, + } + if crosscheck_note: + fixture["crosscheck"] = crosscheck_note + with open(args.out, "w", encoding="utf-8") as fh: + json.dump(fixture, fh, ensure_ascii=False, indent=1) + print(f"wrote {args.out}: {len(rows)} samples for {args.model} ({counter_id}), oracle={oracle}") + if not args.quiet: + for row in rows: + extra = f" spm={row['spm_count']}" if "spm_count" in row and row["spm_count"] != row["count"] else "" + print(f" {row['label']:<28} chars={row['chars']:<7} tokens={row['count']:<7}{extra}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/restart_ingestor.sh b/scripts/restart_ingestor.sh new file mode 100755 index 0000000000..5be33c0ad3 --- /dev/null +++ b/scripts/restart_ingestor.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# +# Restart the ingestor so it runs the freshly built bin/ragflow_server. +# +# Why this is a script and not an ad-hoc command: the re-parse deploy requires stopping the running +# ingestor, and stopping it is the one step that can cost in-flight work. The earlier incidents were +# exactly this - tasks delivered to an ingestor that died before it settled them, which JetStream +# then stops redelivering after MaxDeliver (16). So this script checks that nothing is in flight +# first, refuses to proceed if something is, and prints the queue state before and after. +# +# Usage: +# bash scripts/restart_ingestor.sh +# +set -euo pipefail +cd "$(dirname "$0")/.." + +NATS_BOX="natsio/nats-box" +NATS_NET="docker_ragflow" + +queue_state() { + docker run --rm --network "$NATS_NET" "$NATS_BOX" \ + nats --server nats:4222 consumer ls RAGFLOW_TASKS 2>/dev/null | grep RAGFLOW_CONSUMER || true + docker exec docker-mysql-1 sh -c \ + 'MYSQL_PWD=infini_rag_flow mysql -N -uroot rag_flow -e "SELECT CONCAT(\"documents run=1: \", COUNT(*)) FROM document WHERE run=1;"' 2>/dev/null || true +} + +echo "=== before ===" +queue_state + +inflight=$(docker run --rm --network "$NATS_NET" "$NATS_BOX" \ + nats --server nats:4222 consumer ls RAGFLOW_TASKS 2>/dev/null | grep RAGFLOW_CONSUMER | awk '{print $(NF-2)}' | tr -d '│ ') +if [ "${inflight:-0}" != "0" ]; then + echo "refusing to restart: $inflight message(s) in flight; wait for them to settle" >&2 + exit 1 +fi + +echo "=== restarting ingestor ===" +pkill -f "ragflow_server --ingestor" || true +sleep 3 +setsid nohup ./bin/ragflow_server --ingestor >>logs/ingestor_server.log 2>&1 & +sleep 30 +pgrep -af "ragflow_server --ingestor" | cut -c1-60 + +echo "=== after ===" +queue_state +tail -3 logs/ingestor_server.log | cut -c1-120 diff --git a/scripts/tokenizer_coverage.py b/scripts/tokenizer_coverage.py new file mode 100755 index 0000000000..08becf5383 --- /dev/null +++ b/scripts/tokenizer_coverage.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Report how the embedding models in the catalog are covered by tokenizer counters. + +The catalog (`conf/all_models.json`) declares a `tokenizer` per model when the model +has been matched to a counter in `internal/tokenizer`. Everything else falls back to +the calibrated path at runtime, which is safe but less precise - so the number that +matters operationally is "how many models are tagged, and which untagged ones could be +tagged cheaply because their architecture is already implemented". + +This script answers that from the catalog, with no network and no guesses about +individual models: it groups the tagged models by counter, then groups the untagged +ones by a *name heuristic* into the architecture they most likely use, and prints the +result as markdown so it can be pasted into +`internal/tokenizer/embedding_token_limits.md`. + +Usage: + python3 scripts/tokenizer_coverage.py # human-readable + python3 scripts/tokenizer_coverage.py --markdown # table for the doc + python3 scripts/tokenizer_coverage.py --untagged # list every untagged model +""" + +import argparse +import collections +import json +import sys + +CATALOG = "conf/all_models.json" + +# Counters implemented in internal/tokenizer (see embedding_token_limits.md). A model +# tagged with one of these is counted exactly; anything else is calibrated at runtime. +IMPLEMENTED = { + "cl100k_base": "tiktoken BPE (OpenAI table)", + "xlmr-spm": "XLM-R Unigram SentencePiece", + "bert-wordpiece": "BERT WordPiece", + "qwen-bpe": "byte-level BPE (Qwen pre-tokenizer regex)", + "llama-bpe": "SentencePiece-BPE (▁ prepend/replace, byte fallback)", +} + +# Name heuristics for the untagged models. They are deliberately conservative: a wrong +# guess here means tagging a model with a counter that does not match its tokenizer, +# which is worse than leaving it calibrated. Anything not matched lands in +# "arch-mismatch" (a family we do not implement) or "hosted" (no downloadable artifact). +FAMILY_HINTS = [ + ("xlmr-spm", ("bge-m3", "multilingual-e5", "m3e", "gte-multilingual", "jina-embeddings-v3", "paraphrase-multilingual", "labse", "infgrad", "bge-multilingual")), + ("bert-wordpiece", ("bge-", "e5-", "gte-", "jina-embeddings-v2", "bce-", "text2vec-", "mxbai", "gte-base", "gte-large", "gte-small", "arabic-", "multilingual-e5-small")), + ("qwen-bpe", ("qwen3-embedding", "gte-qwen", "qwen3-reranker")), + ("llama-bpe", ("e5-mistral", "mistral-embed", "nomic-embed", "nemoretriever", "sfr-embedding", "linq-embed", "llama-embed", "bge-en-icl", "stella")), + ("cl100k_base", ("text-embedding-3", "text-embedding-ada")), +] + +HOSTED_HINTS = ( + "gemini", + "text-embedding-004", + "cohere", + "embed-multilingual", + "embed-english", + "voyage", + "titan", + "amazon.", + "baichuan", + "jina-clip", + "cloudsway", + "upstage", + "nvidia", + "gme-", + "ibm", + "watsonx", + "zhipu", + "zai-org", + "yi-", + "ernie", + "doubao", + "minimax", + "tencent", + "siliconflow", + "text-similarity", + "embedding-v1", +) + +# Task and packaging variants of the same model. The catalog lists one entry per +# variant (a retrieval/classification/clustering head, a GGUF or MLX conversion, a +# quantisation), which inflates the model count without adding a tokenizer to verify: +# they all share the base model's tokenizer file. Counting base models is what tells +# us how much tokenizer work is actually left. +VARIANT_SUFFIXES = ( + "-classification", + "-clustering", + "-retrieval", + "-text-matching", + "-reranking", + "-GGUF", + "-gguf", + "-mlx", + "-qat-q4_0-unquantized", + "-qat-q8_0-unquantized", + "-unquantized", + "-int8", + "-fp16", + "-onnx", + "-v0", + "-v1", +) + + +def base_name(name: str) -> str: + """The model a catalog entry is a variant of (`.../x-retrieval-mlx` -> `.../x`).""" + base = name.split(":")[0] + changed = True + while changed: + changed = False + for suffix in VARIANT_SUFFIXES: + if base.endswith(suffix): + base = base[: -len(suffix)] + changed = True + return base + + +def embedding_models(catalog: dict) -> list[dict]: + out = [] + for model in catalog["models"]: + types = model.get("model_types") or [] + if not isinstance(types, list): + types = [types] + if any("embed" in str(t).lower() for t in types): + out.append(model) + return out + + +def guess_family(name: str) -> str: + lowered = name.lower() + for family, hints in FAMILY_HINTS: + if any(hint in lowered for hint in hints): + return family + if any(hint in lowered for hint in HOSTED_HINTS): + return "hosted" + return "other-architecture" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--catalog", default=CATALOG) + parser.add_argument("--markdown", action="store_true", help="print the tables as markdown") + parser.add_argument("--untagged", action="store_true", help="list every untagged model by guess") + args = parser.parse_args() + + try: + with open(args.catalog, encoding="utf-8") as handle: + catalog = json.load(handle) + except (OSError, json.JSONDecodeError) as err: + print(f"cannot read {args.catalog}: {err}", file=sys.stderr) + return 1 + + models = embedding_models(catalog) + tagged = collections.defaultdict(list) + untagged = collections.defaultdict(list) + bases: dict[str, str] = {} + for model in models: + name = model.get("name", "?") + bases.setdefault(base_name(name), name) + counter = model.get("tokenizer") + if counter: + tagged[counter].append(name) + else: + untagged[guess_family(name)].append(base_name(name)) + + total = len(models) + tagged_n = sum(len(v) for v in tagged.values()) + print(f"# Tokenizer coverage of the model catalog ({args.catalog})") + print() + print(f"embedding entries: **{total}** | distinct base models: **{len(bases)}** | tagged entries: **{tagged_n}** | calibrated entries: **{total - tagged_n}**") + print() + print("## Tagged (counted exactly)") + print() + if args.markdown: + print("| counter | entries | distinct base models | scheme |") + print("|---|---|---|---|") + for counter, names in sorted(tagged.items(), key=lambda kv: -len(kv[1])): + bases_here = len({base_name(n) for n in names}) + print(f"| `{counter}` | {len(names)} | {bases_here} | {IMPLEMENTED.get(counter, '**not implemented**')} |") + else: + for counter, names in sorted(tagged.items(), key=lambda kv: -len(kv[1])): + print(f" {counter:<18} {len(names):>3} {IMPLEMENTED.get(counter, 'NOT IMPLEMENTED')}") + print() + print("## Untagged (calibrated at runtime; grouped by name heuristic)") + print() + if args.markdown: + print("| likely architecture | entries | distinct base models | counter exists? |") + print("|---|---|---|---|") + for family, names in sorted(untagged.items(), key=lambda kv: -len(kv[1])): + exists = "yes - only the tag is missing" if family in IMPLEMENTED else ("no - new implementation, plus an oracle" if family != "hosted" else "no artifact to verify against") + print(f"| {family} | {len(names)} | {len(set(names))} | {exists} |") + else: + for family, names in sorted(untagged.items(), key=lambda kv: -len(kv[1])): + print(f" {family:<20} {len(names):>3} entries, {len(set(names)):>3} base models") + if args.untagged: + print() + print("## Every untagged model") + print() + for family, names in sorted(untagged.items(), key=lambda kv: -len(kv[1])): + print(f"### {family} ({len(names)})") + for name in sorted(names): + print(f"- {name}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/covsweep/main.go b/tools/covsweep/main.go new file mode 100644 index 0000000000..e161021442 --- /dev/null +++ b/tools/covsweep/main.go @@ -0,0 +1,1037 @@ +// Command covsweep answers one question for a whole document set: +// +// does the search index still serve every passage of every source document? +// +// It exists because "document.chunk_num equals the number of chunk documents in ES" only proves +// the two counters agree - it cannot see a document whose chunks silently dropped a passage. The +// Python scripts under scripts/ answer the same question for a handful of documents; this tool +// answers it for a whole document set in one pass, fast enough to run after every repair, which is +// what makes the answer useful. +// +// The document set is never assumed: --docs, --corpus, --index and --kb have no built-in defaults, +// so a run either names its inputs or stops before it reads anything. +// +// Method, per document: +// +// 1. pull every chunk text of the document out of ES (kb_id + docnm_kwd filter) and join them; +// 2. project both the source file and the joined chunk text onto their alphanumeric characters +// (lowercased), remembering for every file character which byte of the original it came from; +// 3. match 16-character k-grams between the two projections with a rolling hash; a k-gram the +// chunk text contains covers the file bytes it came from. The projection is what makes this +// honest: the ingest strips punctuation AND closes the gap it leaves, so a source line +// "* China at the Asian Games * China at the Paralympics" is served as "...GamesChina at the +// Paralympics" - comparing bytes or word tokens marks that passage missing, which is how an +// earlier version of this tool reported a uniform ~78% coverage with no document above 90%; +// 4. report a run of uncovered file bytes as a missing passage only when it is at least --min-gap +// bytes long, which is far longer than the remainder of a re-punctuated boundary. +// +// --method window keeps the raw byte-window comparison for the record, where the false-alarm rate +// on markup-heavy documents can be seen. +// +// Verdict semantics: a word counts as SERVED only when the chunk text carries it as its own token. +// Gluing is not tolerated - an entry whose words were run together ("ChinaChina", "HathiTrustThe") +// cannot be matched by a query for "China", so those words are defects even though their characters +// are present. The report splits them into "content never arrived" and "glued into a longer token" +// so the two failure modes can be told apart, but both land in the flagged list: +// +// /coverage_go.jsonl one row per document (word_pct, missing_strict_words, absent/fused splits, gaps) +// /below99.txt documents with any unserved word - strict, includes glued ones +// /fused.txt the glued subset, for measuring a separator fix +// +// Concurrency, and where the time actually goes (the tool prints a phase breakdown on exit): +// +// fetchers --shards goroutines, each paging its own shard of the document set with +// search_after. A page is megabytes of content_with_weight, so each shard also runs a +// decoder goroutine: the fetch asks for page N+1 while the decode parses page N, +// otherwise the socket idles for the whole parse; +// workers --workers goroutines doing tokenising and hashing; +// writer one goroutine appending JSONL and printing progress. +// +// A document is emitted as soon as the chunk texts seen for it reach the chunk_num column of +// --docs, so one slow shard cannot stall the rest; the flip side is that a stale chunk_num either +// truncates a document or emits it with nothing at all (see Usage). +// +// The phase breakdown exists because this workload does not scale the way "add goroutines" +// suggests: the index has few shards, so the fetch half saturates before the CPU does, and the +// compute half is dominated by the largest documents rather than by their number - measuring beats +// guessing. +// +// Usage: +// +// # "\t\t" rows, one per document, with one source file per row's name +// # column in --corpus. Regenerate the list immediately before a run: chunk_num is the per +// # document threshold above, so a stale value truncates a document, or at 0 emits it empty. +// docker exec sh -c 'MYSQL_PWD= mysql -N -uroot \ +// -e "SELECT id,name,chunk_num FROM document WHERE kb_id=\"\";"' > /tmp/cov_docs.tsv +// +// go run ./tools/covsweep --docs /tmp/cov_docs.tsv --corpus \ +// --index --kb --out outputs/coverage +// go run ./tools/covsweep ... --limit 400 # calibrate on a subset +// go run ./tools/covsweep ... --only-file /below99.txt # re-check only the flagged ones +package main + +import ( + "bufio" + "bytes" + "encoding/base64" + "encoding/json" + "flag" + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "regexp" + "runtime" + "sort" + "strings" + "sync" + "sync/atomic" + "syscall" + "time" + "unicode" +) + +type doc struct { + ID string + Name string + Chunks int +} + +type job struct { + Name string + ID string + Chunks int + Texts []string +} + +type gap struct { + Start int `json:"start"` + Length int `json:"length"` + Context string `json:"context"` +} + +type result struct { + Name string `json:"name"` + ID string `json:"id"` + MySQLChunks int `json:"mysql_chunks"` + ChunkCount int `json:"chunk_count"` + FileChars int `json:"file_chars"` + FileAlnum int `json:"file_alnum"` // alphanumeric characters the file projects to + CoveredChar int `json:"covered_chars"` + CoveragePct float64 `json:"coverage_pct"` // order-sensitive: 16-gram containment + GapChars int `json:"gap_chars"` + // WordPct is the order-insensitive answer to "is anything missing": the share of the file's + // word occurrences (multiset) that the chunk text also carries. It is the headline number + // because the ingest legitimately re-orders words inside tables (a source row + // "Longlist | 2021 | Baba Badji | Ghost Letters" is served as the cells + // "		2021Baba BadjiGhost LettersLonglist