infiniflow/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-10-12 02:16:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(go,web): port tokenizer counters, embedding limits and multi-provider web search (#19802)

Browse Source
This commit is contained in:
Zhichang Yu
2026-09-18 19:11:28 +08:00
committed by GitHub
parent 8f8699cb6c
commit 44b8d366ba
194 changed files with 14173 additions and 1746 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/sep-tests.yml
+84 -10
View File
@@ -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/<repo>/<file>). 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/<repo>/<file>). 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: |
Dockerfile
+23
View File
@@ -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/<repo>/<file>; 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
Dockerfile_base
+21
View File
@@ -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/<repo>/<file>; 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
Dockerfile_go
+21
View File
@@ -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 (<workdir>/ragflow_deps/huggingface.co/<repo>/<file>). 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
build.sh
+71
View File
@@ -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
cmd/ragflow_server.go
+59 -10
View File
@@ -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
conf/all_models.json
+42
View File
@@ -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
conf/mapping.json
+6 -4
View File
@@ -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 @@
}
]
}
}
}
docker/docker-compose-base.yml
+37 -8
View File
@@ -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
go.mod
+28 -27
View File
@@ -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
go.sum
+15 -4
View File
@@ -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=
internal/agent/component/llm.go
+35 -168
View File
@@ -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) {
internal/agent/retrievalbridge/runtime_adapter.go
+38 -45
View File
@@ -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
}
internal/agent/retrievalbridge/runtime_adapter_test.go
+76 -6
View File
@@ -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)
}
}
internal/agent/runtime/maps.go
+73
View File
@@ -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
}
internal/agent/runtime/retrieval.go
+22 -20
View File
@@ -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.
internal/agent/tool/agentic_search.go
+5 -4
View File
@@ -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
}
internal/agent/tool/dataset_navigation.go
+9 -9
View File
@@ -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)
internal/agent/tool/dataset_navigation_test.go
+24
View File
@@ -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))
}
internal/agent/tool/retrieval.go
-8
View File
@@ -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"
internal/agent/tool/retrieval_nlp.go
+52 -64
View File
@@ -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 }
internal/agent/tool/retrieval_nlp_test.go
+38
View File
@@ -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) {
internal/agent/tool/retrieval_service.go
+23 -242
View File
@@ -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() }
internal/agent/tool/retrieval_wiring_test.go
+10 -2
View File
@@ -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,
})
internal/common/logger.go
+135 -33
View File
@@ -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/<Path>
// 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=<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)
internal/common/logger_ctx_test.go
+115
View File
@@ -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")
}
}
internal/common/model_assets.go
+76
View File
@@ -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.
// `<dir>/huggingface.co/<repo>/<file>` (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/<repo>/<file>` 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:
//
// <dir>/huggingface.co/<repo>/<file> the ragflow_deps-like root
// <dir>/<repo>/<file> a huggingface.co-like root
// <dir>/<file> 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
}
internal/common/model_assets_test.go
+69
View File
@@ -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)
}
}
internal/dao/ingestion_task_test.go
+47
View File
@@ -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)
}
}
internal/dao/tenant_model.go
+19
View File
@@ -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
internal/engine/elasticsearch/chunk.go
+146 -64
View File
@@ -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,
},
internal/engine/elasticsearch/chunk_helpers_test.go
+17 -2
View File
@@ -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)
internal/engine/elasticsearch/chunk_readback_integration_test.go
+2 -2
View File
@@ -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"
internal/engine/elasticsearch/chunk_refresh_test.go
+83
View File
@@ -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])
}
}
internal/engine/elasticsearch/chunk_test.go
+10 -1
View File
@@ -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)
internal/engine/elasticsearch/kg_test.go
+2 -2
View File
@@ -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")
internal/engine/nats/consumer_selfheal_test.go
+143
View File
@@ -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")
}
}
internal/engine/nats/knowledgecompile.go
+1 -1
View File
@@ -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)
internal/engine/nats/nats.go
+166 -22
View File
@@ -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
}
internal/engine/nats/stream_config_test.go
+12 -11
View File
@@ -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)
internal/engine/nats/syncer.go
+25 -3
View File
@@ -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
internal/entity/models/302ai.go
+1 -1
View File
@@ -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 {
internal/entity/models/aliyun.go
+1 -1
View File
@@ -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)
internal/entity/models/astraflow.go
+1 -1
View File
@@ -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
internal/entity/models/avian.go
+1 -1
View File
@@ -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 {
internal/entity/models/azure_openai.go
+1 -1
View File
@@ -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
internal/entity/models/baichuan.go
+1 -1
View File
@@ -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 {
internal/entity/models/baidu.go
+1 -1
View File
@@ -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 {
internal/entity/models/base_model.go
+16 -3
View File
@@ -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)
internal/entity/models/chat_tools.go
+22 -20
View File
@@ -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++ {
internal/entity/models/cometapi.go
+1 -1
View File
@@ -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
internal/entity/models/deepinfra.go
+1 -1
View File
@@ -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 {
internal/entity/models/deepseek.go
+1 -1
View File
@@ -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)
internal/entity/models/futurmix.go
+1 -1
View File
@@ -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
internal/entity/models/gitee.go
+1 -1
View File
@@ -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)
internal/entity/models/gpustack.go
+1 -1
View File
@@ -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.
internal/entity/models/greenpt.go
+1 -1
View File
@@ -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
internal/entity/models/groq.go
+1 -1
View File
@@ -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 {
internal/entity/models/huaweicloud.go
+1 -1
View File
@@ -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 {
internal/entity/models/huggingface.go
+1 -1
View File
@@ -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 {
internal/entity/models/hunyuan.go
+1 -1
View File
@@ -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
internal/entity/models/jiekouai.go
+1 -1
View File
@@ -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 {
internal/entity/models/jina.go
+1 -1
View File
@@ -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 {
internal/entity/models/llm.go
+534 -50
View File
@@ -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).
internal/entity/models/llm_test.go
+83 -1
View File
@@ -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)
}
}
}
internal/entity/models/lmstudio.go
+1 -1
View File
@@ -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)
internal/entity/models/localai.go
+1 -1
View File
@@ -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
internal/entity/models/longcat.go
+1 -1
View File
@@ -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
internal/entity/models/minimax.go
+1 -1
View File
@@ -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)
internal/entity/models/mistral.go
+1 -1
View File
@@ -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
internal/entity/models/model.go
+59 -5
View File
@@ -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) {
internal/entity/models/model_test.go
+235
View File
@@ -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
}
internal/entity/models/modelscope.go
+1 -1
View File
@@ -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.
internal/entity/models/moonshot.go
+1 -1
View File
@@ -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)
internal/entity/models/mws.go
+1 -1
View File
@@ -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.
internal/entity/models/n1n.go
+1 -1
View File
@@ -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.
internal/entity/models/novita.go
+1 -1
View File
@@ -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
internal/entity/models/nvidia.go
+1 -1
View File
@@ -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 {
internal/entity/models/openai.go
+11 -1
View File
@@ -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
}
internal/entity/models/openai_api_compatible.go
+1 -1
View File
@@ -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
internal/entity/models/openrouter.go
+1 -1
View File
@@ -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 {
internal/entity/models/orcarouter.go
+1 -1
View File
@@ -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 {
internal/entity/models/perplexity.go
+1 -1
View File
@@ -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 {
internal/entity/models/ppio.go
+1 -1
View File
@@ -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 {
internal/entity/models/qiniu.go
+1 -1
View File
@@ -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 {
internal/entity/models/ragcon.go
+1 -1
View File
@@ -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 {
internal/entity/models/response_handler.go
+3 -2
View File
@@ -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))
}
}
internal/entity/models/siliconflow.go
+9 -1
View File
@@ -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))
}
internal/entity/models/stepfun.go
+1 -1
View File
@@ -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
internal/entity/models/togetherai.go
+1 -1
View File
@@ -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 {
internal/entity/models/tokenhub.go
+1 -1
View File
@@ -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 {
internal/entity/models/tokenpony.go
+1 -1
View File
@@ -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
internal/entity/models/types.go
+38 -5
View File
@@ -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
internal/entity/models/upstage.go
+1 -1
View File
@@ -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
internal/entity/models/vllm.go
+1 -1
View File
@@ -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)
internal/entity/models/volcengine.go
+1 -1
View File
@@ -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)
internal/entity/models/xai.go
+1 -1
View File
@@ -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
internal/entity/models/xiaomi.go
+1 -1
View File
@@ -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 {
internal/entity/models/xinference.go
+1 -1
View File
@@ -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.
internal/entity/models/xunfei.go
+1 -1
View File
@@ -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 {

Some files were not shown because too many files have changed in this diff Show More