From 3c0f8fc19297bbd8a36231f008ae389b32216e4b Mon Sep 17 00:00:00 2001 From: Jack Date: Mon, 13 Jul 2026 11:08:04 +0800 Subject: [PATCH] Refactor: refactor dataflow_service.go and guard nats message re-delivery (#16826) ### Summary 1. refactor dataflow_service.go 2. guard nats message re-delivery 3. support document parse cancelling & re-run --- cmd/ragflow_server.go | 20 +- go.mod | 11 +- go.sum | 18 +- internal/common/task.go | 5 + internal/deepdoc/parser/pdf/parser.go | 2 +- .../pdf/parser_pipeline_integration_test.go | 26 +- .../deepdoc/parser/pdf/test_helpers_test.go | 24 + internal/engine/elasticsearch/client.go | 5 + internal/engine/engine.go | 3 + internal/engine/infinity/client.go | 5 + internal/engine/nats/nats.go | 4 + .../component/slice3_integration_test.go | 86 ++ internal/ingestion/component/slice3_test.go | 65 +- internal/ingestion/component/tokenizer.go | 16 +- .../ingestion/component/tokenizer_test.go | 98 ++- internal/ingestion/pipeline/errors.go | 13 +- .../{test_helpers.go => helpers_test.go} | 53 +- internal/ingestion/pipeline/payload.go | 84 ++ internal/ingestion/pipeline/payload_test.go | 198 +++++ internal/ingestion/pipeline/pipeline.go | 153 ++-- internal/ingestion/pipeline/pipeline_test.go | 205 ++++- .../pipeline/template_integration_test.go | 52 +- internal/ingestion/service/doc_state.go | 80 ++ internal/ingestion/service/doc_state_test.go | 122 +++ .../service/execute_task_ack_test.go | 342 ++++++++ .../ingestion/service/execute_task_test.go | 84 +- internal/ingestion/service/heartbeat_test.go | 102 +++ .../ingestion/service/ingestion_service.go | 542 +++++++++---- .../service/ingestor_lifecycle_test.go | 57 ++ .../ingestion/service/process_message_test.go | 180 +++++ internal/ingestion/service/progress_sink.go | 105 +++ .../ingestion/service/progress_sink_test.go | 281 +++++++ ...test.go => real_consumer_pipeline_test.go} | 45 +- internal/ingestion/service/run_task_test.go | 274 +++++++ internal/ingestion/task/chunk_index_writer.go | 70 ++ .../ingestion/task/chunk_index_writer_test.go | 116 +++ internal/ingestion/task/chunk_process.go | 11 +- internal/ingestion/task/chunk_process_test.go | 50 +- internal/ingestion/task/chunk_utils.go | 4 +- internal/ingestion/task/chunk_utils_test.go | 60 +- internal/ingestion/task/constants.go | 2 +- internal/ingestion/task/dataflow_service.go | 525 ------------ .../ingestion/task/dataflow_service_test.go | 754 ------------------ internal/ingestion/task/embedder.go | 95 +++ internal/ingestion/task/embedder_test.go | 177 ++++ internal/ingestion/task/embedding_test.go | 77 -- .../ingestion/task/embedding_token_test.go | 55 -- internal/ingestion/task/golden_compare.go | 14 +- .../ingestion/task/golden_compare_test.go | 4 +- ...aflow_e2e_test.go => pipeline_e2e_test.go} | 47 +- internal/ingestion/task/pipeline_executor.go | 237 ++++++ .../task/pipeline_executor_defaults.go | 89 +++ .../ingestion/task/pipeline_executor_test.go | 442 ++++++++++ ...t.go => pipeline_real_integration_test.go} | 77 +- internal/ingestion/task/pipeline_run_test.go | 96 +++ internal/ingestion/task/real_consumer_test.go | 88 +- internal/ingestion/task/task_context.go | 7 +- internal/ingestion/task/task_context_test.go | 63 +- internal/ingestion/task/task_handler.go | 55 +- internal/ingestion/task/task_handler_test.go | 140 +--- ...w_golden.go => compare_pipeline_golden.go} | 2 +- internal/ingestion/testutil/nats_helper.go | 70 ++ internal/service/chunk/chunk.go | 4 - internal/service/chunk/chunk_test.go | 2 + internal/service/chunk_types.go | 2 - internal/service/dataset.go | 10 +- internal/service/document.go | 61 +- internal/service/document_test.go | 11 +- internal/service/ingestion_task_service.go | 132 ++- .../service/ingestion_task_service_test.go | 293 ++++++- 70 files changed, 4901 insertions(+), 2401 deletions(-) create mode 100644 internal/ingestion/component/slice3_integration_test.go rename internal/ingestion/pipeline/{test_helpers.go => helpers_test.go} (77%) create mode 100644 internal/ingestion/pipeline/payload.go create mode 100644 internal/ingestion/pipeline/payload_test.go create mode 100644 internal/ingestion/service/doc_state.go create mode 100644 internal/ingestion/service/doc_state_test.go create mode 100644 internal/ingestion/service/execute_task_ack_test.go create mode 100644 internal/ingestion/service/heartbeat_test.go create mode 100644 internal/ingestion/service/ingestor_lifecycle_test.go create mode 100644 internal/ingestion/service/process_message_test.go create mode 100644 internal/ingestion/service/progress_sink.go create mode 100644 internal/ingestion/service/progress_sink_test.go rename internal/ingestion/service/{real_consumer_dataflow_test.go => real_consumer_pipeline_test.go} (72%) create mode 100644 internal/ingestion/service/run_task_test.go create mode 100644 internal/ingestion/task/chunk_index_writer.go create mode 100644 internal/ingestion/task/chunk_index_writer_test.go delete mode 100644 internal/ingestion/task/dataflow_service.go delete mode 100644 internal/ingestion/task/dataflow_service_test.go create mode 100644 internal/ingestion/task/embedder.go create mode 100644 internal/ingestion/task/embedder_test.go delete mode 100644 internal/ingestion/task/embedding_test.go delete mode 100644 internal/ingestion/task/embedding_token_test.go rename internal/ingestion/task/{dataflow_e2e_test.go => pipeline_e2e_test.go} (88%) create mode 100644 internal/ingestion/task/pipeline_executor.go create mode 100644 internal/ingestion/task/pipeline_executor_defaults.go create mode 100644 internal/ingestion/task/pipeline_executor_test.go rename internal/ingestion/task/{dataflow_real_pipeline_integration_test.go => pipeline_real_integration_test.go} (89%) create mode 100644 internal/ingestion/task/pipeline_run_test.go rename internal/ingestion/task/tool/{compare_dataflow_golden.go => compare_pipeline_golden.go} (99%) create mode 100644 internal/ingestion/testutil/nats_helper.go diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 63fc5e94d8..177626aec0 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -482,14 +482,8 @@ func runAdmin(args *serverArgs) error { func runIngestor(args *serverArgs) error { // Initialize tokenizer (rag_analyzer) - dictPath := common.GetEnv(common.EnvRAGFlowDictPath) - if dictPath == "" { - dictPath = "/usr/share/infinity/resource" - } - tokenizerCfg := &tokenizer.PoolConfig{ - DictPath: dictPath, - } - if err := tokenizer.Init(tokenizerCfg); err != nil { + // tokenizer.Init handles DictPath fallback: env var → /usr/share/infinity/resource + if err := tokenizer.Init(&tokenizer.PoolConfig{}); err != nil { common.Fatal("Failed to initialize tokenizer", zap.Error(err)) } defer tokenizer.Close() @@ -647,13 +641,9 @@ func runAPI(args *serverArgs) error { local.InitAdminStatus(1, "admin server not connected") // Initialize tokenizer (rag_analyzer) - dictPath := common.GetEnv(common.EnvRAGFlowDictPath) - if dictPath == "" { - dictPath = "/usr/share/infinity/resource" - } - tokenizerCfg := &tokenizer.PoolConfig{ - DictPath: dictPath, - } + // tokenizer.Init fills DictPath from env var or default, so + // tokenizerCfg.DictPath carries the resolved path for downstream use. + tokenizerCfg := &tokenizer.PoolConfig{} if err := tokenizer.Init(tokenizerCfg); err != nil { common.Fatal("Failed to initialize tokenizer", zap.Error(err)) } diff --git a/go.mod b/go.mod index c701f2cfde..20b551fd9a 100644 --- a/go.mod +++ b/go.mod @@ -35,6 +35,7 @@ require ( github.com/kaptinlin/jsonrepair v0.4.8 github.com/lib/pq v1.10.9 github.com/minio/minio-go/v7 v7.0.99 + 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 github.com/peterh/liner v1.2.2 @@ -61,7 +62,6 @@ require ( golang.org/x/text v0.38.0 google.golang.org/api v0.287.1 google.golang.org/genai v1.54.0 - google.golang.org/grpc v1.82.0 gopkg.in/natefinch/lumberjack.v2 v2.2.1 gopkg.in/yaml.v3 v3.0.1 gorm.io/driver/mysql v1.5.2 @@ -85,6 +85,7 @@ 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/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect github.com/apache/thrift v0.23.0 // 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 @@ -131,6 +132,7 @@ require ( github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe // indirect github.com/golang-sql/sqlexp v0.1.0 // indirect github.com/google/go-cmp v0.7.0 // indirect + github.com/google/go-tpm v0.9.8 // 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 @@ -140,7 +142,7 @@ require ( github.com/hashicorp/hcl v1.0.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect - github.com/klauspost/compress v1.18.5 // indirect + github.com/klauspost/compress v1.18.6 // indirect github.com/klauspost/cpuid/v2 v2.2.11 // indirect github.com/klauspost/crc32 v1.3.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect @@ -149,11 +151,13 @@ require ( github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-runewidth v0.0.3 // indirect github.com/minio/crc64nvme v1.1.1 // indirect + 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/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect - github.com/nats-io/nkeys v0.4.15 // indirect + github.com/nats-io/jwt/v2 v2.8.2 // indirect + github.com/nats-io/nkeys v0.4.16 // indirect github.com/nats-io/nuid v1.0.1 // indirect github.com/oapi-codegen/runtime v1.4.0 // indirect github.com/pelletier/go-toml/v2 v2.2.2 // indirect @@ -206,6 +210,7 @@ require ( 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 + google.golang.org/grpc v1.82.0 // indirect google.golang.org/protobuf v1.36.11 // indirect gopkg.in/ini.v1 v1.67.0 // indirect modernc.org/libc v1.22.5 // indirect diff --git a/go.sum b/go.sum index 1db8268e78..fdc2f34f06 100644 --- a/go.sum +++ b/go.sum @@ -87,6 +87,8 @@ github.com/aliyun/credentials-go v1.3.1/go.mod h1:8jKYhQuDawt8x2+fusqa1Y6mPxemTs github.com/aliyun/credentials-go v1.3.6/go.mod h1:1LxUuX7L5YrZUWzBrRyk0SwSdH4OmPrib8NVePL3fxM= github.com/aliyun/credentials-go v1.4.5 h1:O76WYKgdy1oQYYiJkERjlA2dxGuvLRrzuO2ScrtGWSk= github.com/aliyun/credentials-go v1.4.5/go.mod h1:Jm6d+xIgwJVLVWT561vy67ZRP4lPTQxMbEYRuT2Ti1U= +github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op h1:Z/MZK75wC/NSrkgqeNIa7jexam9uWzhLmFTSCPI/kn0= +github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op/go.mod h1:FQyySiasQQM8735Ddel3MRojmy4dA1IqCeyJ5jmPMbI= github.com/apache/thrift v0.23.0 h1:wKR6YnefQSEnxpEfmgTPuJibNG4bF0p2TK34tHLWi3s= github.com/apache/thrift v0.23.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ= @@ -272,6 +274,8 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/go-tpm v0.9.8 h1:slArAR9Ft+1ybZu0lBwpSmpwhRXaa85hWtMinMyRAWo= +github.com/google/go-tpm v0.9.8/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc= github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0= @@ -317,8 +321,8 @@ github.com/kaptinlin/jsonrepair v0.4.8 h1:9oaoEe/vaKgm8ko4TLjBLUEog6tBW6WUzZXLPL github.com/kaptinlin/jsonrepair v0.4.8/go.mod h1:eWRC42KDUT0MHkMplUN6necu59FQFqKOKe+86akpY3g= github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= -github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE= -github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao= +github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.0.1/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.2.11 h1:0OwqZRYI2rFrjS4kvkDnqJkKHdHaRnCm68/DY4OxRzU= github.com/klauspost/cpuid/v2 v2.2.11/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= @@ -350,6 +354,8 @@ github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1f github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= github.com/minio/crc64nvme v1.1.1 h1:8dwx/Pz49suywbO+auHCBpCtlW1OfpcLN7wYgVR6wAI= github.com/minio/crc64nvme v1.1.1/go.mod h1:eVfm2fAzLlxMdUGc0EEBGSMmPwmXD5XiNRpnu9J3bvg= +github.com/minio/highwayhash v1.0.4 h1:asJizugGgchQod2ja9NJlGOWq4s7KsAWr5XUc9Clgl4= +github.com/minio/highwayhash v1.0.4/go.mod h1:GGYsuwP/fPD6Y9hMiXuapVvlIUEhFhMTh0rxU3ik1LQ= github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34= github.com/minio/md5-simd v1.1.2/go.mod h1:MzdKDxYpY2BT9XQFocsiZf/NKVtR7nkE4RoEpN+20RM= github.com/minio/minio-go/v7 v7.0.99 h1:2vH/byrwUkIpFQFOilvTfaUpvAX3fEFhEzO+DR3DlCE= @@ -364,10 +370,14 @@ github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3Rllmb github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8= +github.com/nats-io/jwt/v2 v2.8.2 h1:XXRgB60MSTnqsRwejQurVDs/hcv2dkt+86GjI+I/bMc= +github.com/nats-io/jwt/v2 v2.8.2/go.mod h1:Ag/56sq9OblL4JgdYufDd16Egb17Kr/8WwwuO/forVc= +github.com/nats-io/nats-server/v2 v2.14.3 h1:+xjydPt7rkit67G+04TN0mcO2n+8nveZE7tK/PPV53A= +github.com/nats-io/nats-server/v2 v2.14.3/go.mod h1:5IlCtBzfwyzQzPMjmoJ9W2/LKmnJRtNyuOs/OT+NHDY= github.com/nats-io/nats.go v1.52.0 h1:n3avV4VBsCgsdwh71TppsTwtv+QdPs7ntSKM8qJLGsc= github.com/nats-io/nats.go v1.52.0/go.mod h1:26HypzazeOkyO3/mqd1zZd53STJN0EjCYF9Uy2ZOBno= -github.com/nats-io/nkeys v0.4.15 h1:JACV5jRVO9V856KOapQ7x+EY8Jo3qw1vJt/9Jpwzkk4= -github.com/nats-io/nkeys v0.4.15/go.mod h1:CpMchTXC9fxA5zrMo4KpySxNjiDVvr8ANOSZdiNfUrs= +github.com/nats-io/nkeys v0.4.16 h1:rd5oAuLOb8mnAycB0xleuEBNS1pVVnN0fv/FF34Eypg= +github.com/nats-io/nkeys v0.4.16/go.mod h1:llLgWoI0o4z/Q57q2R1kHfmocyhGV6VG/U18Glg1Afs= github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw= github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c= github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno= diff --git a/internal/common/task.go b/internal/common/task.go index c12e21722e..52258eb4a0 100644 --- a/internal/common/task.go +++ b/internal/common/task.go @@ -30,4 +30,9 @@ type TaskHandle interface { GetMessage() TaskMessage Ack() error Nack() error + + // InProgress resets the AckWait timer without acknowledging the message, + // signalling the broker that the worker is still processing. Call + // periodically during long tasks to avoid in-flight redelivery. + InProgress() error } diff --git a/internal/deepdoc/parser/pdf/parser.go b/internal/deepdoc/parser/pdf/parser.go index 2204a34633..79a0017f05 100644 --- a/internal/deepdoc/parser/pdf/parser.go +++ b/internal/deepdoc/parser/pdf/parser.go @@ -187,7 +187,7 @@ func (p *Parser) processPage(ctx context.Context, engine pdf.PDFEngine, pg int, // render to an unsafe DPI and spike memory on large pages. const maxRetryZoom = 9.0 retryZoom := math.Min(p.Config.Zoom*pdf.DlaScale, maxRetryZoom) - slog.Info("per-page zoom retry", "page", pg, "zoom", retryZoom) + slog.Debug("per-page zoom retry", "page", pg, "zoom", retryZoom) retryImg, retryRenderErr := p.renderAtDPI(ctx, engine, pg, retryZoom*72) if retryRenderErr == nil && retryImg != nil { ocrBoxes, updatedChars, ocrUsed = p.processPageBoxes(ctx, retryImg, chars, pg, retryRenderErr, isScanNoise, docAnalyzer) diff --git a/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go b/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go index 1a2cec409c..0909f07986 100644 --- a/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go +++ b/internal/deepdoc/parser/pdf/parser_pipeline_integration_test.go @@ -1,4 +1,4 @@ -//go:build cgo && integration +//go:build cgo && manual package pdf @@ -372,30 +372,6 @@ func TestIntegration_Idempotency(t *testing.T) { }) } -// cropImageRect crops a rectangular region from an image. -func cropImageRect(img image.Image, x0, y0, x1, y1 int) image.Image { - b := img.Bounds() - if x0 < b.Min.X { - x0 = b.Min.X - } - if y0 < b.Min.Y { - y0 = b.Min.Y - } - if x1 > b.Max.X { - x1 = b.Max.X - } - if y1 > b.Max.Y { - y1 = b.Max.Y - } - out := image.NewRGBA(image.Rect(0, 0, x1-x0, y1-y0)) - for y := y0; y < y1; y++ { - for x := x0; x < x1; x++ { - out.Set(x-x0, y-y0, img.At(x, y)) - } - } - return out -} - const coordEpsilon = 1.0 // pixels const confEpsilon = 0.01 diff --git a/internal/deepdoc/parser/pdf/test_helpers_test.go b/internal/deepdoc/parser/pdf/test_helpers_test.go index 81ea5bcbd0..ee494db15e 100644 --- a/internal/deepdoc/parser/pdf/test_helpers_test.go +++ b/internal/deepdoc/parser/pdf/test_helpers_test.go @@ -6,6 +6,30 @@ import ( pdf "ragflow/internal/deepdoc/parser/pdf/type" ) +// cropImageRect crops a rectangular region from an image. +func cropImageRect(img image.Image, x0, y0, x1, y1 int) image.Image { + b := img.Bounds() + if x0 < b.Min.X { + x0 = b.Min.X + } + if y0 < b.Min.Y { + y0 = b.Min.Y + } + if x1 > b.Max.X { + x1 = b.Max.X + } + if y1 > b.Max.Y { + y1 = b.Max.Y + } + out := image.NewRGBA(image.Rect(0, 0, x1-x0, y1-y0)) + for y := y0; y < y1; y++ { + for x := x0; x < x1; x++ { + out.Set(x-x0, y-y0, img.At(x, y)) + } + } + return out +} + // ── testPageImg: small test image for ocrMergeChars tests ───────────── // 90×120 px at 216 DPI → 30×40 pt in PDF space after /3.0 scaling. diff --git a/internal/engine/elasticsearch/client.go b/internal/engine/elasticsearch/client.go index 87f34d29f4..01be8803bb 100644 --- a/internal/engine/elasticsearch/client.go +++ b/internal/engine/elasticsearch/client.go @@ -105,6 +105,11 @@ func (e *elasticsearchEngine) GetType() string { return "elasticsearch" } +// SupportsPageRank returns true because Elasticsearch supports pagerank. +func (e *elasticsearchEngine) SupportsPageRank() bool { + return true +} + // Ping health check func (e *elasticsearchEngine) Ping(ctx context.Context) error { req := esapi.InfoRequest{} diff --git a/internal/engine/engine.go b/internal/engine/engine.go index b3d6f8dc2c..0e4956d2bf 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -76,6 +76,9 @@ type DocEngine interface { // GetType returns the engine type GetType() string + // SupportsPageRank reports whether the engine supports dataset-level pagerank. + SupportsPageRank() bool + // FilterDocIdsByMetaPushdown runs a metadata filter directly against // the doc metadata index, returning matching doc IDs or nil if push-down // is not supported (caller should fall back to in-memory filtering). diff --git a/internal/engine/infinity/client.go b/internal/engine/infinity/client.go index 2eca460f79..ed141d300e 100644 --- a/internal/engine/infinity/client.go +++ b/internal/engine/infinity/client.go @@ -191,6 +191,11 @@ func (e *infinityEngine) GetType() string { return "infinity" } +// SupportsPageRank returns false because Infinity does not support pagerank. +func (e *infinityEngine) SupportsPageRank() bool { + return false +} + // Ping checks if Infinity is accessible func (e *infinityEngine) Ping(ctx context.Context) error { if e.client == nil || e.client.conn == nil { diff --git a/internal/engine/nats/nats.go b/internal/engine/nats/nats.go index ad4f68f144..3662a96f8b 100644 --- a/internal/engine/nats/nats.go +++ b/internal/engine/nats/nats.go @@ -253,3 +253,7 @@ func (m *NatsMessageHandle) Ack() error { func (m *NatsMessageHandle) Nack() error { return m.message.Nak() } + +func (m *NatsMessageHandle) InProgress() error { + return m.message.InProgress() +} diff --git a/internal/ingestion/component/slice3_integration_test.go b/internal/ingestion/component/slice3_integration_test.go new file mode 100644 index 0000000000..c2e494f500 --- /dev/null +++ b/internal/ingestion/component/slice3_integration_test.go @@ -0,0 +1,86 @@ +//go:build integration +// +build integration + +// +// 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. +// + +// Slice 3 integration tests (require tokenizer pool). +package component + +import ( + "context" + "testing" +) + +// TestTokenizer_FallsBackToContentWithWeight pins the python +// rag/flow/tokenizer.py:111 fallback. A chunk with only +// content_with_weight (no text) must tokenize the fallback text. +func TestTokenizer_FallsBackToContentWithWeight(t *testing.T) { + requireTokenizerPool(t) + c := &TokenizerComponent{} + c.param.SearchMethod = []string{"full_text"} + c.param.Fields = []string{"text"} + c.param.FilenameEmbdWeight = 0 + + out, err := c.Invoke(context.Background(), map[string]any{ + "name": "doc", + "output_format": "chunks", + "chunks": []map[string]any{ + {"content_with_weight": "fallback text", "doc_type_kwd": "text"}, + }, + }) + if err != nil { + t.Fatalf("Tokenizer.Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) == 0 { + t.Fatal("no chunks emitted") + } + if got := chunks[0]["content_ltks"]; got == nil { + t.Errorf("content_ltks missing; content_with_weight fallback did not run") + } else if s, ok := got.(string); !ok || s == "" { + t.Errorf("content_ltks = %v (type %T), want non-empty string", got, got) + } +} + +// TestTokenizer_DoesNotChangeChunkText pins the regression guard +// for the fallback. When both "text" and "content_with_weight" +// are present, "text" wins. +func TestTokenizer_DoesNotChangeChunkText(t *testing.T) { + requireTokenizerPool(t) + c := &TokenizerComponent{} + c.param.SearchMethod = []string{"full_text"} + c.param.Fields = []string{"text"} + c.param.FilenameEmbdWeight = 0 + + out, err := c.Invoke(context.Background(), map[string]any{ + "name": "doc", + "output_format": "chunks", + "chunks": []map[string]any{ + {"text": "primary text", "content_with_weight": "fallback", "doc_type_kwd": "text"}, + }, + }) + if err != nil { + t.Fatalf("Tokenizer.Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) == 0 { + t.Fatal("no chunks emitted") + } + if got, want := chunks[0]["text"], "primary text"; got != want { + t.Errorf("text = %v, want %v (text should win over content_with_weight)", got, want) + } +} diff --git a/internal/ingestion/component/slice3_test.go b/internal/ingestion/component/slice3_test.go index 44d05e3d19..c89ddfd0bc 100644 --- a/internal/ingestion/component/slice3_test.go +++ b/internal/ingestion/component/slice3_test.go @@ -16,78 +16,15 @@ // Slice 3 tests for port-rag-flow-pipeline-to-go.md Phase 5. // -// Pins the Tokenizer content_with_weight fallback and the -// Extractor prompt placeholder substitution. +// Pins the Extractor prompt placeholder substitution. package component import ( - "context" - pipe "ragflow/internal/ingestion/pipeline" "strings" "testing" ) -// TestTokenizer_FallsBackToContentWithWeight pins the python -// rag/flow/tokenizer.py:111 fallback. A chunk with only -// content_with_weight (no text) must tokenize the fallback text. -func TestTokenizer_FallsBackToContentWithWeight(t *testing.T) { - pipe.RequireTokenizerPool(t) - c := &TokenizerComponent{} - c.param.SearchMethod = []string{"full_text"} - c.param.Fields = []string{"text"} - c.param.FilenameEmbdWeight = 0 - - out, err := c.Invoke(context.Background(), map[string]any{ - "name": "doc", - "output_format": "chunks", - "chunks": []map[string]any{ - {"content_with_weight": "fallback text", "doc_type_kwd": "text"}, - }, - }) - if err != nil { - t.Fatalf("Tokenizer.Invoke: %v", err) - } - chunks, _ := out["chunks"].([]map[string]any) - if len(chunks) == 0 { - t.Fatal("no chunks emitted") - } - if got := chunks[0]["content_ltks"]; got == nil { - t.Errorf("content_ltks missing; content_with_weight fallback did not run") - } else if s, ok := got.(string); !ok || s == "" { - t.Errorf("content_ltks = %v (type %T), want non-empty string", got, got) - } -} - -// TestTokenizer_DoesNotChangeChunkText pins the regression guard -// for the fallback. When both "text" and "content_with_weight" -// are present, "text" wins. -func TestTokenizer_DoesNotChangeChunkText(t *testing.T) { - pipe.RequireTokenizerPool(t) - c := &TokenizerComponent{} - c.param.SearchMethod = []string{"full_text"} - c.param.Fields = []string{"text"} - c.param.FilenameEmbdWeight = 0 - - out, err := c.Invoke(context.Background(), map[string]any{ - "name": "doc", - "output_format": "chunks", - "chunks": []map[string]any{ - {"text": "primary text", "content_with_weight": "fallback", "doc_type_kwd": "text"}, - }, - }) - if err != nil { - t.Fatalf("Tokenizer.Invoke: %v", err) - } - chunks, _ := out["chunks"].([]map[string]any) - if len(chunks) == 0 { - t.Fatal("no chunks emitted") - } - if got, want := chunks[0]["text"], "primary text"; got != want { - t.Errorf("text = %v, want %v (text should win over content_with_weight)", got, want) - } -} - // TestSubstitutePromptPlaceholders_ReplacesAtChunks pins the // resume-template pattern `{TitleChunker:FlatMiceFix@chunks}`. // The substitute is the joined chunk text. diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index db860442c9..232f6ffa0d 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -655,27 +655,39 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error { return fmt.Errorf("Tokenizer: keyword tokens marshal: %w", err) } } - if s := ck.Summary; s != "" { + if s := ck.Summary; strings.TrimSpace(s) != "" { st, err := tokenizer.Tokenize(s) if err != nil { return fmt.Errorf("Tokenizer: summary tokenize: %w", err) } + if st == "" { + st = s + } ck.ContentLtks = st smt, err := tokenizer.FineGrainedTokenize(st) if err != nil { return fmt.Errorf("Tokenizer: summary fine-grain: %w", err) } + if smt == "" { + smt = st + } ck.ContentSmLtks = smt - } else if t := ck.Text; t != "" { + } else if t := ck.Text; strings.TrimSpace(t) != "" { tt, err := tokenizer.Tokenize(t) if err != nil { return fmt.Errorf("Tokenizer: text tokenize: %w", err) } + if tt == "" { + tt = t + } ck.ContentLtks = tt smt, err := tokenizer.FineGrainedTokenize(tt) if err != nil { return fmt.Errorf("Tokenizer: text fine-grain: %w", err) } + if smt == "" { + smt = tt + } ck.ContentSmLtks = smt } } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index 1f2950be4c..5123863a4a 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -23,11 +23,8 @@ import ( "bytes" "context" "errors" - "fmt" "log" "math" - "os" - "ragflow/internal/common" "reflect" "strings" "sync/atomic" @@ -39,42 +36,16 @@ import ( "ragflow/internal/tokenizer" ) -var tokenizerPoolInitErr error - -// TestMain initializes the tokenizer pool before any test runs. -// The tokenizer package needs the C++ RAGAnalyzer dictionaries -// (see internal/tokenizer.Init) for `Tokenize` / -// `FineGrainedTokenize`; without it, `tokenizeChunks` errors out -// with "tokenizer pool not initialized". Tests in other packages -// initialize the pool at startup; this package must do the same -// because the Tokenizer component touches tokenizer.Tokenize. -// -// If Init fails (e.g., dict path missing in some CI sandboxes), -// we log the failure but still run the tests. Cases that exercise -// tokenizeChunks will fail rather than skip when the pool is not -// initialized. -func TestMain(m *testing.M) { - cfg := &tokenizer.PoolConfig{ - DictPath: common.GetEnv(common.EnvRAGFlowDictPath), +func requireTokenizerPool(t *testing.T) { + t.Helper() + if err := tokenizer.Init(&tokenizer.PoolConfig{ + DictPath: "/usr/share/infinity/resource", MinSize: 1, MaxSize: 2, IdleTimeout: 30 * time.Second, AcquireTimeout: 5 * time.Second, - } - if cfg.DictPath == "" { - cfg.DictPath = "/usr/share/infinity/resource" - } - tokenizerPoolInitErr = tokenizer.Init(cfg) - if tokenizerPoolInitErr != nil { - fmt.Fprintf(os.Stderr, "tokenizer pool init failed (tests will skip tokenize-dependent cases): %v\n", tokenizerPoolInitErr) - } - os.Exit(m.Run()) -} - -func requireTokenizerPool(t *testing.T) { - t.Helper() - if tokenizerPoolInitErr != nil { - t.Skipf("tokenizer pool unavailable: %v", tokenizerPoolInitErr) + }); err != nil { + t.Skipf("tokenizer pool unavailable: %v", err) } } @@ -1016,3 +987,60 @@ func TestTokenizerComponent_InstanceResolversDoNotLeakAcrossComponents(t *testin t.Fatalf("instance resolvers leaked: vecA=%v vecB=%v", vecA, vecB) } } + +func TestValidateTokenizerOutputs_SymbolOnlyContentLtksIsEmptyFails(t *testing.T) { + // Simulates a chunk whose Text is a symbol/punctuation character that + // the C++ RAGAnalyzer tokenizer cannot produce tokens for (e.g. "·", ")", "("). + // After tokenizeChunks runs, ContentLtks and ContentSmLtks remain empty, + // and validateTokenizerOutputs must detect this as a failure. + ck := schema.ChunkDoc{ + Text: ")", + ContentLtks: "", + ContentSmLtks: "", + } + err := validateTokenizerOutputs([]schema.ChunkDoc{ck}, []string{"full_text"}, []string{"text"}) + if err == nil || !strings.Contains(err.Error(), "missing full_text tokens") { + t.Fatalf("err = %v, want missing full_text tokens", err) + } +} + +func TestTokenizeChunks_SymbolOnlyTextFallsBackToRawText(t *testing.T) { + requireTokenizerPool(t) + chunks := []schema.ChunkDoc{ + {Text: "·"}, // middle dot · — seen in production chunk[15] + {Text: ")"}, + {Text: "("}, + {Text: "*"}, + } + err := tokenizeChunks(chunks, "test") + if err != nil { + t.Fatalf("tokenizeChunks: %v", err) + } + for i, ck := range chunks { + t.Logf("chunk[%d]: text=%q content_ltks=%q content_sm_ltks=%q", + i, ck.Text, ck.ContentLtks, ck.ContentSmLtks) + // After fix: Tokenize returns empty for symbol-only text, + // but the fallback sets ContentLtks = raw text. + if strings.TrimSpace(ck.ContentLtks) == "" { + t.Errorf("chunk[%d]: expected non-empty ContentLtks (raw text fallback) for %q, got empty", + i, ck.Text) + } + } +} + +func TestTokenizeChunks_WhitespaceSummaryShadowsTextBug(t *testing.T) { + requireTokenizerPool(t) + chunks := []schema.ChunkDoc{ + {Summary: " ", Text: "real content here"}, + } + err := tokenizeChunks(chunks, "test") + if err != nil { + t.Fatalf("tokenizeChunks: %v", err) + } + // After fix: TrimSpace(" ") is empty, so the Summary branch is skipped. + // The Text branch is entered and "real content here" is tokenized normally. + if strings.TrimSpace(chunks[0].ContentLtks) == "" { + t.Errorf("whitespace Summary should be skipped, Text %q should be tokenized, but ContentLtks is empty", + chunks[0].Text) + } +} diff --git a/internal/ingestion/pipeline/errors.go b/internal/ingestion/pipeline/errors.go index d90b9d2b31..ad7ea4beab 100644 --- a/internal/ingestion/pipeline/errors.go +++ b/internal/ingestion/pipeline/errors.go @@ -19,10 +19,7 @@ package pipeline import "errors" var ( - errNilDSL = errors.New("pipeline: nil DSL") - errEmptyStages = errors.New("pipeline: DSL has no components") - errUnknownComponent = errors.New("pipeline: unknown component") - errUnknownStage = errors.New("pipeline: unknown stage") + errNilDSL = errors.New("pipeline: nil DSL") ) type stageError struct { @@ -33,11 +30,3 @@ type stageError struct { func (e *stageError) Error() string { return "pipeline: stage " + e.Stage + ": " + e.Reason } - -type sinkError struct { - Reason string -} - -func (e *sinkError) Error() string { - return "pipeline: sink: " + e.Reason -} diff --git a/internal/ingestion/pipeline/test_helpers.go b/internal/ingestion/pipeline/helpers_test.go similarity index 77% rename from internal/ingestion/pipeline/test_helpers.go rename to internal/ingestion/pipeline/helpers_test.go index 56834fa490..300cc67606 100644 --- a/internal/ingestion/pipeline/test_helpers.go +++ b/internal/ingestion/pipeline/helpers_test.go @@ -18,9 +18,7 @@ import ( "encoding/json" "math" "path/filepath" - "ragflow/internal/common" goruntime "runtime" - "sort" "strings" "testing" "time" @@ -39,61 +37,24 @@ func repoRootFromPipelineTest(t *testing.T) string { func RequireTokenizerPool(t *testing.T) { t.Helper() - if tokenizer.IsInitialized() { - return - } - cfg := &tokenizer.PoolConfig{ - DictPath: common.GetEnv(common.EnvRAGFlowDictPath), + if err := tokenizer.Init(&tokenizer.PoolConfig{ + DictPath: "/usr/share/infinity/resource", MinSize: 1, MaxSize: 2, IdleTimeout: 30 * time.Second, AcquireTimeout: 5 * time.Second, - } - if cfg.DictPath == "" { - cfg.DictPath = "/usr/share/infinity/resource" - } - if err := tokenizer.Init(cfg); err != nil { + }); err != nil { t.Skipf("tokenizer pool init failed: %v", err) } } func terminalComponentIDsFromTemplate(t *testing.T, raw []byte) []string { t.Helper() - var tpl map[string]any - if err := json.Unmarshal(raw, &tpl); err != nil { - t.Fatalf("unmarshal template: %v", err) + ids, err := TerminalComponentIDs(raw) + if err != nil { + t.Fatal(err) } - dsl, ok := tpl["dsl"].(map[string]any) - if !ok { - t.Fatalf("template dsl = %T, want map[string]any", tpl["dsl"]) - } - components, ok := dsl["components"].(map[string]any) - if !ok { - t.Fatalf("template components = %T, want map[string]any", dsl["components"]) - } - var terminals []string - for id, rawComp := range components { - comp, ok := rawComp.(map[string]any) - if !ok { - t.Fatalf("component %q = %T, want map[string]any", id, rawComp) - } - switch ds := comp["downstream"].(type) { - case nil: - terminals = append(terminals, id) - case []any: - if len(ds) == 0 { - terminals = append(terminals, id) - } - case []string: - if len(ds) == 0 { - terminals = append(terminals, id) - } - default: - t.Fatalf("component %q downstream = %T, want []any/[]string/nil", id, comp["downstream"]) - } - } - sort.Strings(terminals) - return terminals + return ids } func terminalPayloadFromRunOutput(t *testing.T, out map[string]any, terminalID string) map[string]any { diff --git a/internal/ingestion/pipeline/payload.go b/internal/ingestion/pipeline/payload.go new file mode 100644 index 0000000000..d9e8cd7371 --- /dev/null +++ b/internal/ingestion/pipeline/payload.go @@ -0,0 +1,84 @@ +// +// 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 pipeline + +import ( + "encoding/json" + "fmt" + "sort" +) + +// ExtractPayload extracts the terminal component's output from a pipeline run +// result. When the output carries an "output_format" key the whole map is +// returned as-is; otherwise the payload is looked up from the run output keyed +// by the DSL's single terminal component id. +func ExtractPayload(dsl string, out map[string]any) (map[string]any, error) { + if out == nil { + return nil, nil + } + if _, ok := out["output_format"]; ok { + return out, nil + } + terminalIDs, err := TerminalComponentIDs([]byte(dsl)) + if err != nil { + return nil, err + } + if len(terminalIDs) != 1 { + return nil, fmt.Errorf("pipeline requires exactly 1 terminal, got %d: %v", len(terminalIDs), terminalIDs) + } + payload, ok := out[terminalIDs[0]].(map[string]any) + if !ok { + return nil, fmt.Errorf("run output missing terminal payload %q", terminalIDs[0]) + } + return payload, nil +} + +// TerminalComponentIDs walks a raw DSL JSON and returns the sorted ids of +// components that have no downstream connections (terminals / sinks). +func TerminalComponentIDs(raw []byte) ([]string, error) { + var tpl map[string]any + if err := json.Unmarshal(raw, &tpl); err != nil { + return nil, fmt.Errorf("unmarshal canvas dsl: %w", err) + } + root := tpl + if nested, ok := tpl["dsl"].(map[string]any); ok { + root = nested + } + components, ok := root["components"].(map[string]any) + if !ok { + return nil, fmt.Errorf("canvas dsl missing components map") + } + terminals := make([]string, 0, len(components)) + for id, rawComp := range components { + comp, ok := rawComp.(map[string]any) + if !ok { + return nil, fmt.Errorf("component %q has invalid type %T", id, rawComp) + } + switch downstream := comp["downstream"].(type) { + case nil: + terminals = append(terminals, id) + case []any: + if len(downstream) == 0 { + terminals = append(terminals, id) + } + default: + // Non-slice downstream means the component is connected; ignore it here. + } + } + sort.Strings(terminals) + return terminals, nil +} diff --git a/internal/ingestion/pipeline/payload_test.go b/internal/ingestion/pipeline/payload_test.go new file mode 100644 index 0000000000..964ae94e72 --- /dev/null +++ b/internal/ingestion/pipeline/payload_test.go @@ -0,0 +1,198 @@ +// +// 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 pipeline + +import ( + "strings" + "testing" +) + +func TestExtractPayload(t *testing.T) { + tests := []struct { + name string + dsl string + output map[string]any + wantOK bool + wantErr string + want map[string]any + }{ + { + name: "output_format key present → return output as-is", + dsl: `{}`, + output: map[string]any{"output_format": "x", "data": 1}, + wantOK: true, + want: map[string]any{"output_format": "x", "data": 1}, + }, + { + name: "nil output → nil, nil", + dsl: `{}`, + wantOK: true, + }, + { + name: "single terminal → extract its payload", + dsl: `{"components":{"c1":{"downstream":null},"begin":{"downstream":["c1"]}}}`, + output: map[string]any{"c1": map[string]any{"text": "hi"}, "begin": map[string]any{}}, + wantOK: true, + want: map[string]any{"text": "hi"}, + }, + { + name: "zero terminals → error", + dsl: `{"components":{}}`, + output: map[string]any{}, + wantErr: "requires exactly 1 terminal, got 0", + }, + { + name: "multiple terminals → error", + dsl: `{"components":{"c1":{"downstream":null},"c2":{"downstream":null}}}`, + output: map[string]any{"c1": map[string]any{}, "c2": map[string]any{}}, + wantErr: "requires exactly 1 terminal, got 2", + }, + { + name: "nested dsl template → unwrap and extract", + dsl: `{"id":"t1","dsl":{"components":{"c1":{"downstream":null}},"path":["c1"],"graph":{"nodes":[]}}}`, + output: map[string]any{"c1": map[string]any{"text": "nested"}}, + wantOK: true, + want: map[string]any{"text": "nested"}, + }, + { + name: "DSL missing components key", + dsl: `{}`, + output: map[string]any{}, + wantErr: "missing components map", + }, + { + name: "terminal payload is not a map", + dsl: `{"components":{"c1":{"downstream":null}}}`, + output: map[string]any{"c1": "not-a-map"}, + wantErr: "missing terminal payload", + }, + { + name: "output_format takes priority over terminal extraction", + dsl: `{"components":{"c1":{"downstream":null}}}`, + output: map[string]any{"output_format": "markdown", "text": "content"}, + wantOK: true, + want: map[string]any{"output_format": "markdown", "text": "content"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ExtractPayload(tt.dsl, tt.output) + if tt.wantErr != "" { + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %q, want containing %q", err.Error(), tt.wantErr) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !tt.wantOK { + return + } + if tt.want == nil && got != nil { + t.Fatalf("got = %v, want nil", got) + } + if tt.want != nil { + for k, v := range tt.want { + if got[k] != v { + t.Fatalf("got[%q] = %v, want %v", k, got[k], v) + } + } + } + }) + } +} + +func TestTerminalComponentIDs(t *testing.T) { + tests := []struct { + name string + raw []byte + want []string + wantErr string + }{ + { + name: "single terminal", + raw: []byte(`{"components":{"c1":{"downstream":null}}}`), + want: []string{"c1"}, + }, + { + name: "multiple terminals sorted", + raw: []byte(`{"components":{"c2":{"downstream":null},"c1":{"downstream":null}}}`), + want: []string{"c1", "c2"}, + }, + { + name: "non-terminal with downstream connections excluded", + raw: []byte(`{"components":{"c1":{"downstream":["c2"]},"c2":{"downstream":null}}}`), + want: []string{"c2"}, + }, + { + name: "empty downstream array counts as terminal", + raw: []byte(`{"components":{"c1":{"downstream":[]}}}`), + want: []string{"c1"}, + }, + { + name: "nested dsl template unwrapped", + raw: []byte(`{"dsl":{"components":{"c1":{"downstream":null}}}}`), + want: []string{"c1"}, + }, + { + name: "no components key", + raw: []byte(`{}`), + wantErr: "missing components map", + }, + { + name: "no terminals", + raw: []byte(`{"components":{}}`), + want: []string{}, + }, + { + name: "invalid JSON", + raw: []byte(`{invalid`), + wantErr: "unmarshal", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := TerminalComponentIDs(tt.raw) + if tt.wantErr != "" { + if err == nil { + t.Fatalf("expected error containing %q, got nil", tt.wantErr) + } + if !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %q, want containing %q", err.Error(), tt.wantErr) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(got) != len(tt.want) { + t.Fatalf("got %v (%d), want %v (%d)", got, len(got), tt.want, len(tt.want)) + } + for i := range tt.want { + if got[i] != tt.want[i] { + t.Fatalf("got[%d] = %q, want %q", i, got[i], tt.want[i]) + } + } + }) + } +} diff --git a/internal/ingestion/pipeline/pipeline.go b/internal/ingestion/pipeline/pipeline.go index 818a0c1b75..4a7695ba51 100644 --- a/internal/ingestion/pipeline/pipeline.go +++ b/internal/ingestion/pipeline/pipeline.go @@ -28,9 +28,7 @@ import ( _ "ragflow/internal/agent/component" "ragflow/internal/agent/runtime" "ragflow/internal/common" - "ragflow/internal/dao" redis2 "ragflow/internal/engine/redis" - "ragflow/internal/entity" "ragflow/internal/ingestion/component/globals" "github.com/cloudwego/eino/compose" @@ -53,6 +51,7 @@ type Pipeline struct { // distinguishable error so the caller knows resume is unavailable. requireResume bool factory runtime.ComponentFactory // optional instance-scoped component factory + sink ProgressSink // optional progress sink; nil -> drop events (DB-independent) } // ErrResumeUnavailable is returned by Run when WithRequireResume is set but no @@ -99,6 +98,38 @@ func WithDocumentID(docID string) PipelineOption { return func(p *Pipeline) { p.documentID = docID } } +// ProgressEvent is a structured component lifecycle event emitted by the +// pipeline to a ProgressSink. The pipeline fills every field - including the +// ingestion-proprietary Index (from its componentIndexMap) and the Total +// denominator - so the sink needs no canvas knowledge to persist it. +type ProgressEvent struct { + TaskID string + DocumentID string + Component string + Message string + Index int + Phase int + Total int +} + +// ProgressSink receives pipeline progress for durable persistence. It is the +// single channel through which the pipeline reports component lifecycle +// events and the component-total denominator; the pipeline itself never +// touches the DAO layer. Implementations live in the orchestration layer +// (internal/ingestion/service). A nil sink is valid: events are dropped and +// the pipeline stays DB-independent (unit tests, headless runs). +type ProgressSink interface { + OnComponentTotal(taskID string, total int) + OnComponentProgress(ev ProgressEvent) +} + +// WithProgressSink injects a sink that receives component progress events +// and the component-total denominator. When unset, the pipeline drops +// progress events and stays DB-independent. +func WithProgressSink(s ProgressSink) PipelineOption { + return func(p *Pipeline) { p.sink = s } +} + // NewPipelineFromDSL compiles the canonical ingestion canvas DSL. // It accepts either the inner canvas DSL or the template wrapper whose // top-level `dsl` field carries that canvas. @@ -174,12 +205,6 @@ func cloneMapOrEmpty(m map[string]any) map[string]any { return out } -func stageTimeout() time.Duration { - return defaultStageTimeout -} - -var defaultStageTimeout = 60 * time.Second - // defaultCheckpointTTL is the expiry applied to the eino checkpoint payload // and the RunTracker hash. A finished run's checkpoint is deleted on success; // the TTL only guards against leaks from crashed runs that never clean up. @@ -256,16 +281,12 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, setups ...map // Record the component count as the authoritative denominator for // progress percentage. Best-effort: a DB failure (or headless run // with no DB) must not abort the pipeline — progress is observability. - if dao.DB != nil { - total := len(p.canvas.Components) - if err := dao.NewIngestionTaskDAO().UpdateComponentTotal(p.taskID, total); err != nil { - common.Error(fmt.Sprintf("pipeline: update component_total for task %s failed: %v", p.taskID, err), err) - } + if p.sink != nil { + p.sink.OnComponentTotal(p.taskID, len(p.canvas.Components)) } runState := canvas.NewCanvasState("", p.taskID) runCtx := canvas.WithState(ctx, runState) - runCtx = canvas.WithComponentTimeoutOverride(runCtx, stageTimeout()) // Framework-level progress fan-out: the canvas framework // (realComponentBody) pulls this callback from ctx via // runtime.ProgressCallbackFromContext and records every component @@ -438,22 +459,11 @@ func finalizeResult(current, out map[string]any, runState *canvas.CanvasState) m return merged } -// taskLogProgressCallback returns a runtime.ProgressCallback that appends -// an ingestion_task_log row for every component progress event emitted by -// the canvas framework (component start = 0, done = 1, fail = -1). Progress -// is a framework-level concern owned by realComponentBody; this callback is -// the ingestion-specific sink that turns those events into durable task -// logs that the API can later read back (dao.IngestionTaskLogDAO). -// -// The write is best-effort: a DB failure is logged and never aborts the -// pipeline. When the DAO/DB has not been initialized (unit tests, headless -// runs) the callback is nil, so TrackProgress becomes a no-op and the -// pipeline stays DB-independent. // componentIndexMap builds a deterministic cpnID → 0-based-index map for // the task's canvas. Map iteration order is non-deterministic in Go, so // the cpnIDs are sorted to keep the index stable across runs. The index is -// ingestion-proprietary (plan §5.1 M1) and computed here at the sink, -// never carried on the shared runtime.ProgressEvent. +// ingestion-proprietary and computed here, then carried on the pipeline-local +// ProgressEvent so the sink needs no canvas knowledge. func (p *Pipeline) componentIndexMap() map[string]int { ids := make([]string, 0, len(p.canvas.Components)) for id := range p.canvas.Components { @@ -467,33 +477,18 @@ func (p *Pipeline) componentIndexMap() map[string]int { return m } -// taskLogProgressCallback returns a runtime.ProgressCallback that appends -// an ingestion_task_log row for every component progress event emitted by -// the canvas framework. Progress is a framework-level concern owned by -// realComponentBody; this callback is the ingestion-specific sink that -// turns those events into durable task logs the API can read back -// (dao.IngestionTaskLogDAO). -// -// The event carries the node id (cpnID) and phase. The sink derives: -// - ComponentIndex from the task's ComponentIndexMap (ingestion-only); -// - Message from the phase + cpnID (+ error), mirroring the legacy -// "cpnID Started/Done" / "cpnID: " strings the frontend expects. -// -// The write is best-effort: a DB failure is logged and never aborts the -// pipeline. When the DAO/DB has not been initialized (unit tests, headless -// runs) the callback is nil, so TrackProgress becomes a no-op and the -// pipeline stays DB-independent. +// taskLogProgressCallback returns a runtime.ProgressCallback that forwards +// every component lifecycle event (start/done/fail) to the pipeline's +// ProgressSink. The sink owns all persistence; this callback only shapes the +// event - deriving the message string the frontend expects and the +// ingestion-proprietary component index - so the pipeline never touches the +// DAO layer. Returns nil when no sink is attached, leaving TrackProgress a +// no-op and the pipeline DB-independent (unit tests, headless runs). func (p *Pipeline) taskLogProgressCallback() runtime.ProgressCallback { - if dao.DB == nil { + if p.sink == nil { return nil } - logDAO := dao.NewIngestionTaskLogDAO() - docDAO := dao.NewDocumentDAO() indexMap := p.componentIndexMap() - // total is the authoritative denominator (plan §8). After the - // UserFillUp/legacy-no-op guard in Compile, every component in the canvas - // reports progress, so len(Components) == component_total and the - // aggregate percent can reach 100%. total := len(p.canvas.Components) return func(ev runtime.ProgressEvent) { var msg string @@ -509,54 +504,14 @@ func (p *Pipeline) taskLogProgressCallback() runtime.ProgressCallback { msg = ev.Component + " Error" } } - entry := &entity.IngestionTaskLog{ - TaskID: p.taskID, - Checkpoint: entity.JSONMap{}, - ComponentIndex: indexMap[ev.Component], - Phase: int(ev.Phase), - Component: ev.Component, - Message: msg, - } - if err := logDAO.Create(entry); err != nil { - common.Error(fmt.Sprintf("pipeline: write ingestion_task_log for task %s failed: %v", p.taskID, err), err) - } - // Mirror progress into the document table so the existing - // GET /api/v1/datasets/{dataset_id}/documents endpoint (which reads - // document.progress/run/progress_msg) reflects live Go pipeline - // progress. Best-effort: a DB failure is logged and never aborts the - // pipeline. Skipped when no owning document is bound (headless runs). - if p.documentID == "" { - return - } - progress, run := p.deriveDocumentProgress(logDAO, total) - updates := map[string]interface{}{ - "progress": progress, - "run": run, - "progress_msg": msg, - } - if err := docDAO.UpdateByID(p.documentID, updates); err != nil { - common.Error(fmt.Sprintf("pipeline: mirror progress to document %s for task %s failed: %v", p.documentID, p.taskID, err), err) - } + p.sink.OnComponentProgress(ProgressEvent{ + TaskID: p.taskID, + DocumentID: p.documentID, + Component: ev.Component, + Message: msg, + Index: indexMap[ev.Component], + Phase: int(ev.Phase), + Total: total, + }) } } - -// deriveDocumentProgress computes the document-level progress (0~1) and run -// label ("0".."4", matching Python's document.run enum) from the aggregated -// ingestion_task_log. It reads the freshly-written log row back so the value -// is correct across resume rounds and crash recovery (plan §8). -func (p *Pipeline) deriveDocumentProgress(logDAO *dao.IngestionTaskLogDAO, total int) (float64, string) { - agg, err := logDAO.AggregateProgress(p.taskID, total) - if err != nil || agg == nil || total <= 0 { - return 0, "0" - } - run := "0" // UNSTART - switch { - case agg.Failed > 0: - run = "4" // FAIL - case agg.Done == total: - run = "3" // DONE - case agg.Done > 0 || agg.Running > 0: - run = "1" // RUNNING - } - return agg.Percent / 100, run -} diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index 1863d20aa6..71dc101ee7 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -22,8 +22,13 @@ import ( "fmt" "sync" "testing" + "time" + "ragflow/internal/agent/canvas" "ragflow/internal/agent/runtime" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" ) type mockCanvasStage struct { @@ -169,8 +174,9 @@ func (s *factorySentinelStage) Invoke(_ context.Context, inputs map[string]any) // memCheckpointStore is a thread-safe in-memory canvas.CheckPointStore used // to exercise the resumable run path without Redis. type memCheckpointStore struct { - mu sync.Mutex - data map[string][]byte + mu sync.Mutex + data map[string][]byte + deleted int // number of times Delete was called } func newMemCheckpointStore() *memCheckpointStore { @@ -197,9 +203,16 @@ func (s *memCheckpointStore) Delete(_ context.Context, id string) error { s.mu.Lock() defer s.mu.Unlock() delete(s.data, id) + s.deleted++ return nil } +func (s *memCheckpointStore) deleteCount() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.deleted +} + // TestPipelineRun_InstanceFactoryOverridesDefaultFactory verifies that a // pipeline-scoped component factory can provide task-specific components. func TestPipelineRun_InstanceFactoryOverridesDefaultFactory(t *testing.T) { @@ -376,3 +389,191 @@ func TestPipelineRun_RequireResumeRejectsWithoutStore(t *testing.T) { t.Fatalf("expected ErrResumeUnavailable, got %v", err) } } + +// recordingSink captures OnComponentTotal / OnComponentProgress calls so tests +// can assert the pipeline forwards progress to the sink instead of writing +// the DAO layer directly. +type recordingSink struct { + mu sync.Mutex + total int + totalSet bool + events []ProgressEvent +} + +func (r *recordingSink) OnComponentTotal(taskID string, total int) { + r.mu.Lock() + defer r.mu.Unlock() + r.total = total + r.totalSet = true +} + +func (r *recordingSink) OnComponentProgress(ev ProgressEvent) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, ev) +} + +// TestPipelineRunForwardsProgressToSink verifies the pipeline reports the +// component-total denominator and each component lifecycle event to the +// injected ProgressSink, and carries task/document/total context on every +// event so the sink needs no canvas knowledge. +func TestPipelineRunForwardsProgressToSink(t *testing.T) { + stageA := &mockCanvasStage{output: map[string]any{"a": 1}} + stageB := &mockCanvasStage{output: map[string]any{"b": 2}} + const ( + nameA = "p.SinkStageA" + nameB = "p.SinkStageB" + ) + runtime.MustRegister(nameA, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return stageA, nil }, + runtime.Metadata{Version: "1.0.0"}) + runtime.MustRegister(nameB, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return stageB, nil }, + runtime.Metadata{Version: "1.0.0"}) + + sink := &recordingSink{} + pipe, err := NewPipelineFromDSL([]byte(`{ + "dsl": { + "components": { + "begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]}, + "a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]}, + "b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]} + }, + "path": ["begin", "a", "b"], + "graph": {"nodes": []} + } + }`), "task-sink", WithProgressSink(sink), WithDocumentID("doc-sink")) + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + if _, err := pipe.Run(context.Background(), map[string]any{"name": "doc-sink"}); err != nil { + t.Fatalf("Run: %v", err) + } + + sink.mu.Lock() + defer sink.mu.Unlock() + if !sink.totalSet || sink.total != 3 { + t.Fatalf("OnComponentTotal = (%d, set=%v), want 3", sink.total, sink.totalSet) + } + if len(sink.events) == 0 { + t.Fatal("expected progress events, got none") + } + seen := map[string]bool{} + for _, ev := range sink.events { + if ev.TaskID != "task-sink" { + t.Fatalf("event TaskID = %q, want task-sink", ev.TaskID) + } + if ev.DocumentID != "doc-sink" { + t.Fatalf("event DocumentID = %q, want doc-sink", ev.DocumentID) + } + if ev.Total != 3 { + t.Fatalf("event Total = %d, want 3", ev.Total) + } + seen[ev.Component] = true + } + for _, want := range []string{"a", "b"} { + if !seen[want] { + t.Fatalf("expected progress event for component %q, seen=%v", want, seen) + } + } +} + +// ============================================================================= +// cleanupCheckpoint — direct unit test +// ============================================================================= + +func TestCleanupCheckpoint_DeletesStoreAndClearsTracker(t *testing.T) { + mr := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { client.Close() }) + + store := newMemCheckpointStore() + if err := store.Set(context.Background(), "cp-1", []byte("data")); err != nil { + t.Fatalf("store.Set: %v", err) + } + tracker := canvas.NewRunTrackerWithClient(client, time.Hour) + if err := tracker.AttachInterrupt(context.Background(), "cp-1", "interrupt-1"); err != nil { + t.Fatalf("AttachInterrupt: %v", err) + } + + p := &Pipeline{} + p.cleanupCheckpoint(context.Background(), store, tracker, "cp-1") + + if store.deleteCount() != 1 { + t.Fatalf("store.Delete was not called") + } + id, ok, err := tracker.GetInterruptID(context.Background(), "cp-1") + if err != nil { + t.Fatalf("GetInterruptID: %v", err) + } + if ok && id != "" { + t.Fatalf("interrupt id should be cleared, got %q", id) + } +} + +// ============================================================================= +// runPlain — tracker integration with miniredis +// ============================================================================= + +func TestRunPlain_WithTracker_Success(t *testing.T) { + stage := &mockCanvasStage{output: map[string]any{"result": "ok"}} + const name = "p.RunPlainSuccess" + runtime.MustRegister(name, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return stage, nil }, + runtime.Metadata{Version: "1.0.0"}) + + mr := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { client.Close() }) + tracker := canvas.NewRunTrackerWithClient(client, time.Hour) + + pipe, err := NewPipelineFromDSL([]byte(`{ + "dsl": { + "components": { + "begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]}, + "a": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]} + }, + "path": ["begin", "a"], + "graph": {"nodes": []} + } + }`), "task-tracker-ok", WithRunTracker(tracker)) + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + + _, err = pipe.Run(context.Background(), map[string]any{"name": "doc"}) + if err != nil { + t.Fatalf("Run: %v", err) + } +} + +func TestRunPlain_WithTracker_Error(t *testing.T) { + const name = "p.RunPlainErr" + runtime.MustRegister(name, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return &errCanvasStage{}, nil }, + runtime.Metadata{Version: "1.0.0"}) + + mr := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { client.Close() }) + tracker := canvas.NewRunTrackerWithClient(client, time.Hour) + + pipe, err := NewPipelineFromDSL([]byte(`{ + "dsl": { + "components": { + "begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["err"]}, + "err": {"obj": {"component_name": "`+name+`", "params": {}}, "upstream": ["begin"]} + }, + "path": ["begin", "err"], + "graph": {"nodes": []} + } + }`), "task-tracker-err", WithRunTracker(tracker)) + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + + _, err = pipe.Run(context.Background(), map[string]any{"name": "doc"}) + if err == nil { + t.Fatal("expected stage error, got nil") + } +} diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index dfb5dc966f..274e2c1426 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -1,6 +1,3 @@ -//go:build embedding -// +build embedding - // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,7 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -// TODO: This test should be part of unit test once "file name" issue for embedding is fixed. +// These tests exercise the real File -> Parser -> Chunker -> Tokenizer chain +// end-to-end. They double as the regression test for run-level `name` +// propagation: the Tokenizer reads `name` from CanvasState.Globals (via +// globals.GlobalOrInput), so the filename resolved by File reaches the +// Tokenizer and title weighting runs - the q__vec first element is +// 0.1*len(name) + 0.9*len(content), and embedding_token_consumption includes +// the one title encode (len(name)) plus the per-chunk content encodes. package pipeline import ( @@ -31,6 +34,7 @@ import ( "testing" "ragflow/internal/agent/runtime" + "ragflow/internal/common" componentpkg "ragflow/internal/ingestion/component" _ "ragflow/internal/ingestion/component/chunker" "ragflow/internal/storage" @@ -112,14 +116,14 @@ func TestPipelineRun_TemplateGeneral_RealComponents(t *testing.T) { t.Fatalf("chunks[%d].content_sm_ltks missing or empty: %v", i, chunks[i]["content_sm_ltks"]) } vec := floatSliceFromAny(t, chunks[i]["q_4_vec"]) - wantFirst := expectedFixedEmbedderFirst("", wantText) + wantFirst := expectedFixedEmbedderFirst(filename, wantText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[%d].q_4_vec = %v, want first=%v", i, vec, wantFirst) } totalTokens += len(wantText) } - if got := payload["embedding_token_consumption"]; got != totalTokens { - t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens) + if got := payload["embedding_token_consumption"]; got != totalTokens+len(filename) { + t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens+len(filename)) } state := stateFromRunOutput(t, out) @@ -182,7 +186,7 @@ func TestPipelineRun_TemplateOne_RealComponents(t *testing.T) { wantTexts := []string{"Alpha paragraph.", "Beta paragraph."} wantMergedText := "Alpha paragraph.\nBeta paragraph." - assertTokenizerTerminalChunk(t, payload, "", wantMergedText) + assertTokenizerTerminalChunk(t, payload, filename, wantMergedText) state := stateFromRunOutput(t, out) fileState, ok := state["File"] @@ -284,11 +288,11 @@ func TestPipelineRun_TemplateOne_RealComponents_PDFDeepdocChunking(t *testing.T) } vec := floatSliceFromAny(t, chunks[0]["q_4_vec"]) trimmedChunkText := strings.TrimSpace(chunkText) - wantFirst := expectedFixedEmbedderFirst("", trimmedChunkText) + wantFirst := expectedFixedEmbedderFirst(fixture.Name, trimmedChunkText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[0].q_4_vec = %v, want first=%v", vec, wantFirst) } - wantTokens := len(trimmedChunkText) + wantTokens := len(fixture.Name) + len(trimmedChunkText) if got := payload["embedding_token_consumption"]; got != wantTokens { t.Fatalf("embedding_token_consumption = %v, want %d", got, wantTokens) } @@ -388,14 +392,14 @@ func TestPipelineRun_TemplateManual_RealComponents(t *testing.T) { } vec := floatSliceFromAny(t, chunks[i]["q_4_vec"]) wantEmbedText := strings.TrimSpace(wantText) - wantFirst := expectedFixedEmbedderFirst("", wantEmbedText) + wantFirst := expectedFixedEmbedderFirst(filename, wantEmbedText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[%d].q_4_vec = %v, want first=%v", i, vec, wantFirst) } totalTokens += len(wantEmbedText) } - if got := payload["embedding_token_consumption"]; got != totalTokens { - t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens) + if got := payload["embedding_token_consumption"]; got != totalTokens+len(filename) { + t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens+len(filename)) } state := stateFromRunOutput(t, out) @@ -484,14 +488,14 @@ func TestPipelineRun_TemplateLaws_RealComponents(t *testing.T) { } vec := floatSliceFromAny(t, chunks[i]["q_4_vec"]) wantEmbedText := strings.TrimSpace(wantText) - wantFirst := expectedFixedEmbedderFirst("", wantEmbedText) + wantFirst := expectedFixedEmbedderFirst(filename, wantEmbedText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[%d].q_4_vec = %v, want first=%v", i, vec, wantFirst) } totalTokens += len(wantEmbedText) } - if got := payload["embedding_token_consumption"]; got != totalTokens { - t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens) + if got := payload["embedding_token_consumption"]; got != totalTokens+len(filename) { + t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens+len(filename)) } state := stateFromRunOutput(t, out) @@ -564,14 +568,14 @@ func TestPipelineRun_TemplatePaper_RealComponents(t *testing.T) { } wantEmbedText := strings.TrimSpace(wantText) vec := floatSliceFromAny(t, chunks[i]["q_4_vec"]) - wantFirst := expectedFixedEmbedderFirst("", wantEmbedText) + wantFirst := expectedFixedEmbedderFirst(filename, wantEmbedText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[%d].q_4_vec = %v, want first=%v", i, vec, wantFirst) } totalTokens += len(wantEmbedText) } - if got := payload["embedding_token_consumption"]; got != totalTokens { - t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens) + if got := payload["embedding_token_consumption"]; got != totalTokens+len(filename) { + t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens+len(filename)) } state := stateFromRunOutput(t, out) @@ -649,14 +653,14 @@ func TestPipelineRun_TemplateBook_RealComponents(t *testing.T) { } wantEmbedText := strings.TrimSpace(wantText) vec := floatSliceFromAny(t, chunks[i]["q_4_vec"]) - wantFirst := expectedFixedEmbedderFirst("", wantEmbedText) + wantFirst := expectedFixedEmbedderFirst(filename, wantEmbedText) if len(vec) != 4 || !approxFloat(vec[0], wantFirst) { t.Fatalf("chunks[%d].q_4_vec = %v, want first=%v", i, vec, wantFirst) } totalTokens += len(wantEmbedText) } - if got := payload["embedding_token_consumption"]; got != totalTokens { - t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens) + if got := payload["embedding_token_consumption"]; got != totalTokens+len(filename) { + t.Fatalf("embedding_token_consumption = %v, want %d", got, totalTokens+len(filename)) } state := stateFromRunOutput(t, out) @@ -678,7 +682,7 @@ func TestPipelineRun_TemplateBook_RealComponents(t *testing.T) { func TestPipelineRun_TemplateResume_RealComponents(t *testing.T) { RequireTokenizerPool(t) apiKey := common.GetEnv(common.EnvOpenAIApiKey) - baseURL := common.GetEnv(common.EnvOpenAIBaseUrl) + baseURL := common.GetEnv(common.EnvOpenAIBaseURL) model := common.GetEnv(common.EnvOpenAIModel) if apiKey == "" || baseURL == "" || model == "" { t.Skip("missing required env (OPENAI_API_KEY/OPENAI_BASE_URL/OPENAI_MODEL); skipping real resume extractor integration test") @@ -1102,7 +1106,7 @@ func assertTokenizerTerminalChunk(t *testing.T, payload map[string]any, name, wa if got := vec[0]; !approxFloat(got, wantFirst) { t.Fatalf("chunks[0].q_4_vec[0] = %v, want %v", got, wantFirst) } - wantTokens := len(wantMergedText) + wantTokens := len(name) + len(wantMergedText) if got := payload["embedding_token_consumption"]; got != wantTokens { t.Fatalf("embedding_token_consumption = %v, want %d", got, wantTokens) } diff --git a/internal/ingestion/service/doc_state.go b/internal/ingestion/service/doc_state.go new file mode 100644 index 0000000000..8ac39ca6d2 --- /dev/null +++ b/internal/ingestion/service/doc_state.go @@ -0,0 +1,80 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "fmt" + + "ragflow/internal/common" + taskpkg "ragflow/internal/ingestion/task" + servicepkg "ragflow/internal/service" +) + +// docStateSvc is the subset of *service.DocumentService needed to finalize a +// pipeline run's effect on document state. Extracted as an interface so tests +// can inject a stub without constructing a real DocumentService (which depends +// on initialized server config). +type docStateSvc interface { + GetDocumentMetadataByID(docID string) (map[string]any, error) + SetDocumentMetadata(docID string, meta map[string]any) error + IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error +} + +// docStateUpdater applies a pipeline run's results to document state: it +// merges the pipeline-produced metadata (filling only keys not already present) +// and bumps the document/dataset chunk and token counters. Both steps are +// best-effort; failures are logged and do not fail the task. +type docStateUpdater struct { + docSvc docStateSvc +} + +// newDocStateUpdater creates a docStateUpdater with the real DocumentService +// injected at construction time. Tests inject stubs via the docSvc field. +func newDocStateUpdater() *docStateUpdater { + return &docStateUpdater{ + docSvc: servicepkg.NewDocumentService(), + } +} + +func (u *docStateUpdater) apply(r *taskpkg.PipelineResult) { + if r == nil { + return + } + if len(r.Metadata) > 0 { + if err := mergeDocMetadata(u.docSvc, r.DocID, r.Metadata); err != nil { + common.Warn(fmt.Sprintf("failed to update document metadata: %v", err)) + } + } + if err := u.docSvc.IncrementChunkNum(r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, 0); err != nil { + common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err)) + } +} + +// mergeDocMetadata reads existing metadata, fills in keys not already present +// (existing keys are preserved, not overwritten), and writes the merged map back. +func mergeDocMetadata(svc docStateSvc, docID string, metadata map[string]any) error { + existing, err := svc.GetDocumentMetadataByID(docID) + if err != nil { + existing = make(map[string]any) + } + for k, v := range metadata { + if _, exists := existing[k]; !exists { + existing[k] = v + } + } + return svc.SetDocumentMetadata(docID, existing) +} diff --git a/internal/ingestion/service/doc_state_test.go b/internal/ingestion/service/doc_state_test.go new file mode 100644 index 0000000000..77503d19da --- /dev/null +++ b/internal/ingestion/service/doc_state_test.go @@ -0,0 +1,122 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "testing" + + taskpkg "ragflow/internal/ingestion/task" +) + +type stubDocStateSvc struct { + metaData map[string]any + gotDocID string + gotKbID string + gotChunkNum int + gotTokenNum int + gotDuration float64 + setCalled bool + incrementCalled bool +} + +func (s *stubDocStateSvc) GetDocumentMetadataByID(docID string) (map[string]any, error) { + if s.metaData == nil { + return make(map[string]any), nil + } + return s.metaData, nil +} + +func (s *stubDocStateSvc) SetDocumentMetadata(docID string, meta map[string]any) error { + s.setCalled = true + if s.metaData == nil { + s.metaData = make(map[string]any) + } + for k, v := range meta { + s.metaData[k] = v + } + return nil +} + +func (s *stubDocStateSvc) IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error { + s.incrementCalled = true + s.gotDocID = docID + s.gotKbID = kbID + s.gotChunkNum = chunkNum + s.gotTokenNum = tokenNum + s.gotDuration = duration + return nil +} + +func TestDocStateUpdater_NilResultIsNoop(t *testing.T) { + svc := &stubDocStateSvc{} + u := &docStateUpdater{docSvc: svc} + u.apply(nil) + if svc.setCalled || svc.incrementCalled { + t.Fatal("nil result must not touch document state") + } +} + +func TestDocStateUpdater_EmptyMetadataSkipsMerge(t *testing.T) { + svc := &stubDocStateSvc{} + u := &docStateUpdater{docSvc: svc} + u.apply(&taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 3, TokenConsumption: 100}) + if svc.setCalled { + t.Fatal("empty metadata must not call SetDocumentMetadata") + } + if !svc.incrementCalled { + t.Fatal("counter must still be bumped") + } + if svc.gotChunkNum != 3 || svc.gotTokenNum != 100 { + t.Fatalf("chunk=%d token=%d, want 3/100", svc.gotChunkNum, svc.gotTokenNum) + } +} + +func TestDocStateUpdater_MergesNewKeys(t *testing.T) { + svc := &stubDocStateSvc{metaData: map[string]any{"existing": "old"}} + u := &docStateUpdater{docSvc: svc} + u.apply(&taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"new_key": "value"}, ChunkCount: 1, TokenConsumption: 10}) + if svc.metaData["existing"] != "old" { + t.Errorf("existing key should be preserved: got %q", svc.metaData["existing"]) + } + if svc.metaData["new_key"] != "value" { + t.Errorf("new_key = %q, want \"value\"", svc.metaData["new_key"]) + } +} + +func TestDocStateUpdater_PreservesExistingKey(t *testing.T) { + svc := &stubDocStateSvc{metaData: map[string]any{"author": "Alice"}} + u := &docStateUpdater{docSvc: svc} + u.apply(&taskpkg.PipelineResult{DocID: "doc-1", Metadata: map[string]any{"author": "Bob"}, ChunkCount: 1, TokenConsumption: 10}) + if svc.metaData["author"] != "Alice" { + t.Errorf("existing key must NOT be overwritten: got %q", svc.metaData["author"]) + } +} + +func TestDocStateUpdater_IncrementArgs(t *testing.T) { + svc := &stubDocStateSvc{} + u := &docStateUpdater{docSvc: svc} + u.apply(&taskpkg.PipelineResult{DocID: "doc-1", KbID: "kb-1", ChunkCount: 10, TokenConsumption: 100}) + if svc.gotDocID != "doc-1" || svc.gotKbID != "kb-1" { + t.Fatalf("docID=%q kbID=%q, want doc-1/kb-1", svc.gotDocID, svc.gotKbID) + } + if svc.gotChunkNum != 10 || svc.gotTokenNum != 100 { + t.Fatalf("chunk=%d token=%d, want 10/100", svc.gotChunkNum, svc.gotTokenNum) + } + if svc.gotDuration != 0 { + t.Fatalf("duration=%f, want 0", svc.gotDuration) + } +} diff --git a/internal/ingestion/service/execute_task_ack_test.go b/internal/ingestion/service/execute_task_ack_test.go new file mode 100644 index 0000000000..b942eace89 --- /dev/null +++ b/internal/ingestion/service/execute_task_ack_test.go @@ -0,0 +1,342 @@ +package service + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/ingestion/testutil" +) + +// fakeTaskHandle records Ack/Nack/InProgress calls so tests can assert the worker +// acknowledges the MQ message at the right time without a live broker. The +// counters are atomic because executeTask drives them from multiple goroutines +// (the heartbeat goroutine calls InProgress; the defer calls Ack/Nack) while the +// test goroutine reads them - plain ints would be a data race under -race. +type fakeTaskHandle struct { + msg common.TaskMessage + acks atomic.Int64 + nacks atomic.Int64 + inProgress atomic.Int64 +} + +func (f *fakeTaskHandle) GetMessage() common.TaskMessage { return f.msg } +func (f *fakeTaskHandle) Ack() error { f.acks.Add(1); return nil } +func (f *fakeTaskHandle) Nack() error { f.nacks.Add(1); return nil } +func (f *fakeTaskHandle) InProgress() error { f.inProgress.Add(1); return nil } + +func newAckTaskCtx(ctx context.Context, taskID, docID string, handle *fakeTaskHandle) *taskpkg.TaskContext { + taskCtx := taskpkg.NewTaskContextForScheduling( + ctx, + &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING}, + ) + taskCtx.Handle = handle + return taskCtx +} + +// TestExecuteTask_AcksMessageOnCompletion: a successfully completed task must +// Ack its MQ message so consumer-group queues do not redeliver (and double- +// execute) it. Regression test for the missing ack on the success path. +func TestExecuteTask_AcksMessageOnCompletion(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return nil + } + + handle := &fakeTaskHandle{} + ingestor.executeTask(newAckTaskCtx(context.Background(), taskID, docID, handle)) + + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("expected 1 Ack / 0 Nack on completion, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load final task: %v", err) + } + if finalTask.Status != common.COMPLETED { + t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED) + } +} + +// TestExecuteTask_AcksMessageOnFailure: a task whose runDocumentTask fails is +// marked FAILED (a terminal status) and must still be Acked - redelivery would +// only re-find it FAILED and skip. +func TestExecuteTask_AcksMessageOnFailure(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return errors.New("boom") + } + + handle := &fakeTaskHandle{} + ingestor.executeTask(newAckTaskCtx(context.Background(), taskID, docID, handle)) + + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("expected 1 Ack / 0 Nack on failure, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load final task: %v", err) + } + if finalTask.Status != common.FAILED { + t.Fatalf("final status = %s, want %s", finalTask.Status, common.FAILED) + } +} + +// TestExecuteTask_AcksMessageOnContextCancel: a task with a cancelled context +// (e.g. Redis cancel flag or Ingestor shutdown) is now terminal — the cancel +// is durably recorded (progress=-1, run=CANCEL) and the message is Acked to +// prevent indefinite redeliveries of an already-cancelled task. +func TestExecuteTask_AcksMessageOnContextCancel(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + var runCalled bool + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + runCalled = true + return nil + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + handle := &fakeTaskHandle{} + ingestor.executeTask(newAckTaskCtx(ctx, taskID, docID, handle)) + + if runCalled { + t.Fatal("expected runDocumentTask to be skipped on cancelled ctx") + } + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("expected 1 Ack / 0 Nack on cancel, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestExecuteTask_HeartbeatsInProgressDuringLongTask: a long-running task +// must call InProgress periodically while processing so the broker does not +// redeliver the unacked message mid-task (AckWait timer stays fresh). +// Regression test for in-flight double execution of slow tasks. +func TestExecuteTask_HeartbeatsInProgressDuringLongTask(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.heartbeatInterval = 5 * time.Millisecond + + started := make(chan struct{}) + proceed := make(chan struct{}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + close(started) // heartbeat goroutine is running by now + <-proceed // simulate a long task + return nil + } + + handle := &fakeTaskHandle{} + go ingestor.executeTask(newAckTaskCtx(context.Background(), taskID, docID, handle)) + + <-started + time.Sleep(30 * time.Millisecond) // let the ticker fire a few times + close(proceed) // release the long task + + if handle.inProgress.Load() == 0 { + t.Fatal("expected InProgress heartbeats while runDocumentTask was blocked, got 0") + } + + // Poll for Ack completion (executeTask must finish MarkCompleted + Ack). + deadline := time.Now().Add(2 * time.Second) + for handle.acks.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if handle.acks.Load() != 1 { + t.Fatalf("expected 1 Ack on completion after heartbeat, got acks=%d nacks=%d inProgress=%d", + handle.acks.Load(), handle.nacks.Load(), handle.inProgress.Load()) + } +} + +// TestClaimTask_FirstTrueThenFalse: claiming a task for the first time must +// succeed; a second claim while the first worker is still processing must +// fail. This is the local guard that catches MQ redeliveries. +func TestClaimTask_FirstTrueThenFalse(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + + if !ingestor.claimTask("task-1") { + t.Fatal("first claim should succeed") + } + + // Same task claimed again — must fail (another worker is already on it). + if ingestor.claimTask("task-1") { + t.Fatal("second claim should fail: task already claimed by another worker") + } + + // Different task — should succeed. + if !ingestor.claimTask("task-2") { + t.Fatal("different task should succeed") + } +} + +// TestClaimTask_AfterReleaseCanReclaim: after a worker finishes and releases +// the task, a fresh claim (e.g. on restart) must succeed again. +func TestClaimTask_AfterReleaseCanReclaim(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.claimTask("task-1") + ingestor.releaseTask("task-1") + + if !ingestor.claimTask("task-1") { + t.Fatal("after release should be able to reclaim (previous worker finished)") + } +} + +// TestExecuteTask_ReleasesTaskFromCurrentTasks: executeTask must call +// releaseTask in its defer so the task is removed from currentTasks and +// a subsequent claim succeeds. +func TestExecuteTask_ReleasesTaskFromCurrentTasks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return nil + } + ingestor.claimTask(taskID) + + handle := &fakeTaskHandle{} + ingestor.executeTask(newAckTaskCtx(context.Background(), taskID, docID, handle)) + + if _, stillActive := ingestor.currentTasks[taskID]; stillActive { + t.Fatal("expected task released from currentTasks after executeTask finished") + } + // After release, a new claim must succeed. + if !ingestor.claimTask(taskID) { + t.Fatal("expected reclaim after executeTask to succeed") + } +} + +// TestSettleMessage_AckOnTerminal: body returns true -> Ack, no Nack. +func TestSettleMessage_AckOnTerminal(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + ingestor.settleMessage(taskCtx, func(ctx context.Context) bool { return true }) + + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("body=true: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestSettleMessage_NackOnNonTerminal: body returns false -> Nack, no Ack. +func TestSettleMessage_NackOnNonTerminal(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + ingestor.settleMessage(taskCtx, func(ctx context.Context) bool { return false }) + + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("body=false: expected 1 Nack/0 Ack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestSettleMessage_NackOnPanic: if body panics, the message is still Nacked +// (terminal=false default) and the panic propagates. +func TestSettleMessage_NackOnPanic(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + panicked := true + func() { + defer func() { recover() }() + ingestor.settleMessage(taskCtx, func(ctx context.Context) bool { + panic("boom") + }) + panicked = false + }() + + if !panicked { + t.Fatal("expected panic to propagate, got none") + } + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("body=panic: expected 1 Nack/0 Ack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestAckOrNack_AckOnTerminal: terminal=true -> Ack called, Nack not called. +func TestAckOrNack_AckOnTerminal(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + ingestor.ackOrNack(taskCtx, true) + + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("terminal=true: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestAckOrNack_NackOnNonTerminal: terminal=false -> Nack called, Ack not called. +func TestAckOrNack_NackOnNonTerminal(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + ingestor.ackOrNack(taskCtx, false) + + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("terminal=false: expected 1 Nack/0 Ack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestAckOrNack_NoOpWhenNoHandle: nil handle -> no ack/nack, no panic. +func TestAckOrNack_NoOpWhenNoHandle(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + taskCtx := taskpkg.NewTaskContextForScheduling( + context.Background(), + &entity.IngestionTask{ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.RUNNING}, + ) + // taskCtx.Handle is nil + // Must not panic + ingestor.ackOrNack(taskCtx, true) + ingestor.ackOrNack(taskCtx, false) +} diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go index 1767d0a1bd..c9d5056509 100644 --- a/internal/ingestion/service/execute_task_test.go +++ b/internal/ingestion/service/execute_task_test.go @@ -25,12 +25,11 @@ func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - // Create a task log with invalid checkpoint (current_step is a string instead of number) + // Create a task log with invalid checkpoint (run_count is a string instead of number) err := db.Create(&entity.IngestionTaskLog{ TaskID: taskID, Checkpoint: entity.JSONMap{ - "current_step": "not-a-number", // intentionally wrong type - "total_step": 5, + "run_count": "not-a-number", // intentionally wrong type }, }).Error if err != nil { @@ -53,18 +52,19 @@ func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) { // Execute the task - this should NOT panic or fatal exit (this is our main validation!) ingestor.executeTask(taskCtx) - // Verify runDocumentTask was not called (we returned early due to checkpoint parse failure) - if runDocumentTaskCalled { - t.Fatal("expected runDocumentTask to not be called due to checkpoint parse failure") + // Corrupted run_count values are skipped by IncrementRunCount, so the task + // proceeds to runDocumentTask and completes normally. + if !runDocumentTaskCalled { + t.Fatal("expected runDocumentTask to be called (bad run_count is skipped, not fatal)") } - // Verify task status was set to FAILED + // Verify task status was set to COMPLETED finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) if err != nil { t.Fatalf("load final ingestion task: %v", err) } - if finalTask.Status != common.FAILED { - t.Fatalf("final status = %s, want %s", finalTask.Status, common.FAILED) + if finalTask.Status != common.COMPLETED { + t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED) } } @@ -95,7 +95,7 @@ func TestDefaultRunDocumentTask_RequiresConfiguredPipelineID(t *testing.T) { } } -func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { +func TestExecuteTask_PipelineRoutesToTaskHandler(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() @@ -106,19 +106,19 @@ func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { ) ingestor := NewIngestor("test", 1, []string{"pdf"}) - var routedToDataflow bool + var routedToPipeline bool var gotTaskID string var gotProgress []float64 var gotMsgs []string ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { - routedToDataflow = true + routedToPipeline = true gotTaskID = ingestionTask.ID wrapped := func(prog float64, msg string) { gotProgress = append(gotProgress, prog*100) gotMsgs = append(gotMsgs, msg) } - wrapped(0.82, "mock dataflow start") - wrapped(1.0, "mock dataflow done") + wrapped(0.82, "mock pipeline start") + wrapped(1.0, "mock pipeline done") return nil } @@ -129,8 +129,8 @@ func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { ingestor.executeTask(taskCtx) - if !routedToDataflow { - t.Fatal("expected executeTask to route dataflow task to runDocumentTask") + if !routedToPipeline { + t.Fatal("expected executeTask to route pipeline task to runDocumentTask") } if gotTaskID != taskID { t.Fatalf("runDocumentTask got task ID %q, want %q", gotTaskID, taskID) @@ -145,7 +145,55 @@ func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { if len(gotProgress) != 2 || gotProgress[0] != 82 || gotProgress[1] != 100 { t.Fatalf("gotProgress = %v, want [82 100]", gotProgress) } - if len(gotMsgs) != 2 || gotMsgs[1] != "mock dataflow done" { - t.Fatalf("gotMsgs = %v, want final message %q", gotMsgs, "mock dataflow done") + if len(gotMsgs) != 2 || gotMsgs[1] != "mock pipeline done" { + t.Fatalf("gotMsgs = %v, want final message %q", gotMsgs, "mock pipeline done") + } +} + +// TestExecuteTask_CancelBeforePipeline verifies that when cancelCheck returns +// true at task start, the task is cancelled before AdvanceStep, +// runDocumentTask is never called, and document progress is set to -1 with a +// cancel marker. Mirrors Python's cancel flow where has_canceled() returns +// true in Pipeline.callback(). +func TestExecuteTask_CancelBeforePipeline(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.cancelCheck = func(taskID string) bool { return true } + + var runDocumentTaskCalled bool + ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { + runDocumentTaskCalled = true + return nil + } + + taskCtx := taskpkg.NewTaskContextForScheduling( + context.Background(), + &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING}, + ) + ingestor.executeTask(taskCtx) + + if runDocumentTaskCalled { + t.Fatal("expected runDocumentTask to NOT be called when cancel is detected before pipeline") + } + + doc, err := dao.NewDocumentDAO().GetByID(docID) + if err != nil { + t.Fatalf("load document: %v", err) + } + if doc.Progress != -1 { + t.Fatalf("document.progress = %v, want -1 (cancelled)", doc.Progress) + } + if doc.Run == nil || *doc.Run != string(entity.TaskStatusCancel) { + t.Fatalf("document.run = %v, want %s (CANCEL)", doc.Run, entity.TaskStatusCancel) + } + if doc.ProgressMsg == nil || *doc.ProgressMsg == "" { + t.Fatal("document.progress_msg should contain cancel marker, got empty") } } diff --git a/internal/ingestion/service/heartbeat_test.go b/internal/ingestion/service/heartbeat_test.go new file mode 100644 index 0000000000..243306ec15 --- /dev/null +++ b/internal/ingestion/service/heartbeat_test.go @@ -0,0 +1,102 @@ +package service + +import ( + "context" + "sync" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" +) + +// controllableHandle is a TaskHandle whose InProgress behavior is delegated, +// for testing heartbeat timing and shutdown without a live broker. +type controllableHandle struct { + inProgressFn func() error +} + +func (h *controllableHandle) GetMessage() common.TaskMessage { return common.TaskMessage{} } +func (h *controllableHandle) Ack() error { return nil } +func (h *controllableHandle) Nack() error { return nil } +func (h *controllableHandle) InProgress() error { return h.inProgressFn() } + +// TestStartHeartbeat_TicksInProgressUntilStop: with a short interval the +// heartbeat goroutine calls InProgress repeatedly; stop() halts it. +func TestStartHeartbeat_TicksInProgressUntilStop(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.heartbeatInterval = 2 * time.Millisecond + + handle := &fakeTaskHandle{} + taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) + + stop := ingestor.startHeartbeat(taskCtx) + time.Sleep(15 * time.Millisecond) + stop() + + if handle.inProgress.Load() == 0 { + t.Fatal("expected InProgress heartbeats, got 0") + } +} + +// TestStartHeartbeat_StopWaitsForInFlightInProgress: stop() must block until an +// in-flight InProgress call returns, so the caller can ack/nack with no +// concurrent InProgress on the same message. Regression guard for the +// close-without-wait heartbeat shutdown (problem 1). +func TestStartHeartbeat_StopWaitsForInFlightInProgress(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.heartbeatInterval = time.Millisecond + + started := make(chan struct{}) + release := make(chan struct{}) + var startedOnce sync.Once + h := &controllableHandle{ + inProgressFn: func() error { + startedOnce.Do(func() { close(started) }) + <-release + return nil + }, + } + taskCtx := taskpkg.NewTaskContextForScheduling( + context.Background(), + &entity.IngestionTask{ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.RUNNING}, + ) + taskCtx.Handle = h + + stop := ingestor.startHeartbeat(taskCtx) + <-started // heartbeat goroutine is inside InProgress + + stopDone := make(chan struct{}) + go func() { stop(); close(stopDone) }() + + select { + case <-stopDone: + t.Fatal("stop() returned before in-flight InProgress completed") + case <-time.After(20 * time.Millisecond): + } + + close(release) // let InProgress complete + + select { + case <-stopDone: + // good: stop() returned only after InProgress completed + case <-time.After(time.Second): + t.Fatal("stop() did not return after in-flight InProgress completed") + } +} + +// TestStartHeartbeat_NoOpWhenNoHandle: with no MQ handle (standalone/test path) +// startHeartbeat returns a no-op stop and starts no goroutine. +func TestStartHeartbeat_NoOpWhenNoHandle(t *testing.T) { + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.heartbeatInterval = time.Millisecond + + taskCtx := taskpkg.NewTaskContextForScheduling( + context.Background(), + &entity.IngestionTask{ID: "task-1"}, + ) + // taskCtx.Handle is nil + stop := ingestor.startHeartbeat(taskCtx) + stop() // must not block or panic +} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index e1f891a3b2..bf12c629b0 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -22,31 +22,29 @@ import ( "fmt" "ragflow/internal/utility" "sync" + "time" "ragflow/internal/common" - "ragflow/internal/dao" "ragflow/internal/engine" + redis2 "ragflow/internal/engine/redis" "ragflow/internal/entity" taskpkg "ragflow/internal/ingestion/task" servicepkg "ragflow/internal/service" - "google.golang.org/grpc" - "gorm.io/gorm" + "github.com/cenkalti/backoff/v5" ) type Ingestor struct { - id string - name string - serverAddr string - conn *grpc.ClientConn - ctx context.Context - cancel context.CancelFunc - reconnectMu sync.Mutex + id string + name string + ctx context.Context + cancel context.CancelFunc // Configuration maxConcurrency int32 supportedDocTypes []string version string + heartbeatInterval time.Duration // Runtime state currentTasks map[string]*taskpkg.TaskContext @@ -60,35 +58,42 @@ type Ingestor struct { workerWg sync.WaitGroup startOnce sync.Once - ingestionTaskDAO *dao.IngestionTaskDAO - ingestionTaskLogDAO *dao.IngestionTaskLogDAO - ingestionTaskSvc *servicepkg.IngestionTaskService + ingestionTaskSvc *servicepkg.IngestionTaskService + docState *docStateUpdater // runDocumentTask dispatches to the migrated task handler path. // Tests may override this to verify branch routing without invoking // the full downstream stack. runDocumentTask func(ctx context.Context, ingestionTask *entity.IngestionTask) error + + // cancelCheck is polled periodically (every 3s) during task execution. + // When it returns true the task's context is cancelled, which causes the + // pipeline to stop at the next ctx.Err() check. Defaults to a Redis + // cancel-flag lookup that mirrors Python's has_canceled(). Tests may + // override this to simulate cancel without Redis. + cancelCheck func(taskID string) bool } func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *Ingestor { ctx, cancel := context.WithCancel(context.Background()) id := utility.GenerateUUID() ingestor := &Ingestor{ - id: id, - name: name, - ctx: ctx, - cancel: cancel, - maxConcurrency: maxConcurrency, - supportedDocTypes: supportedTypes, - version: "1.0.0", - currentTasks: make(map[string]*taskpkg.TaskContext), - taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), - ShutdownCh: make(chan struct{}, 1), - ingestionTaskDAO: dao.NewIngestionTaskDAO(), - ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), - ingestionTaskSvc: servicepkg.NewIngestionTaskService(), + id: id, + name: name, + ctx: ctx, + cancel: cancel, + maxConcurrency: maxConcurrency, + supportedDocTypes: supportedTypes, + version: "1.0.0", + currentTasks: make(map[string]*taskpkg.TaskContext), + taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), + ShutdownCh: make(chan struct{}, 1), + ingestionTaskSvc: servicepkg.NewIngestionTaskService(), + docState: newDocStateUpdater(), + heartbeatInterval: 10 * time.Second, } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask + ingestor.cancelCheck = ingestor.defaultCancelCheck return ingestor } @@ -115,77 +120,95 @@ func (e *Ingestor) Start() error { continue } for _, taskHandle := range taskHandles { - taskMessage := taskHandle.GetMessage() - common.Info(fmt.Sprintf("Received task id: %s, type: %s", taskMessage.TaskID, taskMessage.TaskType)) - if taskMessage.TaskType != common.TaskTypeIngestionTask { - common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err - } - continue - } - var task *entity.IngestionTask - task, err = e.ingestionTaskSvc.StartRunning(taskMessage.TaskID) - if err != nil { - if errors.Is(err, common.ErrTaskNotFound) { - common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err - } - continue - } else { - common.Error(fmt.Sprintf("error setting task %s to running", taskMessage.TaskID), err) - return err - } - } - if task == nil { - common.Info(fmt.Sprintf("task %s is already removed", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - return err - } - continue - } - - switch task.Status { - case common.COMPLETED, common.STOPPED, common.FAILED: - common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err - } - continue - case common.STOPPING, common.CREATED: - err = fmt.Errorf("task %s is in unexpected status %s", taskMessage.TaskID, task.Status) + if err := e.processMessage(taskHandle); err != nil { return err - case common.RUNNING: - } - - // Construct TaskContext with parent context - taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task) - - // Push to task channel; if full, reject the task (backpressure) - select { - case e.taskChan <- taskCtx: - common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.ID, len(e.taskChan), cap(e.taskChan))) - default: - common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID)) - - err = taskHandle.Nack() - if err != nil { - common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) - return err - } } } } } +// processMessage handles a single incoming MQ message: filter by type, +// activate the task (state transition), guard against duplicate execution +// (claim), and enqueue to the worker pool (or backpressure-reject). +func (e *Ingestor) processMessage(handle common.TaskHandle) error { + taskMessage := handle.GetMessage() + common.Info(fmt.Sprintf("Received task id: %s, type: %s", taskMessage.TaskID, taskMessage.TaskType)) + + if taskMessage.TaskType != common.TaskTypeIngestionTask { + common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) + if err := handle.Ack(); err != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) + return err + } + return nil + } + + task, err := e.ingestionTaskSvc.StartRunning(taskMessage.TaskID) + if err != nil { + if errors.Is(err, common.ErrTaskNotFound) { + common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) + if ackErr := handle.Ack(); ackErr != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) + return ackErr + } + return nil + } + common.Error(fmt.Sprintf("error setting task %s to running", taskMessage.TaskID), err) + return err + } + if task == nil { + common.Info(fmt.Sprintf("task %s is already removed", taskMessage.TaskID)) + if ackErr := handle.Ack(); ackErr != nil { + return ackErr + } + return nil + } + + switch task.Status { + case common.COMPLETED, common.STOPPED, common.FAILED: + common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) + if ackErr := handle.Ack(); ackErr != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) + return ackErr + } + return nil + case common.STOPPING, common.CREATED: + return fmt.Errorf("task %s is in unexpected status %s", taskMessage.TaskID, task.Status) + case common.RUNNING: + // Guard against MQ redelivery: if another worker in this + // process is already processing this task, ack the redelivered + // message and skip instead of scheduling it again. + if !e.claimTask(task.ID) { + common.Warn(fmt.Sprintf("task %s redelivered while worker still processing, ack skip (task_id=%s doc_id=%s kb_id=%s)", + taskMessage.TaskID, task.ID, task.DocumentID, task.DatasetID)) + if ackErr := handle.Ack(); ackErr != nil { + common.Error(fmt.Sprintf("error ack redelivered task %s", taskMessage.TaskID), ackErr) + return ackErr + } + return nil + } + } + + // Construct TaskContext and carry the MQ handle so the worker can + // Ack/Nack when the task reaches a terminal status. + taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task) + taskCtx.Handle = handle + + // Push to task channel; if full, reject the task (backpressure). + select { + case e.taskChan <- taskCtx: + common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.ID, len(e.taskChan), cap(e.taskChan))) + default: + common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID)) + e.releaseTask(task.ID) + if nackErr := handle.Nack(); nackErr != nil { + common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), nackErr) + return nackErr + } + } + return nil +} + func (e *Ingestor) startWorkerPool() { e.startOnce.Do(func() { for i := int32(0); i < e.maxConcurrency; i++ { @@ -211,84 +234,282 @@ func (e *Ingestor) workerLoop(id int32) { } func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { - ctx := taskCtx.Ctx task := taskCtx.IngestionTask common.Info(fmt.Sprintf("Starting task %s", task.ID)) - latestLog, err := e.ingestionTaskLogDAO.LatestLogByTaskID(task.ID) - if err != nil { - if !errors.Is(err, gorm.ErrRecordNotFound) { - common.Error(fmt.Sprintf("Failed to get latest task log for task %s", task.ID), err) - if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { - common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) - } - return - } - latestLog = &entity.IngestionTaskLog{ - ID: 0, - TaskID: task.ID, - Checkpoint: entity.JSONMap{ - "current_step": 1, - "total_step": 5, - }, - } - err = e.ingestionTaskLogDAO.Create(latestLog) - if err != nil { - common.Error(fmt.Sprintf("Failed to create task log for task %s", task.ID), err) - if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { - common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) - } - return - } + // Release the claim when executeTask returns so that a future + // redelivery (after restart) can re-claim the task. + defer e.releaseTask(task.ID) + + // Derive a per-task cancelable context so that an external cancel + // signal (Redis cancel flag, mirrored from Python's {task_id}-cancel) + // can stop only this task without affecting the whole Ingestor. + perTaskCtx, perTaskCancel := context.WithCancel(taskCtx.Ctx) + taskCtx.Ctx = perTaskCtx + cancelDone := make(chan struct{}) + pollExited := make(chan struct{}) + go func() { + defer close(pollExited) + e.pollCancel(task.ID, perTaskCancel, cancelDone) + }() + defer func() { + close(cancelDone) + <-pollExited + perTaskCancel() + }() + + // Synchronous check: if already cancelled (e.g. flag set between MQ + // delivery and worker claim), stop before the pipeline even starts. + if e.cancelCheck(task.ID) { + common.Info(fmt.Sprintf("Task %s cancel flag detected before pipeline start, cancelling", task.ID)) + perTaskCancel() } - var checkpointMap map[string]interface{} - checkpointMap = latestLog.Checkpoint - currentStep, ok := common.GetInt(checkpointMap["current_step"]) - if !ok { - common.Error(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID), nil) - if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { - common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) - } - return - } - totalStep, ok := common.GetInt(checkpointMap["total_step"]) - if !ok { - common.Error(fmt.Sprintf("Failed to get total step from task log for task %s", task.ID), nil) - if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { - common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) - } - return - } - // Bump checkpoint to signal we started processing - checkpointMap["current_step"] = currentStep + 1 - latestLog.Checkpoint = checkpointMap - if err := e.ingestionTaskLogDAO.Update(latestLog); err != nil { - common.Error(fmt.Sprintf("Failed to persist checkpoint for task %s", task.ID), err) - } - _ = totalStep // unused in this temporary integration + e.settleMessage(taskCtx, func(ctx context.Context) bool { + return e.runTask(ctx, task) + }) +} +// markStopped transitions the task to STOPPED (terminal). It first calls +// RequestStop to handle RUNNING → STOPPING, then MarkStopped for the final +// STOPPING → STOPPED transition. Finally it cleans up the Redis cancel flag +// so that a future retry of the same task does not immediately re-cancel. +func (e *Ingestor) markStopped(taskID string) bool { + if _, err := e.ingestionTaskSvc.RequestStop(taskID); err != nil { + common.Error(fmt.Sprintf("markStopped: RequestStop task %s: %v", taskID, err), err) + } + if err := e.ingestionTaskSvc.MarkStopped(taskID); err != nil { + common.Error(fmt.Sprintf("markStopped: MarkStopped task %s: %v", taskID, err), err) + return false + } + if rc := redis2.Get(); rc != nil { + rc.Delete(fmt.Sprintf("%s-cancel", taskID)) + } + return true +} + +// markFailed persists FAILED status for the task and reports whether the +// terminal status was durably written, so the caller can decide Ack vs Nack. +func (e *Ingestor) markFailed(taskID string) bool { + if uErr := e.ingestionTaskSvc.MarkFailed(taskID); uErr != nil { + common.Error(fmt.Sprintf("Failed to set task %s to FAILED", taskID), uErr) + return false + } + return true +} + +// runTask executes the task's business logic — run-count advance, document +// pipeline, and completion — behind the heartbeat. It returns whether the +// task reached a durably-persisted terminal status. +func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool { select { case <-ctx.Done(): common.Info(fmt.Sprintf("Task %s cancelled", task.ID)) - return + e.markCancelProgress(task) + e.markStopped(task.ID) + return true default: } - if err := e.runDocumentTask(ctx, task); err != nil { - common.Error(fmt.Sprintf("Task %s failed", task.ID), err) - if uErr := e.ingestionTaskSvc.MarkFailed(task.ID); uErr != nil { - common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) - } - return + + if err := e.ingestionTaskSvc.IncrementRunCount(task.ID); err != nil { + common.Error(fmt.Sprintf("Failed to increment run count for task %s", task.ID), err) + return e.markFailed(task.ID) } - err = e.ingestionTaskSvc.MarkCompleted(task.ID) + if err := e.runDocumentTask(ctx, task); err != nil { + if errors.Is(err, context.Canceled) { + common.Info(fmt.Sprintf("Task %s cancelled during pipeline", task.ID)) + e.markCancelProgress(task) + return e.markStopped(task.ID) + } + if errors.Is(err, context.DeadlineExceeded) { + common.Info(fmt.Sprintf("Task %s timed out during pipeline", task.ID)) + e.markTimeoutProgress(task) + return e.markFailed(task.ID) + } + common.Error(fmt.Sprintf("Task %s failed", task.ID), err) + return e.markFailed(task.ID) + } + + _, err := backoff.Retry(ctx, func() (struct{}, error) { + return struct{}{}, e.ingestionTaskSvc.MarkCompleted(task.ID) + }, backoff.WithMaxTries(3)) if err != nil { common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err) - return + return false } common.Info(fmt.Sprintf("Task %s completed", task.ID)) + return true +} + +// settleMessage runs body under a heartbeat, then settles the MQ message. The +// heartbeat is stopped (and waited on) before ack/nack — see startHeartbeat. +// terminal is derived from body's return value; on panic terminal defaults to +// false (non-terminal → Nack) so the broker redelivers after restart. +func (e *Ingestor) settleMessage(taskCtx *taskpkg.TaskContext, body func(context.Context) bool) (terminal bool) { + stop := e.startHeartbeat(taskCtx) + defer func() { + stop() // stop heartbeat (and wait) before ack/nack + e.ackOrNack(taskCtx, terminal) + }() + terminal = body(taskCtx.Ctx) + return +} + +// ackOrNack settles the MQ message according to the terminal flag: Ack if the +// task reached a durably-persisted terminal status, Nack otherwise so the +// broker redelivers and resumes. A nil handle (standalone/test path) is a no-op. +func (e *Ingestor) ackOrNack(taskCtx *taskpkg.TaskContext, terminal bool) { + if taskCtx.Handle == nil { + return + } + if terminal { + if err := taskCtx.Handle.Ack(); err != nil { + common.Error(fmt.Sprintf("ack task %s", taskCtx.IngestionTask.ID), err) + } + return + } + if err := taskCtx.Handle.Nack(); err != nil { + common.Error(fmt.Sprintf("nack task %s", taskCtx.IngestionTask.ID), err) + } +} + +// defaultCancelCheck reads the Redis cancel flag that Python sets via +// REDIS_CONN.set(f"{task_id}-cancel", "x"). Falls back to checking the +// task status in DB when Redis is unavailable — a STOPPING status +// (set by RequestStop) is treated as a cancel signal. +func (e *Ingestor) defaultCancelCheck(taskID string) bool { + rc := redis2.Get() + if rc != nil { + if ok, _ := rc.Exist(fmt.Sprintf("%s-cancel", taskID)); ok { + return true + } + } + task, err := e.ingestionTaskSvc.GetTask(taskID) + if err != nil { + return false + } + return task.Status == common.STOPPING +} + +// pollCancel ticks every 3s to check the cancel flag. When cancelCheck +// returns true it cancels the per-task context, which causes the pipeline's +// next ctx.Err() check to abort and runTask to record progress=-1. The +// goroutine exits when done is closed (executeTask returns). +func (e *Ingestor) pollCancel(taskID string, cancel context.CancelFunc, done <-chan struct{}) { + // Check immediately so the test path (which sets cancelCheck to a func + // that returns true) does not need to wait for the first tick. + if e.cancelCheck(taskID) { + common.Info(fmt.Sprintf("Task %s cancel flag detected during polling, cancelling pipeline", taskID)) + cancel() + return + } + ticker := time.NewTicker(3 * time.Second) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ticker.C: + if e.cancelCheck(taskID) { + cancel() + return + } + } + } +} + +// markCancelProgress writes the cancelled-progress markers to the document +// row. Mirrors Python's cancel_all_task_of: progress=-1, run=CANCEL, and an +// appended timestamped cancel message (progress_msg += cancelMsg). +func (e *Ingestor) markCancelProgress(task *entity.IngestionTask) { + svc := servicepkg.NewDocumentService() + doc, err := svc.GetDocumentByID(task.DocumentID) + if err != nil { + common.Error(fmt.Sprintf("markCancelProgress: load document %s: %v", task.DocumentID, err), err) + return + } + cancelMsg := fmt.Sprintf("\n%s Task stopped by user.", time.Now().Format("15:04:05")) + existingMsg := "" + if doc.ProgressMsg != nil { + existingMsg = *doc.ProgressMsg + } + _ = svc.UpdateRunProgress(task.DocumentID, -1.0, string(entity.TaskStatusCancel), existingMsg+cancelMsg) +} + +// markTimeoutProgress writes the timeout-progress markers to the document +// row. Unlike cancellation (markCancelProgress), this records a TIMEOUT +// failure rather than a user-initiated stop. +func (e *Ingestor) markTimeoutProgress(task *entity.IngestionTask) { + svc := servicepkg.NewDocumentService() + doc, err := svc.GetDocumentByID(task.DocumentID) + if err != nil { + common.Error(fmt.Sprintf("markTimeoutProgress: load document %s: %v", task.DocumentID, err), err) + return + } + timeoutMsg := fmt.Sprintf("\n%s Task timed out.", time.Now().Format("15:04:05")) + existingMsg := "" + if doc.ProgressMsg != nil { + existingMsg = *doc.ProgressMsg + } + _ = svc.UpdateRunProgress(task.DocumentID, -1.0, string(entity.TaskStatusFail), existingMsg+timeoutMsg) +} + +// claimTask registers a worker claim on a task ID. Returns false if another +// worker has already claimed it (e.g. MQ redelivery), true on first claim. +// startHeartbeat launches a goroutine that calls Handle.InProgress every +// heartbeatInterval to keep the broker AckWait timer fresh during long tasks. +// It returns a stop function that signals the goroutine to exit and BLOCKS +// until it has, so the caller can ack/nack with no in-flight InProgress on the +// same message. Returns a no-op stop when there is no handle or no interval +// (standalone/test path). +func (e *Ingestor) startHeartbeat(taskCtx *taskpkg.TaskContext) func() { + if taskCtx.Handle == nil || e.heartbeatInterval <= 0 { + return func() {} + } + var wg sync.WaitGroup + wg.Add(1) + done := make(chan struct{}) + go func() { + defer wg.Done() + ticker := time.NewTicker(e.heartbeatInterval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := taskCtx.Handle.InProgress(); err != nil { + common.Error(fmt.Sprintf("heartbeat task %s", taskCtx.IngestionTask.ID), err) + } + case <-done: + return + case <-taskCtx.Ctx.Done(): + return + } + } + }() + return func() { + close(done) + wg.Wait() + } +} + +func (e *Ingestor) claimTask(taskID string) bool { + e.tasksMu.Lock() + defer e.tasksMu.Unlock() + if _, ok := e.currentTasks[taskID]; ok { + return false + } + e.currentTasks[taskID] = nil // placeholder; replaced after scheduling + return true +} + +// releaseTask removes the claim so a future redelivery (after process restart) +// can re-claim the task. +func (e *Ingestor) releaseTask(taskID string) { + e.tasksMu.Lock() + delete(e.currentTasks, taskID) + e.tasksMu.Unlock() } func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *entity.IngestionTask) error { @@ -300,7 +521,24 @@ func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *en return fmt.Errorf("ingestion task %s: no pipeline_id configured for document %s or dataset %s", ingestionTask.ID, docTaskCtx.Doc.ID, docTaskCtx.KB.ID) } docTaskCtx.Ctx = ctx - return taskpkg.NewTaskHandler(docTaskCtx).Handle() + // The sink owns all document/ingestion_task_log/ingestion_task.component_total + // writes for this run; inject it into the executor so the pipeline reports + // progress to the service layer instead of touching the DAO directly. + sink := newProgressSink(e.ingestionTaskSvc) + result, err := taskpkg.NewTaskHandler(docTaskCtx). + WithPipelineExecutorFactory(func(c *taskpkg.TaskContext, canvasID string) (*taskpkg.PipelineExecutor, error) { + ex, err := taskpkg.NewPipelineExecutor(c, canvasID, 0) + if err != nil { + return nil, err + } + return ex.WithProgressSink(sink), nil + }). + Handle() + if err != nil { + return err + } + e.docState.apply(result) + return nil } // Stop gracefully shuts down the ingestor diff --git a/internal/ingestion/service/ingestor_lifecycle_test.go b/internal/ingestion/service/ingestor_lifecycle_test.go new file mode 100644 index 0000000000..bdd1baf2b1 --- /dev/null +++ b/internal/ingestion/service/ingestor_lifecycle_test.go @@ -0,0 +1,57 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "testing" + "time" +) + +// TestStartWorkerPool_StartOnceIdempotent verifies that calling startWorkerPool +// twice only starts maxConcurrency workers (sync.Once gate). +func TestStartWorkerPool_StartOnceIdempotent(t *testing.T) { + const concurrency int32 = 3 + ingestor := NewIngestor("test-idempotent", concurrency, nil) + // Stop background worker loops immediately so we can count workers. + ingestor.cancel() + ingestor.startWorkerPool() + ingestor.startWorkerPool() + ingestor.workerWg.Wait() +} + +// TestStop_GracefulShutdown verifies that Stop cancels the context and waits +// for all worker goroutines to exit without hanging. +func TestStop_GracefulShutdown(t *testing.T) { + const concurrency int32 = 2 + ingestor := NewIngestor("test-shutdown", concurrency, nil) + + // Start workers; they will block on the task channel since nothing is pushed. + ingestor.startWorkerPool() + + done := make(chan struct{}) + go func() { + ingestor.Stop() + close(done) + }() + + select { + case <-done: + // workers exited cleanly + case <-time.After(5 * time.Second): + t.Fatal("Stop() timed out waiting for workers to exit") + } +} diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go new file mode 100644 index 0000000000..8b302c3df4 --- /dev/null +++ b/internal/ingestion/service/process_message_test.go @@ -0,0 +1,180 @@ +package service + +import ( + "testing" + + "ragflow/internal/common" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/ingestion/testutil" +) + +func newFakeHandle(taskID, taskType string) *fakeTaskHandle { + return &fakeTaskHandle{msg: common.TaskMessage{TaskID: taskID, TaskType: taskType}} +} + +// TestProcessMessage_NonIngestionTaskAcks: a non-ingestion task is acked and +// skipped without touching the task DB or enqueuing. +func TestProcessMessage_NonIngestionTaskAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := newFakeHandle("task-1", "not-ingestion") + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (continue), got: %v", err) + } + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("non-ingestion: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 0 { + t.Fatal("expected no task enqueued") + } +} + +// TestProcessMessage_TaskNotFoundAcks: when StartRunning returns +// ErrTaskNotFound the message is acked and skipped. +func TestProcessMessage_TaskNotFoundAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + // No task seeded in DB — StartRunning returns ErrTaskNotFound. + handle := newFakeHandle("no-such-task", common.TaskTypeIngestionTask) + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (continue), got: %v", err) + } + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("not-found: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestProcessMessage_AlreadyCompletedAcks: a task already in a terminal state +// (COMPLETED) is acked and skipped — no enqueue, no status change. +func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Set the task COMPLETED so the status switch skips it. + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID). + Update("status", common.COMPLETED).Error; err != nil { + t.Fatalf("set COMPLETED: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (continue), got: %v", err) + } + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("completed: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 0 { + t.Fatal("expected no task enqueued for completed task") + } +} + +// TestProcessMessage_ClaimFailsAcks: a RUNNING task whose claim fails +// (redelivery guard) is acked without enqueuing — another worker is already +// processing it. +func TestProcessMessage_ClaimFailsAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + // Pre-claim the task so processMessage sees a claim conflict. + ingestor.claimTask(taskID) + + handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (continue), got: %v", err) + } + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("claim-fail: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 0 { + t.Fatal("expected no task enqueued when claim fails") + } +} + +// TestProcessMessage_ClaimSucceedsEnqueues: a RUNNING task with a successful +// claim is enqueued to the worker pool and the message is NOT settled yet +// (ack/nack is deferred to the worker). +func TestProcessMessage_ClaimSucceedsEnqueues(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (continue), got: %v", err) + } + // Ack/Nack must not be called — settlement is deferred to the worker. + if handle.acks.Load() != 0 || handle.nacks.Load() != 0 { + t.Fatalf("enqueued: expected 0 Ack/0 Nack (deferred), got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + // Task must be in the channel. + if len(ingestor.taskChan) != 1 { + t.Fatalf("expected 1 task enqueued, got %d", len(ingestor.taskChan)) + } + // Drain and verify. + taskCtx := <-ingestor.taskChan + if taskCtx.IngestionTask.ID != taskID { + t.Fatalf("enqueued task ID = %s, want %s", taskCtx.IngestionTask.ID, taskID) + } +} + +// TestProcessMessage_ChannelFullNacks: when the task channel is at capacity +// backpressure rejects the task with Nack, releases the claim, and returns nil +// so the message is redelivered and a future attempt can re-claim it. +func TestProcessMessage_ChannelFullNacks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // maxConcurrency=2 → channel cap=4. Fill it completely. + ingestor := NewIngestor("test", 2, []string{"pdf"}) + for i := 0; i < cap(ingestor.taskChan); i++ { + ingestor.taskChan <- taskpkg.NewTaskContextForScheduling(nil, &entity.IngestionTask{ID: "filler"}) + } + + handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) + + err := ingestor.processMessage(handle) + if err != nil { + t.Fatalf("expected nil (nack ok → continue), got: %v", err) + } + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("channel-full: expected 1 Nack/0 Ack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + + // Claim must be released so a future redelivery can re-claim it. + if !ingestor.claimTask(taskID) { + t.Fatal("claim was not released on channel-full — task would be stuck forever") + } + ingestor.releaseTask(taskID) + + // Drain the fillers. + for i := 0; i < cap(ingestor.taskChan); i++ { + <-ingestor.taskChan + } +} diff --git a/internal/ingestion/service/progress_sink.go b/internal/ingestion/service/progress_sink.go new file mode 100644 index 0000000000..1e295c4649 --- /dev/null +++ b/internal/ingestion/service/progress_sink.go @@ -0,0 +1,105 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "fmt" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/ingestion/pipeline" + servicepkg "ragflow/internal/service" +) + +// progressSink implements pipeline.ProgressSink. It is the single writer of +// the document / ingestion_task_log / ingestion_task.component_total tables +// for a pipeline run: the pipeline reports component lifecycle events here +// and the sink persists them through the service layer (IngestionTaskService +// + DocumentService), never the DAO directly. All writes are best-effort - +// failures are logged and never abort the run, mirroring the legacy +// pipeline-internal sink semantics. +type progressSink struct { + taskSvc *servicepkg.IngestionTaskService + docSvc docProgressSvc +} + +// docProgressSvc is the subset of *service.DocumentService the sink needs to +// mirror run progress into the document row. Extracted as an interface so +// tests can inject a stub and assert the mirror call without depending on the +// full DocumentService surface. +type docProgressSvc interface { + UpdateRunProgress(docID string, progress float64, run, progressMsg string) error +} + +func newProgressSink(taskSvc *servicepkg.IngestionTaskService) *progressSink { + // Eagerly construct the DocumentService so docSvc is immutable after this + // point. eino's compose graph runs parallel branches concurrently, so + // OnComponentProgress (and thus docService) can fire from multiple + // goroutines; a lazy check-then-act here would be a data race. The sink + // owns no server-config dependency, so this is safe in any environment. + return &progressSink{ + taskSvc: taskSvc, + docSvc: servicepkg.NewDocumentService(), + } +} + +// docService returns the DocumentService bound at sink construction. It is an +// accessor (not lazy) so concurrent progress callbacks read a stable value. +func (s *progressSink) docService() docProgressSvc { + return s.docSvc +} + +func (s *progressSink) OnComponentTotal(taskID string, total int) { + if err := s.taskSvc.UpdateComponentTotal(taskID, total); err != nil { + common.Error(fmt.Sprintf("progressSink: update component_total for task %s failed: %v", taskID, err), err) + } +} + +func (s *progressSink) OnComponentProgress(ev pipeline.ProgressEvent) { + if err := s.taskSvc.RecordComponentProgress(ev.TaskID, ev.Component, ev.Index, ev.Phase, ev.Message); err != nil { + common.Error(fmt.Sprintf("progressSink: record component progress for task %s failed: %v", ev.TaskID, err), err) + } + if ev.DocumentID == "" { + return + } + agg, err := s.taskSvc.AggregateTaskProgress(ev.TaskID, ev.Total) + if err != nil || agg == nil || ev.Total <= 0 { + return + } + progress, run := deriveDocumentProgress(agg, ev.Total) + if err := s.docService().UpdateRunProgress(ev.DocumentID, progress, run, ev.Message); err != nil { + common.Error(fmt.Sprintf("progressSink: mirror progress to document %s for task %s failed: %v", ev.DocumentID, ev.TaskID, err), err) + } +} + +// deriveDocumentProgress computes the document-level progress (0..1) and run +// label ("0".."4", matching Python's document.run enum) from the aggregated +// ingestion_task_log. This logic is owned by the sink (the document-table +// writer), not the pipeline. +func deriveDocumentProgress(agg *dao.TaskProgress, total int) (float64, string) { + run := string(entity.TaskStatusUnstart) + switch { + case agg.Failed > 0: + run = string(entity.TaskStatusFail) + case agg.Done == total: + run = string(entity.TaskStatusDone) + case agg.Done > 0 || agg.Running > 0: + run = string(entity.TaskStatusRunning) + } + return agg.Percent / 100, run +} diff --git a/internal/ingestion/service/progress_sink_test.go b/internal/ingestion/service/progress_sink_test.go new file mode 100644 index 0000000000..4b2d8e21d3 --- /dev/null +++ b/internal/ingestion/service/progress_sink_test.go @@ -0,0 +1,281 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "sync" + "testing" + + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/ingestion/pipeline" + "ragflow/internal/ingestion/testutil" + servicepkg "ragflow/internal/service" +) + +// TestProgressSink_CanConstructDocumentServiceWithoutServerConfig ensures the +// sink's DocumentService dependency can be built in a headless/test environment +// where server config is not initialized. NewDocumentService historically read +// server.GetConfig().DocEngine.Type, which nil-dereferenced without config; the +// sink must not pull the process-wide config just to mirror run progress. +func TestProgressSink_CanConstructDocumentServiceWithoutServerConfig(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + // No server config is initialized in the test env; this must not panic. + svc := servicepkg.NewDocumentService() + if svc == nil { + t.Fatal("expected non-nil DocumentService") + } +} + +// TestProgressSink_EagerlyConstructsDocumentService ensures the sink builds its +// DocumentService at construction time rather than lazily on the first progress +// event. Lazy construction is a data race under eino's parallel-branch progress +// callbacks (see TestProgressSink_OnComponentProgress_NoDataRace); eager +// construction makes docSvc immutable after newProgressSink returns. +func TestProgressSink_EagerlyConstructsDocumentService(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + sink := newProgressSink(servicepkg.NewIngestionTaskService()) + if sink.docSvc == nil { + t.Fatal("expected sink to eagerly construct its DocumentService, got nil (lazy)") + } +} + +// TestProgressSink_DocService_NoDataRace guards against regressing to lazy +// DocumentService construction. eino's compose graph runs parallel branches +// concurrently (compose/chain_parallel.go, branch.go), so the progress callback +// can fire from multiple goroutines; docService() must return a pre-built, +// immutable DocumentService, not lazily check-then-act on s.docSvc. +// +// The race is hit directly on docService() rather than through +// OnComponentProgress because the latter serializes on the single test-DB +// connection before reaching docService(), which masks the race. +func TestProgressSink_DocService_NoDataRace(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + // Deliberately do NOT inject a stub docSvc: the sink's own DocumentService + // must already be constructed (not lazily built mid-call) when the + // goroutines below race into docService(). + sink := newProgressSink(servicepkg.NewIngestionTaskService()) + + const n = 30 + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + _ = sink.docService() + }() + } + close(start) + wg.Wait() +} + +type stubDocProgressSvc struct { + gotDocID string + gotProgress float64 + gotRun string + gotMsg string + calls int +} + +func (s *stubDocProgressSvc) UpdateRunProgress(docID string, progress float64, run, progressMsg string) error { + s.calls++ + s.gotDocID = docID + s.gotProgress = progress + s.gotRun = run + s.gotMsg = progressMsg + return nil +} + +// TestProgressSinkPersistsViaService verifies the sink is the single writer of +// ingestion_task.component_total, ingestion_task_log, and document run-progress +// - all through the service layer, not the DAO. +func TestProgressSinkPersistsViaService(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + sink := newProgressSink(servicepkg.NewIngestionTaskService()) + stub := &stubDocProgressSvc{} + sink.docSvc = stub + + sink.OnComponentTotal(taskID, 2) + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.ComponentTotal != 2 { + t.Fatalf("component_total = %d, want 2", task.ComponentTotal) + } + + sink.OnComponentProgress(pipeline.ProgressEvent{ + TaskID: taskID, + DocumentID: docID, + Component: "Parser", + Index: 0, + Phase: 1, + Message: "Parser Done", + Total: 2, + }) + + logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) + if err != nil { + t.Fatalf("list logs: %v", err) + } + if len(logs) != 1 { + t.Fatalf("expected 1 component-progress row, got %d", len(logs)) + } + if logs[0].Component != "Parser" || logs[0].Phase != 1 || logs[0].Message != "Parser Done" { + t.Fatalf("unexpected log row: %+v", logs[0]) + } + + // 1 of 2 components done -> RUNNING (run "1"), progress 0.5. + if stub.calls != 1 { + t.Fatalf("UpdateRunProgress calls = %d, want 1", stub.calls) + } + if stub.gotDocID != docID { + t.Fatalf("docID = %q, want %q", stub.gotDocID, docID) + } + if stub.gotProgress != 0.5 { + t.Fatalf("progress = %v, want 0.5", stub.gotProgress) + } + if stub.gotRun != "1" { + t.Fatalf("run = %q, want 1 (RUNNING)", stub.gotRun) + } + if stub.gotMsg != "Parser Done" { + t.Fatalf("progress_msg = %q, want Parser Done", stub.gotMsg) + } +} + +// TestProgressSinkEmptyDocumentIDSkipsMirror verifies the log row is still +// recorded when no owning document is bound, but the document mirror is skipped. +func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + sink := newProgressSink(servicepkg.NewIngestionTaskService()) + stub := &stubDocProgressSvc{} + sink.docSvc = stub + + sink.OnComponentProgress(pipeline.ProgressEvent{ + TaskID: taskID, + Component: "Chunker", + Index: 1, + Phase: 1, + Message: "Chunker Done", + Total: 2, + }) + + logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) + if err != nil { + t.Fatalf("list logs: %v", err) + } + if len(logs) != 1 { + t.Fatalf("expected 1 component-progress row, got %d", len(logs)) + } + if stub.calls != 0 { + t.Fatalf("UpdateRunProgress calls = %d, want 0 (no document bound)", stub.calls) + } +} + +// TestDeriveDocumentProgress exercises every branch of the run-label derivation +// logic. The function is called from OnComponentProgress with a non-nil agg +// (guarded by the caller), so the nil case is documented as a known panic. +func TestDeriveDocumentProgress(t *testing.T) { + tests := []struct { + name string + agg *dao.TaskProgress + total int + wantRun string + wantProg float64 + }{ + { + name: "failed component → fail", + agg: &dao.TaskProgress{Failed: 1, Done: 0, Running: 0, Percent: 0}, + total: 5, + wantRun: string(entity.TaskStatusFail), + wantProg: 0.0, + }, + { + name: "all done → done", + agg: &dao.TaskProgress{Failed: 0, Done: 5, Running: 0, Percent: 100}, + total: 5, + wantRun: string(entity.TaskStatusDone), + wantProg: 1.0, + }, + { + name: "partial done → running", + agg: &dao.TaskProgress{Failed: 0, Done: 3, Running: 0, Percent: 60}, + total: 5, + wantRun: string(entity.TaskStatusRunning), + wantProg: 0.6, + }, + { + name: "running only → running", + agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 2, Percent: 0}, + total: 5, + wantRun: string(entity.TaskStatusRunning), + wantProg: 0.0, + }, + { + name: "nothing started → unstart", + agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 0, Percent: 0}, + total: 5, + wantRun: string(entity.TaskStatusUnstart), + wantProg: 0.0, + }, + { + name: "total zero, nothing done → done (0==0)", + agg: &dao.TaskProgress{Failed: 0, Done: 0, Running: 0, Percent: 0}, + total: 0, + wantRun: string(entity.TaskStatusDone), + wantProg: 0.0, + }, + { + name: "failed overrides done=total", + agg: &dao.TaskProgress{Failed: 1, Done: 5, Running: 0, Percent: 100}, + total: 5, + wantRun: string(entity.TaskStatusFail), + wantProg: 1.0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + prog, run := deriveDocumentProgress(tt.agg, tt.total) + if prog != tt.wantProg { + t.Errorf("progress = %v, want %v", prog, tt.wantProg) + } + if run != tt.wantRun { + t.Errorf("run = %q, want %q", run, tt.wantRun) + } + }) + } +} diff --git a/internal/ingestion/service/real_consumer_dataflow_test.go b/internal/ingestion/service/real_consumer_pipeline_test.go similarity index 72% rename from internal/ingestion/service/real_consumer_dataflow_test.go rename to internal/ingestion/service/real_consumer_pipeline_test.go index cc2a88c99e..820467be3e 100644 --- a/internal/ingestion/service/real_consumer_dataflow_test.go +++ b/internal/ingestion/service/real_consumer_pipeline_test.go @@ -1,27 +1,21 @@ -//go:build manual -// +build manual +//go:build integration package service import ( "context" "encoding/json" - "errors" "testing" "ragflow/internal/common" "ragflow/internal/dao" - "ragflow/internal/engine/nats" "ragflow/internal/entity" taskpkg "ragflow/internal/ingestion/task" "ragflow/internal/ingestion/testutil" ) -func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { - natsEngine := nats.NewNatsEngine("localhost", 4222) - if err := natsEngine.Init(); err != nil { - t.Fatalf("NATS Init: %v", err) - } +func TestRealConsumer_PipelineMessageRoutesToExecuteTask(t *testing.T) { + natsEngine := testutil.SetupNatsEngine(t) if err := natsEngine.InitConsumer("tasks.>"); err != nil { t.Fatalf("InitConsumer: %v", err) } @@ -44,7 +38,7 @@ func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { testutil.WithDocID("doc-q-1"), testutil.WithTaskID("ingest-q-1"), testutil.WithPipelineID("flow-queue-1"), - testutil.WithDocName("queue-dataflow.pdf"), + testutil.WithDocName("queue-pipeline.pdf"), ) // testutil.SeedTestData already created an IngestionTask with status RUNNING. @@ -85,46 +79,33 @@ func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { } ingestionTaskDAO := dao.NewIngestionTaskDAO() - task, err := ingestionTaskDAO.SetRunningByIngestor(taskMsg.TaskID) + _, err = ingestionTaskDAO.UpdateStatusIfCurrent(taskMsg.TaskID, common.CREATED, common.RUNNING) if err != nil { - if errors.Is(err, common.ErrTaskNotFound) { - t.Fatalf("task not found after publish: %s", taskMsg.TaskID) - } t.Fatalf("SetRunningByIngestor: %v", err) } + task, err := ingestionTaskDAO.GetByID(taskMsg.TaskID) + if err != nil || task == nil { + t.Fatalf("task not found after publish: %s", taskMsg.TaskID) + } if task.Status != common.RUNNING { t.Fatalf("task status after SetRunningByIngestor = %s, want %s", task.Status, common.RUNNING) } ingestor := NewIngestor("queue-test", 1, []string{"pdf"}) - var routedToDataflow bool - var progressEvents []string + var routedToPipeline bool taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), task, ) - var finalProgress float64 - taskCtx.ProgressFunc = func(prog float64, msg string) { - finalProgress = prog - progressEvents = append(progressEvents, msg) - } ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { - routedToDataflow = true - taskCtx.ProgressFunc(0.82, "mock queue dataflow start") - taskCtx.ProgressFunc(1.0, "mock queue dataflow done") + routedToPipeline = true return nil } ingestor.executeTask(taskCtx) - if !routedToDataflow { - t.Fatal("expected executeTask to route queue-consumed dataflow task to runDocumentTask") - } - if finalProgress != 1.0 { - t.Fatalf("finalProgress = %v, want 1.0", finalProgress) - } - if len(progressEvents) != 2 { - t.Fatalf("progressEvents = %v, want 2 events", progressEvents) + if !routedToPipeline { + t.Fatal("expected executeTask to route queue-consumed pipeline task to runDocumentTask") } if err := taskHandle.Ack(); err != nil { t.Fatalf("Ack: %v", err) diff --git a/internal/ingestion/service/run_task_test.go b/internal/ingestion/service/run_task_test.go new file mode 100644 index 0000000000..57609889cd --- /dev/null +++ b/internal/ingestion/service/run_task_test.go @@ -0,0 +1,274 @@ +package service + +import ( + "context" + "errors" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/ingestion/testutil" +) + +// TestRunTask_ContextCancelledBeforeCheckpoint: a cancelled context makes +// runTask return true (terminal: durably recorded cancel) immediately, without +// bumping the checkpoint or calling runDocumentTask. The task status is +// transitioned to STOPPED so it does not stay RUNNING forever. +func TestRunTask_ContextCancelledBeforeCheckpoint(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + var runDocCalled bool + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + runDocCalled = true + return nil + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + terminal := ingestor.runTask(ctx, &entity.IngestionTask{ + ID: taskID, DocumentID: "doc-1", DatasetID: "kb-1", + }) + + if !terminal { + t.Fatal("expected true (terminal: durably recorded cancel) on cancelled ctx") + } + if runDocCalled { + t.Fatal("expected runDocumentTask to be skipped on cancelled ctx") + } + // Checkpoint must not have been bumped — no log row should exist. + logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) + if err != nil { + t.Fatalf("list logs: %v", err) + } + if len(logs) != 0 { + t.Fatalf("expected 0 checkpoint rows (ctx cancelled before checkpoint), got %d", len(logs)) + } + // Task must be STOPPED, not left in RUNNING. + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.STOPPED { + t.Fatalf("task status = %s, want STOPPED", task.Status) + } +} + +// TestRunTask_CheckpointFailureMarksFailed: a corrupted run_count value is +// skipped by IncrementRunCount (it scans past unparseable rows). The task +// proceeds normally and completes. +func TestRunTask_CorruptedRunCountSkipped(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Seed a bad checkpoint: run_count is a string, not a number. + if err := db.Create(&entity.IngestionTaskLog{ + TaskID: taskID, + Checkpoint: entity.JSONMap{ + "run_count": "not-a-number", + }, + }).Error; err != nil { + t.Fatalf("insert bad log: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + var runDocCalled bool + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + runDocCalled = true + return nil + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if !terminal { + t.Fatal("expected true (terminal: completed despite corrupted run_count)") + } + if !runDocCalled { + t.Fatal("expected runDocumentTask to be called (bad run_count is skipped, not fatal)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.COMPLETED { + t.Fatalf("task status = %s, want COMPLETED", task.Status) + } +} + +// TestRunTask_RunDocumentTaskFailureMarksFailed: when runDocumentTask errors, +// runTask marks the task FAILED and returns durably-written. +func TestRunTask_RunDocumentTaskFailureMarksFailed(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return errors.New("boom") + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if !terminal { + t.Fatal("expected true (terminal: durably marked FAILED)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.FAILED { + t.Fatalf("task status = %s, want FAILED", task.Status) + } +} + +// TestRunTask_PipelineCancelledMarksStopped: when runDocumentTask returns +// context.Canceled (the pipeline detected a cancel signal), runTask treats +// it as a cancel, not a failure, and transitions the task to STOPPED. +func TestRunTask_PipelineCancelledMarksStopped(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Simulate RequestStop was already called: task in STOPPING. + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID). + Update("status", common.STOPPING).Error; err != nil { + t.Fatalf("set task STOPPING: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return context.Canceled + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.STOPPING, + }) + + if !terminal { + t.Fatal("expected true (terminal: durably marked STOPPED)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.STOPPED { + t.Fatalf("task status = %s, want STOPPED", task.Status) + } +} + +// TestRunTask_ComponentTimeoutMarksFailed: when runDocumentTask returns +// context.DeadlineExceeded (component Invoke hit its per-class timeout), +// runTask marks the task FAILED, not STOPPED. A component timeout is a +// system resource exhaustion, not a user-initiated cancellation. +func TestRunTask_ComponentTimeoutMarksFailed(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return context.DeadlineExceeded + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if !terminal { + t.Fatal("expected true (terminal: durably marked FAILED)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.FAILED { + t.Fatalf("task status = %s, want FAILED (DeadlineExceeded is a failure, not a cancel)", task.Status) + } +} + +// TestRunTask_MarkCompletedFailure: when runDocumentTask succeeds but +// MarkCompleted fails (status conflict), runTask returns false (non-terminal) +// so the message is Nacked for retry. +func TestRunTask_MarkCompletedFailure(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Pre-set the task COMPLETED so the RUNNING→COMPLETED transition inside + // MarkCompleted fails. + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID). + Update("status", common.COMPLETED).Error; err != nil { + t.Fatalf("set task COMPLETED: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return nil + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if terminal { + t.Fatal("expected false (non-terminal) on MarkCompleted failure") + } + + // Task must still be COMPLETED (MarkCompleted failed to transition it). + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.COMPLETED { + t.Fatalf("task status = %s, want COMPLETED (unchanged)", task.Status) + } +} + +// TestRunTask_SuccessfulCompletion: when everything succeeds, runTask returns +// true (terminal) and the task is COMPLETED. +func TestRunTask_SuccessfulCompletion(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + return nil + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if !terminal { + t.Fatal("expected true (terminal: durably completed)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.COMPLETED { + t.Fatalf("task status = %s, want COMPLETED", task.Status) + } +} diff --git a/internal/ingestion/task/chunk_index_writer.go b/internal/ingestion/task/chunk_index_writer.go new file mode 100644 index 0000000000..72cc653175 --- /dev/null +++ b/internal/ingestion/task/chunk_index_writer.go @@ -0,0 +1,70 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import "context" + +// InsertFunc is the signature of the chunk insertion backend (e.g. engine.InsertChunks). +type InsertFunc func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) + +// chunkIndexWriter batches chunks and writes them to the search engine in +// bulkSize-sized batches. Progress is reported every 128 batches. +type chunkIndexWriter struct { + insertFunc InsertFunc + baseName string + datasetID string + bulkSize int +} + +// newChunkIndexWriter creates a chunkIndexWriter. When bulkSize is <= 0 the +// entire chunk slice is sent in one call. +func newChunkIndexWriter( + insertFunc InsertFunc, + baseName string, + datasetID string, + bulkSize int, +) *chunkIndexWriter { + return &chunkIndexWriter{ + insertFunc: insertFunc, + baseName: baseName, + datasetID: datasetID, + bulkSize: bulkSize, + } +} + +// Write inserts chunks in batches. An empty or nil slice is forwarded to the +// backend as-is. +func (w *chunkIndexWriter) Write(ctx context.Context, chunks []map[string]any) error { + if len(chunks) == 0 { + _, err := w.insertFunc(ctx, chunks, w.baseName, w.datasetID) + return err + } + bulkSize := w.bulkSize + if bulkSize <= 0 { + bulkSize = len(chunks) + } + for b := 0; b < len(chunks); b += bulkSize { + end := b + bulkSize + if end > len(chunks) { + end = len(chunks) + } + if _, err := w.insertFunc(ctx, chunks[b:end], w.baseName, w.datasetID); err != nil { + return err + } + } + return nil +} diff --git a/internal/ingestion/task/chunk_index_writer_test.go b/internal/ingestion/task/chunk_index_writer_test.go new file mode 100644 index 0000000000..0d335d7ac4 --- /dev/null +++ b/internal/ingestion/task/chunk_index_writer_test.go @@ -0,0 +1,116 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "context" + "testing" +) + +func TestChunkIndexWriter_EmptyChunks(t *testing.T) { + called := false + w := newChunkIndexWriter( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + called = true + if len(chunks) != 0 { + t.Errorf("expected empty chunks, got %d", len(chunks)) + } + return nil, nil + }, + "test-base", + "kb-1", + 10, + ) + if err := w.Write(context.Background(), nil); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !called { + t.Fatal("insertFunc was not called for empty chunks") + } +} + +func TestChunkIndexWriter_SingleBatch(t *testing.T) { + var batchSizes []int + w := newChunkIndexWriter( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + batchSizes = append(batchSizes, len(chunks)) + if baseName != "test-base" { + t.Errorf("baseName = %q, want test-base", baseName) + } + if datasetID != "kb-1" { + t.Errorf("datasetID = %q, want kb-1", datasetID) + } + return nil, nil + }, + "test-base", + "kb-1", + 10, + ) + chunks := make([]map[string]any, 5) + if err := w.Write(context.Background(), chunks); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(batchSizes) != 1 { + t.Fatalf("expected 1 batch, got %d: %v", len(batchSizes), batchSizes) + } + if batchSizes[0] != 5 { + t.Fatalf("batch size = %d, want 5", batchSizes[0]) + } +} + +func TestChunkIndexWriter_MultipleBatches(t *testing.T) { + var batchSizes []int + w := newChunkIndexWriter( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + batchSizes = append(batchSizes, len(chunks)) + return nil, nil + }, + "base", + "kb-1", + 3, + ) + chunks := make([]map[string]any, 7) + if err := w.Write(context.Background(), chunks); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(batchSizes) != 3 { + t.Fatalf("expected 3 batches for 7 chunks with bulkSize=3, got %d: %v", len(batchSizes), batchSizes) + } + if batchSizes[0] != 3 || batchSizes[1] != 3 || batchSizes[2] != 1 { + t.Fatalf("batch sizes = %v, want [3,3,1]", batchSizes) + } +} + +func TestChunkIndexWriter_BulkSizeZero(t *testing.T) { + var lastBatchSize int + w := newChunkIndexWriter( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + lastBatchSize = len(chunks) + return nil, nil + }, + "base", + "kb-1", + 0, // bulkSize=0 → should use len(chunks) + ) + chunks := make([]map[string]any, 20) + if err := w.Write(context.Background(), chunks); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if lastBatchSize != 20 { + t.Fatalf("batch size = %d, want 20 (bulkSize=0 should degrade to len(chunks))", lastBatchSize) + } +} diff --git a/internal/ingestion/task/chunk_process.go b/internal/ingestion/task/chunk_process.go index 4b746c1157..4ba8d591a6 100644 --- a/internal/ingestion/task/chunk_process.go +++ b/internal/ingestion/task/chunk_process.go @@ -44,7 +44,10 @@ func TruncateTexts(texts []string, maxLength int) []string { } enc, err := tiktoken.GetEncoding("cl100k_base") if err != nil { - // Fallback: if tiktoken fails, return as-is + // Fallback: if tiktoken fails, return as-is. + // NOTE: this path cannot be triggered in unit tests because + // tiktoken.GetEncoding always succeeds in normal environments. + // The fallback is simple (make+copy) and verified by code review. result := make([]string, len(texts)) copy(result, texts) return result @@ -123,9 +126,9 @@ func GetEmbeddingTokenConsumption(output map[string]any) int { } } -// ProcessChunksForDataflow mutates chunks into the pre-index structure used by -// dataflow and returns merged metadata. -func ProcessChunksForDataflow( +// ProcessChunksForPipeline mutates chunks into the pre-index structure used by +// the pipeline and returns merged metadata. +func ProcessChunksForPipeline( chunks []map[string]any, docID string, kbID string, diff --git a/internal/ingestion/task/chunk_process_test.go b/internal/ingestion/task/chunk_process_test.go index 52cc1855e6..33b5451ea4 100644 --- a/internal/ingestion/task/chunk_process_test.go +++ b/internal/ingestion/task/chunk_process_test.go @@ -201,12 +201,12 @@ func TestRenameTextToContentWithWeight_NoTextKey(t *testing.T) { } // ============================================================================= -// ProcessChunksForDataflow — Python: processChunks() +// ProcessChunksForPipeline — Python: processChunks() // ============================================================================= -func TestProcessChunksForDataflow_SetsDocIDAndKBID(t *testing.T) { +func TestProcessChunksForPipeline_SetsDocIDAndKBID(t *testing.T) { chunks := []map[string]any{{"text": "hello world"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if chunks[0]["doc_id"] != "doc-1" { t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"]) @@ -220,18 +220,18 @@ func TestProcessChunksForDataflow_SetsDocIDAndKBID(t *testing.T) { } } -func TestProcessChunksForDataflow_SetsDocNameKwd(t *testing.T) { +func TestProcessChunksForPipeline_SetsDocNameKwd(t *testing.T) { chunks := []map[string]any{{"text": "hello"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if chunks[0]["docnm_kwd"] != "test-doc.pdf" { t.Errorf("docnm_kwd = %q, want \"test-doc.pdf\"", chunks[0]["docnm_kwd"]) } } -func TestProcessChunksForDataflow_SetsTimeFields(t *testing.T) { +func TestProcessChunksForPipeline_SetsTimeFields(t *testing.T) { now := time.Now() chunks := []map[string]any{{"text": "hello"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", now) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", now) if timeStr, ok := chunks[0]["create_time"].(string); ok { if timeStr != now.Format("2006-01-02 15:04:05") { @@ -250,31 +250,31 @@ func TestProcessChunksForDataflow_SetsTimeFields(t *testing.T) { } } -func TestProcessChunksForDataflow_GeneratesID(t *testing.T) { +func TestProcessChunksForPipeline_GeneratesID(t *testing.T) { chunks := []map[string]any{{"text": "hello"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) id, ok := chunks[0]["id"].(string) if !ok || id == "" { t.Errorf("id should be non-empty string, got %v", chunks[0]["id"]) } } -func TestProcessChunksForDataflow_NoPanicOnListText(t *testing.T) { +func TestProcessChunksForPipeline_NoPanicOnListText(t *testing.T) { chunks := []map[string]any{{"text": []any{"bad-shape"}}} - res := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + res := ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if res == nil { t.Errorf("should return valid result") } } -func TestProcessChunksForDataflow_RemovesInternalPipelineFields(t *testing.T) { +func TestProcessChunksForPipeline_RemovesInternalPipelineFields(t *testing.T) { chunks := []map[string]any{{ "text": "hello", "_pdf_positions": []any{[]any{0, 1, 2, 3, 4}}, "image": "data:image/png;base64,abc", }} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if _, exists := chunks[0]["_pdf_positions"]; exists { t.Fatalf("_pdf_positions should be removed before indexing: %v", chunks[0]["_pdf_positions"]) } @@ -283,17 +283,17 @@ func TestProcessChunksForDataflow_RemovesInternalPipelineFields(t *testing.T) { } } -func TestProcessChunksForDataflow_PreservesExistingID(t *testing.T) { +func TestProcessChunksForPipeline_PreservesExistingID(t *testing.T) { chunks := []map[string]any{{"text": "hello", "id": "existing-id"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if chunks[0]["id"] != "existing-id" { t.Errorf("existing id should be preserved, got %q", chunks[0]["id"]) } } -func TestProcessChunksForDataflow_QuestionsProcessing(t *testing.T) { +func TestProcessChunksForPipeline_QuestionsProcessing(t *testing.T) { chunks := []map[string]any{{"text": "hello", "questions": "Q1\nQ2\nQ3"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if _, exists := chunks[0]["questions"]; exists { t.Error("questions key should be removed") @@ -310,9 +310,9 @@ func TestProcessChunksForDataflow_QuestionsProcessing(t *testing.T) { } } -func TestProcessChunksForDataflow_KeywordsProcessing(t *testing.T) { +func TestProcessChunksForPipeline_KeywordsProcessing(t *testing.T) { chunks := []map[string]any{{"text": "hello", "keywords": "kw1,kw2;kw3"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if _, exists := chunks[0]["keywords"]; exists { t.Error("keywords key should be removed") @@ -326,9 +326,9 @@ func TestProcessChunksForDataflow_KeywordsProcessing(t *testing.T) { } } -func TestProcessChunksForDataflow_SummaryProcessing(t *testing.T) { +func TestProcessChunksForPipeline_SummaryProcessing(t *testing.T) { chunks := []map[string]any{{"text": "hello", "summary": "This is a summary."}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if _, exists := chunks[0]["summary"]; exists { t.Error("summary key should be removed") @@ -341,9 +341,9 @@ func TestProcessChunksForDataflow_SummaryProcessing(t *testing.T) { } } -func TestProcessChunksForDataflow_TextRenamed(t *testing.T) { +func TestProcessChunksForPipeline_TextRenamed(t *testing.T) { chunks := []map[string]any{{"text": "hello world"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if _, exists := chunks[0]["text"]; exists { t.Error("text key should be removed") @@ -353,9 +353,9 @@ func TestProcessChunksForDataflow_TextRenamed(t *testing.T) { } } -func TestProcessChunksForDataflow_PreservesContentWithWeight(t *testing.T) { +func TestProcessChunksForPipeline_PreservesContentWithWeight(t *testing.T) { chunks := []map[string]any{{"content_with_weight": "already set", "text": "hello"}} - _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + _ = ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) if chunks[0]["content_with_weight"] != "already set" { t.Errorf("content_with_weight = %q, want \"already set\"", chunks[0]["content_with_weight"]) } diff --git a/internal/ingestion/task/chunk_utils.go b/internal/ingestion/task/chunk_utils.go index 1476d9daf8..f03a81d658 100644 --- a/internal/ingestion/task/chunk_utils.go +++ b/internal/ingestion/task/chunk_utils.go @@ -102,10 +102,10 @@ func deepCopyChunks(chunks []map[string]any) []map[string]any { return out } -// PrepareTextsForDataflowEmbedding extracts texts for embedding from chunks. +// PrepareTextsForPipelineEmbedding extracts texts for embedding from chunks. // Priority: questions > summary > text. // Mirrors Python: EmbeddingUtils.prepare_texts_for_dataflow_embedding() -func PrepareTextsForDataflowEmbedding(chunks []map[string]any) []string { +func PrepareTextsForPipelineEmbedding(chunks []map[string]any) []string { if chunks == nil { return nil } diff --git a/internal/ingestion/task/chunk_utils_test.go b/internal/ingestion/task/chunk_utils_test.go index 8a421e980e..4b26306190 100644 --- a/internal/ingestion/task/chunk_utils_test.go +++ b/internal/ingestion/task/chunk_utils_test.go @@ -219,14 +219,14 @@ func TestNormalizeChunks_NilInput(t *testing.T) { } // ============================================================================= -// PrepareTextsForDataflowEmbedding +// PrepareTextsForPipelineEmbedding // ============================================================================= func TestPrepareTexts_QuestionsPriority(t *testing.T) { chunks := []map[string]any{ {"questions": "Q1\nQ2", "summary": "a summary", "text": "plain text"}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if len(result) != 1 { t.Fatalf("len = %d, want 1", len(result)) } @@ -239,7 +239,7 @@ func TestPrepareTexts_SummaryFallback(t *testing.T) { chunks := []map[string]any{ {"summary": "a summary", "text": "plain text"}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if result[0] != "a summary" { t.Errorf("summary should be used when no questions: got %q", result[0]) } @@ -249,7 +249,7 @@ func TestPrepareTexts_TextFallback(t *testing.T) { chunks := []map[string]any{ {"text": "plain text"}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if result[0] != "plain text" { t.Errorf("text should be used when no questions/summary: got %q", result[0]) } @@ -259,7 +259,7 @@ func TestPrepareTexts_EmptyStringFallback(t *testing.T) { chunks := []map[string]any{ {"text": ""}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if len(result) > 0 { t.Errorf("expected empty string, got %q", result[0]) } @@ -271,7 +271,7 @@ func TestPrepareTexts_MultipleChunks(t *testing.T) { {"summary": "S2", "text": "t2"}, {"text": "t3"}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if len(result) != 3 { t.Fatalf("len = %d, want 3", len(result)) } @@ -287,14 +287,14 @@ func TestPrepareTexts_MultipleChunks(t *testing.T) { } func TestPrepareTexts_NilChunks(t *testing.T) { - result := PrepareTextsForDataflowEmbedding(nil) + result := PrepareTextsForPipelineEmbedding(nil) if result != nil { t.Errorf("expected nil for nil chunks, got %v", result) } } func TestPrepareTexts_EmptyChunks(t *testing.T) { - result := PrepareTextsForDataflowEmbedding([]map[string]any{}) + result := PrepareTextsForPipelineEmbedding([]map[string]any{}) if len(result) != 0 { t.Errorf("expected empty slice, got len=%d", len(result)) } @@ -304,7 +304,7 @@ func TestPrepareTexts_MissingTextKey(t *testing.T) { chunks := []map[string]any{ {"other_key": "value"}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if len(result) > 0 { t.Errorf("expected empty string for missing text key, got %q", result[0]) } @@ -314,7 +314,7 @@ func TestPrepareTexts_NoPanicOnListText(t *testing.T) { chunks := []map[string]any{ {"text": []any{"bad-shape"}}, } - result := PrepareTextsForDataflowEmbedding(chunks) + result := PrepareTextsForPipelineEmbedding(chunks) if len(result) > 0 { t.Errorf("expected empty string for missing text key, got %q", result[0]) } @@ -327,3 +327,43 @@ func TestMustGetChunkTextString_NoPanicOnStringSlice(t *testing.T) { } } +func TestGetEmbeddingTokenConsumption_Int(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: 42} + result := GetEmbeddingTokenConsumption(input) + if result != 42 { + t.Errorf("got %d, want 42", result) + } +} +func TestGetEmbeddingTokenConsumption_Float64(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: float64(42)} + result := GetEmbeddingTokenConsumption(input) + if result != 42 { + t.Errorf("got %d, want 42", result) + } +} +func TestGetEmbeddingTokenConsumption_MissingKey(t *testing.T) { + result := GetEmbeddingTokenConsumption(map[string]any{}) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} +func TestGetEmbeddingTokenConsumption_NilMap(t *testing.T) { + result := GetEmbeddingTokenConsumption(nil) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} +func TestGetEmbeddingTokenConsumption_WrongType(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: "not a number"} + result := GetEmbeddingTokenConsumption(input) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} +func TestGetEmbeddingTokenConsumption_Zero(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: 0} + result := GetEmbeddingTokenConsumption(input) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} diff --git a/internal/ingestion/task/constants.go b/internal/ingestion/task/constants.go index fe6adefd41..7d9eefc6af 100644 --- a/internal/ingestion/task/constants.go +++ b/internal/ingestion/task/constants.go @@ -18,7 +18,7 @@ package task // Special doc_id values used in task messages (mirrors Python constants). const ( - // CANVAS_DEBUG_DOC_ID marks a canvas debug dataflow task — no persistence. + // CANVAS_DEBUG_DOC_ID marks a canvas debug pipeline task — no persistence. CANVAS_DEBUG_DOC_ID = "dataflow_x" // GRAPH_RAPTOR_FAKE_DOC_ID is the fake doc_id used for RAPTOR-generated chunks. diff --git a/internal/ingestion/task/dataflow_service.go b/internal/ingestion/task/dataflow_service.go deleted file mode 100644 index af79a8c13d..0000000000 --- a/internal/ingestion/task/dataflow_service.go +++ /dev/null @@ -1,525 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package task - -import ( - "context" - "encoding/json" - "fmt" - componentpkg "ragflow/internal/ingestion/component" - "ragflow/internal/utility" - "sort" - "strings" - "time" - - "ragflow/internal/common" - "ragflow/internal/dao" - "ragflow/internal/engine" - "ragflow/internal/entity" - "ragflow/internal/entity/models" - pipelinepkg "ragflow/internal/ingestion/pipeline" - "ragflow/internal/service" -) - -type embedder struct { - model *models.EmbeddingModel -} - -func (e *embedder) MaxTokens() int { - if e == nil || e.model == nil { - return 0 - } - return e.model.MaxTokens -} - -func (e *embedder) Encode(texts []string) ([]componentpkg.EmbeddingResult, error) { - config := &models.EmbeddingConfig{Dimension: 0} - embeds, err := e.model.ModelDriver.Embed(e.model.ModelName, texts, e.model.APIConfig, config) - if err != nil { - return nil, err - } - vecs := make([]componentpkg.EmbeddingResult, len(embeds)) - for i, v := range embeds { - vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: v.TokenCount} - } - return vecs, nil -} - -type ProgressFunc func(prog float64, msg string) - -type docService interface { - UpdateDocument(id string, req *service.UpdateDocumentRequest) error - GetDocumentMetadataByID(docID string) (map[string]any, error) - SetDocumentMetadata(docID string, meta map[string]any) error -} - -type chunkCounter interface { - IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error -} - -type defaultDocService struct{} -type defaultChunkCounter struct{} - -func (d *defaultDocService) UpdateDocument(id string, req *service.UpdateDocumentRequest) error { - return service.NewDocumentService().UpdateDocument(id, req) -} - -func (d *defaultDocService) GetDocumentMetadataByID(docID string) (map[string]any, error) { - return service.NewDocumentService().GetDocumentMetadataByID(docID) -} - -func (d *defaultDocService) SetDocumentMetadata(docID string, meta map[string]any) error { - return service.NewDocumentService().SetDocumentMetadata(docID, meta) -} - -func (d *defaultChunkCounter) IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { - return service.NewDocumentService().IncrementChunkNum(docID, kbID, chunkNum, tokenConsumption, duration) -} - -type PipelineExecutor struct { - taskCtx *TaskContext - dataflowID string - embeddingBatchSize int - docBulkSize int - progressFunc ProgressFunc - - docSvc docService - chunkCounter chunkCounter - insertChunksFunc func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) - logCreateFunc func(log *entity.PipelineOperationLog) error - loadDSLFunc func(ctx context.Context, dataflowID string) (string, string, error) - runPipelineFunc func(ctx context.Context, dsl string) (map[string]any, string, error) -} - -// newEmbedderResolver builds the production embedder resolver used by the -// Tokenizer component. It honors an explicit embedding-model id (from the -// Tokenizer's setups) and falls back to the dataset's configured embd_id when -// none is given. Kept as a constructor over injectable deps so the resolution -// logic stays unit-testable without a live model provider / DB. -func newEmbedderResolver( - getEmbeddingModel func(tenantID, embdID string) (*models.EmbeddingModel, error), - getKnowledgebaseByID func(kbID string) (*entity.Knowledgebase, error), -) componentpkg.EmbedderResolver { - return func(tenantID, kbID, embeddingModel string) (componentpkg.Embedder, error) { - embdID := strings.TrimSpace(embeddingModel) - if embdID == "" { - if strings.TrimSpace(kbID) == "" { - return nil, fmt.Errorf("embedding requested but neither embedding_model nor kb_id provided") - } - kb, err := getKnowledgebaseByID(kbID) - if err != nil { - return nil, err - } - if kb == nil || strings.TrimSpace(kb.EmbdID) == "" { - return nil, fmt.Errorf("embedding requested but dataset has no embd_id configured") - } - embdID = kb.EmbdID - } - model, err := getEmbeddingModel(tenantID, embdID) - if err != nil { - return nil, err - } - return &embedder{model: model}, nil - } -} - -// init wires the production embedder resolver into the component package. The -// component package must not import internal/service (dependency direction), -// so the concrete resolver is injected here — the task package is the -// composition root for ingestion runs. -func init() { - componentpkg.DefaultEmbedderResolver = newEmbedderResolver( - service.NewModelProviderService().GetEmbeddingModel, - dao.NewKnowledgebaseDAO().GetByID, - ) -} - -func validateDataflowTaskContext(taskCtx *TaskContext) error { - if taskCtx == nil { - return fmt.Errorf("dataflow service: nil task context") - } - if taskCtx.Doc.ID == "" { - return fmt.Errorf("dataflow service: empty document id") - } - if taskCtx.Doc.KbID == "" { - return fmt.Errorf("dataflow service: empty document knowledgebase id") - } - if taskCtx.Doc.Name == nil || *taskCtx.Doc.Name == "" { - return fmt.Errorf("dataflow service: empty document name") - } - if taskCtx.KB.ID == "" { - return fmt.Errorf("dataflow service: empty knowledgebase id") - } - if taskCtx.Tenant.ID == "" { - return fmt.Errorf("dataflow service: empty tenant id") - } - return nil -} - -func NewDataflowService( - taskCtx *TaskContext, - dataflowID string, - embeddingBatchSize int, - docBulkSize int, -) (*PipelineExecutor, error) { - if err := validateDataflowTaskContext(taskCtx); err != nil { - return nil, err - } - if strings.TrimSpace(dataflowID) == "" { - return nil, fmt.Errorf("dataflow service: empty dataflow id") - } - progressFn := func(prog float64, msg string) {} - if taskCtx != nil && taskCtx.ProgressFunc != nil { - progressFn = taskCtx.ProgressFunc - } - svc := &PipelineExecutor{ - taskCtx: taskCtx, - dataflowID: dataflowID, - embeddingBatchSize: embeddingBatchSize, - docBulkSize: docBulkSize, - progressFunc: progressFn, - docSvc: &defaultDocService{}, - chunkCounter: &defaultChunkCounter{}, - insertChunksFunc: func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { - return engine.Get().InsertChunks(ctx, chunks, baseName, datasetID) - }, - logCreateFunc: dao.NewPipelineOperationLogDAO().Create, - } - svc.loadDSLFunc = svc.defaultLoadDSL - svc.runPipelineFunc = svc.defaultRunPipeline - return svc, nil -} - -func (s *PipelineExecutor) WithProgressFunc(fn ProgressFunc) *PipelineExecutor { - s.progressFunc = fn - return s -} - -func (s *PipelineExecutor) WithInsertChunksFunc(f func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error)) *PipelineExecutor { - s.insertChunksFunc = f - return s -} - -func (s *PipelineExecutor) WithLogCreateFunc(f func(log *entity.PipelineOperationLog) error) *PipelineExecutor { - s.logCreateFunc = f - return s -} - -func (s *PipelineExecutor) WithDocService(d docService) *PipelineExecutor { - s.docSvc = d - return s -} - -func (s *PipelineExecutor) WithChunkCounter(c chunkCounter) *PipelineExecutor { - s.chunkCounter = c - return s -} - -func (s *PipelineExecutor) WithLoadDSLFunc(f func(ctx context.Context, dataflowID string) (string, string, error)) *PipelineExecutor { - s.loadDSLFunc = f - return s -} - -func (s *PipelineExecutor) WithRunPipelineFunc(f func(ctx context.Context, dsl string) (map[string]any, string, error)) *PipelineExecutor { - s.runPipelineFunc = f - return s -} - -func (s *PipelineExecutor) KB() *entity.Knowledgebase { return &s.taskCtx.KB } -func (s *PipelineExecutor) Doc() *entity.Document { return &s.taskCtx.Doc } -func (s *PipelineExecutor) Tenant() *entity.Tenant { return &s.taskCtx.Tenant } - -func (s *PipelineExecutor) Execute(ctx context.Context) error { - if err := ctx.Err(); err != nil { - return err - } - - dsl, correctedID, err := s.loadDSLFunc(ctx, s.dataflowID) - if err != nil { - return err - } - if correctedID != "" { - s.dataflowID = correctedID - } - - pipelineOutput, pipelineDSL, err := s.runPipelineFunc(ctx, dsl) - if err != nil { - return err - } - - if s.taskCtx.Doc.ID == CANVAS_DEBUG_DOC_ID { - s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") - return nil - } - - if err := s.processOutput(ctx, pipelineOutput); err != nil { - return err - } - - if pipelineDSL != "" { - s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") - } - return nil -} - -func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map[string]any) error { - taskStart := time.Now() - if pipelineOutput == nil { - return nil - } - if err := ctx.Err(); err != nil { - return err - } - - chunks := s.normalizeChunks(pipelineOutput) - if chunks == nil { - return nil - } - - embeddingTokenConsumption := GetEmbeddingTokenConsumption(pipelineOutput) - - metadata := s.processChunks(chunks) - - if len(metadata) > 0 { - if err := s.updateDocumentMetadata(s.taskCtx.Doc.ID, metadata); err != nil { - common.Warn(fmt.Sprintf("failed to update document metadata: %v", err)) - } - } - - indexStart := time.Now() - s.progress(0.82, "[DOC Engine]:\nStart to index...") - if err := s.insertChunks(ctx, chunks); err != nil { - return err - } - - if err := s.incrementChunkNum(s.taskCtx.Doc.ID, s.taskCtx.Doc.KbID, len(chunks), embeddingTokenConsumption, 0); err != nil { - common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err)) - } - indexDuration := time.Since(indexStart).Seconds() - taskDuration := time.Since(taskStart).Seconds() - s.progress(1.0, fmt.Sprintf("Indexing done (%.2fs). Task done (%.2fs)", indexDuration, taskDuration)) - - return nil -} - -func (s *PipelineExecutor) normalizeChunks(output map[string]any) []map[string]any { - return NormalizeChunks(output) -} - -func (s *PipelineExecutor) processChunks(chunks []map[string]any) map[string]any { - return ProcessChunksForDataflow( - chunks, - s.taskCtx.Doc.ID, - s.taskCtx.Doc.KbID, - *s.taskCtx.Doc.Name, - time.Now(), - ) -} - -func (s *PipelineExecutor) insertChunks(ctx context.Context, chunks []map[string]any) error { - baseName := fmt.Sprintf("ragflow_%s", s.taskCtx.Tenant.ID) - if len(chunks) == 0 { - _, err := s.insertChunksFunc(ctx, chunks, baseName, s.taskCtx.Doc.KbID) - return err - } - bulkSize := s.docBulkSize - if bulkSize <= 0 { - bulkSize = len(chunks) - } - for b := 0; b < len(chunks); b += bulkSize { - end := b + bulkSize - if end > len(chunks) { - end = len(chunks) - } - if _, err := s.insertChunksFunc(ctx, chunks[b:end], baseName, s.taskCtx.Doc.KbID); err != nil { - return err - } - if (b/bulkSize)%128 == 0 { - s.progress(0.8+0.1*float64(b+1)/float64(len(chunks)), "") - } - } - return nil -} - -func (s *PipelineExecutor) updateDocumentMetadata(docID string, metadata map[string]any) error { - if len(metadata) == 0 { - return nil - } - existing, err := s.docSvc.GetDocumentMetadataByID(docID) - if err != nil { - existing = make(map[string]any) - } - for k, v := range metadata { - if _, exists := existing[k]; !exists { - existing[k] = v - } - } - return s.docSvc.SetDocumentMetadata(docID, existing) -} - -func (s *PipelineExecutor) recordPipelineLog(docID, dsl, status string) { - var dslMap entity.JSONMap - if err := json.Unmarshal([]byte(dsl), &dslMap); err != nil { - dslMap = entity.JSONMap{"raw": dsl} - } - log := &entity.PipelineOperationLog{ - ID: utility.GenerateUUID(), - TenantID: s.Tenant().ID, - KbID: s.KB().ID, - DocumentID: docID, - PipelineID: &s.dataflowID, - TaskType: string(entity.PipelineTaskTypeParse), - DSL: dslMap, - ParserID: s.taskCtx.Doc.ParserID, - DocumentName: *s.Doc().Name, - DocumentSuffix: s.taskCtx.Doc.Suffix, - DocumentType: s.taskCtx.Doc.Type, - SourceFrom: s.taskCtx.Doc.SourceType, - OperationStatus: status, - } - if err := s.logCreateFunc(log); err != nil { - common.Warn(fmt.Sprintf("failed to record pipeline log: %v", err)) - } -} - -func (s *PipelineExecutor) incrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { - if s.chunkCounter == nil { - return fmt.Errorf("dataflow service: chunk counter is nil") - } - return s.chunkCounter.IncrementChunkNum(docID, kbID, chunkNum, tokenConsumption, duration) -} - -func (s *PipelineExecutor) progress(prog float64, msg string) { - if s.progressFunc != nil { - s.progressFunc(prog, msg) - } -} - -func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, dataflowID string) (string, string, error) { - if s == nil || s.taskCtx == nil { - return "", "", fmt.Errorf("dataflow service: nil task context") - } - if dataflowID == "" { - return "", "", fmt.Errorf("dataflow service: empty dataflow id") - } - canvas, err := dao.NewUserCanvasDAO().GetByID(dataflowID) - if err != nil { - return "", "", fmt.Errorf("load dataflow canvas %s: %w", dataflowID, err) - } - - canvasTitle := "" - if canvas.Title != nil { - canvasTitle = *canvas.Title - } - common.Info(fmt.Sprintf("load dataflow canvas %s, name %s", dataflowID, canvasTitle)) - - raw, err := json.Marshal(canvas.DSL) - if err != nil { - return "", "", fmt.Errorf("marshal canvas dsl %s: %w", dataflowID, err) - } - return string(raw), dataflowID, nil -} - -func (s *PipelineExecutor) defaultRunPipeline(ctx context.Context, dsl string) (map[string]any, string, error) { - if s == nil || s.taskCtx == nil { - return nil, dsl, fmt.Errorf("dataflow service: nil task context") - } - - // Use doc ID as pipeline ID if available, otherwise a placeholder - pipelineID := "pipeline_" + s.taskCtx.Doc.ID - if s.taskCtx.IngestionTask != nil && s.taskCtx.IngestionTask.ID != "" { - pipelineID = s.taskCtx.IngestionTask.ID - } - pipe, err := pipelinepkg.NewPipelineFromDSL([]byte(dsl), pipelineID) - if err != nil { - return nil, dsl, fmt.Errorf("compile pipeline dsl: %w", err) - } - inputs := map[string]any{} - if s.taskCtx.Doc.ID != "" { - inputs["doc_id"] = s.taskCtx.Doc.ID - } - if s.taskCtx.File != nil { - inputs["file"] = s.taskCtx.File - } - inputs["tenant_id"] = s.taskCtx.Tenant.ID - inputs["kb_id"] = s.taskCtx.KB.ID - - output, err := pipe.Run(ctx, inputs) - if err != nil { - return nil, dsl, err - } - payload, err := extractDataflowPipelinePayload(dsl, output) - if err != nil { - return nil, dsl, err - } - return payload, dsl, nil -} - -func extractDataflowPipelinePayload(dsl string, out map[string]any) (map[string]any, error) { - if out == nil { - return nil, nil - } - if _, ok := out["output_format"]; ok { - return out, nil - } - terminalIDs, err := terminalComponentIDsFromDSL([]byte(dsl)) - if err != nil { - return nil, err - } - if len(terminalIDs) != 1 { - return nil, fmt.Errorf("dataflow pipeline requires exactly 1 terminal, got %d: %v", len(terminalIDs), terminalIDs) - } - payload, ok := out[terminalIDs[0]].(map[string]any) - if !ok { - return nil, fmt.Errorf("run output missing terminal payload %q", terminalIDs[0]) - } - return payload, nil -} - -func terminalComponentIDsFromDSL(raw []byte) ([]string, error) { - var tpl map[string]any - if err := json.Unmarshal(raw, &tpl); err != nil { - return nil, fmt.Errorf("unmarshal dataflow dsl: %w", err) - } - root := tpl - if nested, ok := tpl["dsl"].(map[string]any); ok { - root = nested - } - components, ok := root["components"].(map[string]any) - if !ok { - return nil, fmt.Errorf("dataflow dsl missing components map") - } - terminals := make([]string, 0, len(components)) - for id, rawComp := range components { - comp, ok := rawComp.(map[string]any) - if !ok { - return nil, fmt.Errorf("component %q has invalid type %T", id, rawComp) - } - switch downstream := comp["downstream"].(type) { - case nil: - terminals = append(terminals, id) - case []any: - if len(downstream) == 0 { - terminals = append(terminals, id) - } - default: - // Non-slice downstream means the component is connected; ignore it here. - } - } - sort.Strings(terminals) - return terminals, nil -} diff --git a/internal/ingestion/task/dataflow_service_test.go b/internal/ingestion/task/dataflow_service_test.go deleted file mode 100644 index 4f9102f094..0000000000 --- a/internal/ingestion/task/dataflow_service_test.go +++ /dev/null @@ -1,754 +0,0 @@ -package task - -import ( - "context" - "encoding/json" - "strings" - "testing" - "time" - - "github.com/glebarez/sqlite" - "gorm.io/gorm" - "gorm.io/gorm/logger" - - "ragflow/internal/dao" - "ragflow/internal/entity" - "ragflow/internal/entity/models" - "ragflow/internal/service" -) - -// ============================================================================= -// Test helpers -// ============================================================================= - -func strPtr(s string) *string { return &s } - -func makeTaskCtx() *TaskContext { - return &TaskContext{ - IngestionTask: &entity.IngestionTask{ - ID: "task-1", - DocumentID: "doc-1", - }, - Doc: entity.Document{ - ID: "doc-1", - KbID: "kb-1", - Name: strPtr("test-doc.pdf"), - Suffix: ".pdf", - Type: "pdf", - }, - KB: entity.Knowledgebase{ - ID: "kb-1", - TenantID: "tenant-1", - EmbdID: "embd-1", - }, - Tenant: entity.Tenant{ - ID: "tenant-1", - }, - ProgressFunc: func(prog float64, msg string) {}, - } -} - -func setupDataflowServiceTestDB(t *testing.T) func() { - t.Helper() - db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) - if err != nil { - t.Fatalf("open sqlite: %v", err) - } - if err := db.AutoMigrate(&entity.UserCanvas{}, &entity.PipelineOperationLog{}); err != nil { - t.Fatalf("auto-migrate sqlite: %v", err) - } - origDB := dao.DB - dao.DB = db - return func() { dao.DB = origDB } -} - -func mustNewDataflowService(t *testing.T, taskCtx *TaskContext, dataflowID string, embeddingBatchSize int, docBulkSize int) *PipelineExecutor { - t.Helper() - svc, err := NewDataflowService(taskCtx, dataflowID, embeddingBatchSize, docBulkSize) - if err != nil { - t.Fatalf("NewDataflowService: %v", err) - } - return svc -} - -func TestDataflowService_DefaultLoadDSL_UsesUserCanvas(t *testing.T) { - cleanup := setupDataflowServiceTestDB(t) - defer cleanup() - - dslMap := entity.JSONMap{"dsl": map[string]any{"graph": map[string]any{"nodes": []any{}, "edges": []any{}}}} - title := "title 1" - if err := dao.NewUserCanvasDAO().Create(&entity.UserCanvas{Title: &title, ID: "canvas-1", UserID: "u1", Permission: "me", CanvasCategory: "agent_canvas", DSL: dslMap}); err != nil { - t.Fatalf("create user canvas: %v", err) - } - - ctx := makeTaskCtx() - svc := mustNewDataflowService(t, ctx, "canvas-1", 0, 0) - gotDSL, correctedID, err := svc.loadDSLFunc(context.Background(), "canvas-1") - if err != nil { - t.Fatalf("loadDSLFunc: %v", err) - } - if correctedID != "canvas-1" { - t.Fatalf("correctedID = %q, want canvas-1", correctedID) - } - var decoded map[string]any - if err := json.Unmarshal([]byte(gotDSL), &decoded); err != nil { - t.Fatalf("unmarshal dsl: %v", err) - } - if _, ok := decoded["dsl"].(map[string]any); !ok { - t.Fatalf("decoded dsl = %v, want top-level dsl map", decoded) - } -} - -// ============================================================================= -// NewDataflowService — constructor -// ============================================================================= - -func TestNewDataflowService_Basic(t *testing.T) { - svc, err := NewDataflowService(makeTaskCtx(), "flow-1", 0, 0) - if err != nil { - t.Fatalf("NewDataflowService: %v", err) - } - if svc == nil { - t.Fatal("NewDataflowService returned nil") - } - if svc.taskCtx == nil { - t.Error("taskCtx should not be nil") - } - if svc.docSvc == nil || svc.chunkCounter == nil || svc.insertChunksFunc == nil || svc.logCreateFunc == nil || svc.loadDSLFunc == nil || svc.runPipelineFunc == nil { - t.Fatal("expected production dependencies to be fully initialized") - } -} - -func TestNewDataflowService_RejectsNilTaskContext(t *testing.T) { - _, err := NewDataflowService(nil, "flow-1", 0, 0) - if err == nil { - t.Fatal("expected error for nil task context") - } -} - -func TestNewDataflowService_RejectsEmptyDataflowID(t *testing.T) { - _, err := NewDataflowService(makeTaskCtx(), "", 0, 0) - if err == nil { - t.Fatal("expected error for empty dataflow id") - } -} - -func TestNewDataflowService_RejectsIncompleteTaskContext(t *testing.T) { - tests := []struct { - name string - mutate func(*TaskContext) - }{ - {name: "missing doc id", mutate: func(ctx *TaskContext) { ctx.Doc.ID = "" }}, - {name: "missing kb id", mutate: func(ctx *TaskContext) { ctx.Doc.KbID = "" }}, - {name: "missing doc name", mutate: func(ctx *TaskContext) { ctx.Doc.Name = nil }}, - {name: "missing knowledgebase id", mutate: func(ctx *TaskContext) { ctx.KB.ID = "" }}, - {name: "missing tenant id", mutate: func(ctx *TaskContext) { ctx.Tenant.ID = "" }}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctx := makeTaskCtx() - tt.mutate(ctx) - _, err := NewDataflowService(ctx, "flow-1", 0, 0) - if err == nil { - t.Fatal("expected validation error") - } - }) - } -} - -func TestNewDataflowService_CustomBatchSizes(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 64, 128) - if svc.embeddingBatchSize != 64 { - t.Errorf("embeddingBatchSize = %d, want 64", svc.embeddingBatchSize) - } - if svc.docBulkSize != 128 { - t.Errorf("docBulkSize = %d, want 128", svc.docBulkSize) - } -} - -func TestNewDataflowService_WithProgressFunc(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithProgressFunc(func(prog float64, msg string) {}) - if svc.progressFunc == nil { - t.Error("progressFunc should be set via WithProgressFunc") - } -} - -func TestNewDataflowService_DataflowID(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "my-flow-id", 0, 0) - if svc.dataflowID != "my-flow-id" { - t.Errorf("dataflowID = %q, want %q", svc.dataflowID, "my-flow-id") - } -} - -func TestKB_Doc_Tenant_Accessors(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - if svc.KB().ID != "kb-1" { - t.Errorf("KB().ID = %q, want \"kb-1\"", svc.KB().ID) - } - if svc.Doc().ID != "doc-1" { - t.Errorf("Doc().ID = %q, want \"doc-1\"", svc.Doc().ID) - } - if svc.Tenant().ID != "tenant-1" { - t.Errorf("Tenant().ID = %q, want \"tenant-1\"", svc.Tenant().ID) - } -} - -// ============================================================================= -// processChunks -// ============================================================================= - -func TestDataflowService_ProcessChunks_WrapsProcessChunksForDataflow(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - chunks := []map[string]any{{"text": "hello world"}} - meta := svc.processChunks(chunks) - - // Verify the wrapper method works correctly and chunks are processed - if chunks[0]["doc_id"] != "doc-1" { - t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"]) - } - if meta != nil { - // No need to verify the detailed content of meta as ProcessChunksForDataflow already has comprehensive tests - } -} - -// ============================================================================= -// progress -// ============================================================================= - -func TestProgress_CallsCallback(t *testing.T) { - var called bool - var lastProg float64 - var lastMsg string - - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - svc.WithProgressFunc(func(prog float64, msg string) { - called = true - lastProg = prog - lastMsg = msg - }) - svc.progress(0.82, "Start to embedding...") - - if !called { - t.Error("progress callback was not called") - } - if lastProg != 0.82 { - t.Errorf("prog = %f, want 0.82", lastProg) - } - if lastMsg != "Start to embedding..." { - t.Errorf("msg = %q, want \"Start to embedding...\"", lastMsg) - } -} - -func TestProgress_NilCallbackDoesNotPanic(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - svc.progress(0.5, "test") -} - -// ============================================================================= -// insertChunks -// ============================================================================= - -func TestInsertChunks_EmptyChunks(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithInsertChunksFunc( - func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }, - ) - err := svc.insertChunks(context.Background(), nil) - if err != nil { - t.Errorf("expected no error for nil chunks, got %v", err) - } -} - -func TestInsertChunks_BaseNameAndDatasetID(t *testing.T) { - var capturedBaseName, capturedDatasetID string - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithInsertChunksFunc( - func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - capturedBaseName = baseName - capturedDatasetID = datasetID - return nil, nil - }, - ) - chunks := []map[string]any{{"text": "hello"}} - err := svc.insertChunks(context.Background(), chunks) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if capturedBaseName != "ragflow_tenant-1" { - t.Errorf("baseName = %q, want \"ragflow_tenant-1\"", capturedBaseName) - } - if capturedDatasetID != "kb-1" { - t.Errorf("datasetID = %q, want \"kb-1\"", capturedDatasetID) - } -} - -// ============================================================================= -// updateDocumentMetadata -// ============================================================================= - -func TestUpdateDocumentMetadata_EmptyMetadata(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(&stubDocService{}) - err := svc.updateDocumentMetadata("doc-1", nil) - if err != nil { - t.Errorf("empty metadata should not error, got %v", err) - } -} - -func TestUpdateDocumentMetadata_MergesNewKeys(t *testing.T) { - mds := &stubDocService{metaData: map[string]any{"existing": "old"}} - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(mds) - err := svc.updateDocumentMetadata("doc-1", map[string]any{"new_key": "value"}) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if mds.metaData["existing"] != "old" { - t.Errorf("existing key should be preserved: got %q", mds.metaData["existing"]) - } - if mds.metaData["new_key"] != "value" { - t.Errorf("new_key = %q, want \"value\"", mds.metaData["new_key"]) - } -} - -func TestUpdateDocumentMetadata_PreservesExistingKey(t *testing.T) { - mds := &stubDocService{metaData: map[string]any{"author": "Alice"}} - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(mds) - err := svc.updateDocumentMetadata("doc-1", map[string]any{"author": "Bob"}) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if mds.metaData["author"] != "Alice" { - t.Errorf("existing key should NOT be overwritten: got %q", mds.metaData["author"]) - } -} - -// ============================================================================= -// recordPipelineLog -// ============================================================================= - -func TestRecordPipelineLog(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( - func(log *entity.PipelineOperationLog) error { return nil }, - ) - svc.recordPipelineLog("doc-1", `{"components": {}}`, "done") -} - -// ============================================================================= -// updateDocumentMetadata -// ============================================================================= - -// incrementChunkNum -// ============================================================================= - -func TestIncrementChunkNum(t *testing.T) { - counter := &stubChunkCounter{} - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). - WithChunkCounter(counter) - err := svc.incrementChunkNum("doc-1", "kb-1", 10, 100, 0) - if err != nil { - t.Errorf("unexpected error: %v", err) - } - if counter.lastDocID != "doc-1" { - t.Errorf("docID = %q, want \"doc-1\"", counter.lastDocID) - } - if counter.lastKbID != "kb-1" { - t.Errorf("kbID = %q, want \"kb-1\"", counter.lastKbID) - } - if counter.lastChunkNum != 10 { - t.Errorf("chunkNum = %d, want 10", counter.lastChunkNum) - } - if counter.lastTokenNum != 100 { - t.Errorf("tokenNum = %d, want 100", counter.lastTokenNum) - } -} - -func TestIncrementChunkNum_ProcessDuration(t *testing.T) { - counter := &stubChunkCounter{} - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). - WithChunkCounter(counter) - err := svc.incrementChunkNum("doc-1", "kb-1", 5, 50, 12.5) - if err != nil { - t.Errorf("unexpected error: %v", err) - } - if counter.lastDuration != 12.5 { - t.Errorf("duration = %f, want 12.5", counter.lastDuration) - } - if counter.lastChunkNum != 5 { - t.Errorf("chunkNum = %d, want 5", counter.lastChunkNum) - } - if counter.lastTokenNum != 50 { - t.Errorf("tokenNum = %d, want 50", counter.lastTokenNum) - } -} - -// ============================================================================= -// RunDataflow — Python: run_dataflow (line 94) -// ============================================================================= - -func TestRunDataflow_NilOutput(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - err := svc.processOutput(context.Background(), nil) - if err != nil { - t.Errorf("expected nil error for nil output, got %v", err) - } -} - -func TestRunDataflow_EmptyOutput(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( - func(log *entity.PipelineOperationLog) error { return nil }, - ) - err := svc.processOutput(context.Background(), map[string]any{}) - if err != nil { - t.Errorf("expected nil error for empty output, got %v", err) - } -} - -func TestRunDataflow_NormalizedEmpty(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( - func(log *entity.PipelineOperationLog) error { return nil }, - ) - err := svc.processOutput(context.Background(), map[string]any{"markdown": ""}) - if err != nil { - t.Errorf("expected nil error for empty normalized output, got %v", err) - } -} - -func TestRunDataflow_FullFlow(t *testing.T) { - var progressCalls []float64 - var progressMsgs []string - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }). - WithDocService(&stubDocService{}). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). - WithChunkCounter(&stubChunkCounter{}). - WithProgressFunc(func(prog float64, msg string) { - progressCalls = append(progressCalls, prog) - progressMsgs = append(progressMsgs, msg) - }) - - output := map[string]any{ - "chunks": []map[string]any{ - {"text": "hello"}, - {"text": "world"}, - }, - } - err := svc.processOutput(context.Background(), output) - if err != nil { - t.Errorf("unexpected error: %v", err) - } - if len(progressCalls) < 2 { - t.Fatalf("expected multiple progress calls, got %v", progressCalls) - } - if progressCalls[len(progressCalls)-1] != 1.0 { - t.Fatalf("final progress = %v, want 1.0", progressCalls[len(progressCalls)-1]) - } - lastMsg := progressMsgs[len(progressMsgs)-1] - if !strings.Contains(lastMsg, "Indexing done (") || !strings.Contains(lastMsg, "Task done (") { - t.Fatalf("final progress msg = %q, want indexing/task timing message", lastMsg) - } -} - -func TestRunDataflow_AlreadyHasVectors(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }). - WithDocService(&stubDocService{}). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). - WithChunkCounter(&stubChunkCounter{}). - WithProgressFunc(func(prog float64, msg string) {}) - - output := map[string]any{ - "chunks": []map[string]any{ - {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, - }, - } - err := svc.processOutput(context.Background(), output) - if err != nil { - t.Errorf("unexpected error: %v", err) - } -} - -func TestRunDataflow_ContextCanceled(t *testing.T) { - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - ctx, cancel := context.WithCancel(context.Background()) - cancel() - - err := svc.processOutput(ctx, map[string]any{ - "chunks": []map[string]any{{"text": "hello"}}, - }) - if err == nil { - t.Error("expected context canceled error") - } -} - -func TestExtractDataflowPipelinePayload_UnwrapsSingleTerminalOutput(t *testing.T) { - dsl := `{"dsl":{"components":{"Parser:A":{"downstream":["Tokenizer:B"]},"Tokenizer:B":{"downstream":[]}}}}` - out := map[string]any{ - "Tokenizer:B": map[string]any{ - "output_format": "chunks", - "chunks": []map[string]any{{"text": "hello world"}}, - }, - "state": map[string]any{"Tokenizer:B": map[string]any{"output_format": "chunks"}}, - } - - payload, err := extractDataflowPipelinePayload(dsl, out) - if err != nil { - t.Fatalf("extractDataflowPipelinePayload: %v", err) - } - if got := payload["output_format"]; got != "chunks" { - t.Fatalf("output_format = %v, want chunks", got) - } - chunks, ok := payload["chunks"].([]map[string]any) - if !ok { - t.Fatalf("chunks = %T, want []map[string]any", payload["chunks"]) - } - if len(chunks) != 1 || chunks[0]["text"] != "hello world" { - t.Fatalf("chunks = %v, want single hello world chunk", chunks) - } -} - -func TestExtractDataflowPipelinePayload_ErrorsOnMultipleTerminals(t *testing.T) { - dsl := `{"dsl":{"components":{"Tokenizer:A":{"downstream":[]},"Tokenizer:B":{"downstream":[]}}}}` - _, err := extractDataflowPipelinePayload(dsl, map[string]any{ - "Tokenizer:A": map[string]any{"output_format": "chunks"}, - "Tokenizer:B": map[string]any{"output_format": "chunks"}, - }) - if err == nil { - t.Fatal("expected error for multiple terminals") - } - if !strings.Contains(err.Error(), "exactly 1 terminal") { - t.Fatalf("err = %v, want exactly 1 terminal", err) - } -} - -func TestDataflowService_Run_MainFlowWithStubs(t *testing.T) { - logged := false - inserted := false - var progressCalls []float64 - var progressMsgs []string - - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). - WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { - return `{"nodes":[{"id":"n1"}],"edges":[]}`, "flow-corrected", nil - }). - WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { - return map[string]any{ - "chunks": []map[string]any{ - {"text": "hello world"}, - }, - }, dsl, nil - }). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - inserted = true - return nil, nil - }). - WithDocService(&stubDocService{}). - WithChunkCounter(&stubChunkCounter{}). - WithProgressFunc(func(prog float64, msg string) { - progressCalls = append(progressCalls, prog) - progressMsgs = append(progressMsgs, msg) - }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { - logged = true - if log.PipelineID == nil || *log.PipelineID != "flow-corrected" { - t.Fatalf("PipelineID = %v, want flow-corrected", log.PipelineID) - } - return nil - }) - - err := svc.Execute(context.Background()) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if !inserted { - t.Fatal("expected insertChunks to be called") - } - if !logged { - t.Fatal("expected pipeline log to be created") - } - if len(progressCalls) == 0 || progressCalls[len(progressCalls)-1] != 1.0 { - t.Fatalf("expected final progress 1.0, got %v", progressCalls) - } - lastMsg := progressMsgs[len(progressMsgs)-1] - if !strings.Contains(lastMsg, "Indexing done (") || !strings.Contains(lastMsg, "Task done (") { - t.Fatalf("final progress msg = %q, want indexing/task timing message", lastMsg) - } -} - -func TestInsertChunks_ReportsBatchProgress(t *testing.T) { - var progressCalls []float64 - svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 1). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - time.Sleep(5 * time.Millisecond) - return nil, nil - }). - WithProgressFunc(func(prog float64, msg string) { - progressCalls = append(progressCalls, prog) - }) - - chunks := []map[string]any{ - {"text": "a"}, - {"text": "b"}, - } - if err := svc.insertChunks(context.Background(), chunks); err != nil { - t.Fatalf("unexpected error: %v", err) - } - if len(progressCalls) == 0 { - t.Fatal("expected indexing progress callbacks") - } -} - -// ============================================================================= -// Stub implementations for testing -// ============================================================================= - -type stubDocService struct { - err error - metaData map[string]any - lastReq *service.UpdateDocumentRequest - lastDocID string -} - -type stubChunkCounter struct { - err error - lastDocID string - lastKbID string - lastChunkNum int - lastTokenNum int - lastDuration float64 - callCount int -} - -func (s *stubDocService) UpdateDocument(id string, req *service.UpdateDocumentRequest) error { - s.lastDocID = id - s.lastReq = req - return s.err -} - -func (s *stubDocService) GetDocumentMetadataByID(docID string) (map[string]any, error) { - if s.err != nil { - return nil, s.err - } - if s.metaData == nil { - return make(map[string]any), nil - } - return s.metaData, nil -} - -func (s *stubDocService) SetDocumentMetadata(docID string, meta map[string]any) error { - if s.metaData == nil { - s.metaData = make(map[string]any) - } - for k, v := range meta { - s.metaData[k] = v - } - return s.err -} - -func (s *stubChunkCounter) IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { - s.lastDocID = docID - s.lastKbID = kbID - s.lastChunkNum = chunkNum - s.lastTokenNum = tokenConsumption - s.lastDuration = duration - s.callCount++ - return s.err -} - -// Compile-time checks -var ( - _ docService = (*stubDocService)(nil) - _ chunkCounter = (*stubChunkCounter)(nil) -) - -func makeEmbeddingModelForResolver() *models.EmbeddingModel { - return models.NewEmbeddingModel(&stubDriver{}, strPtr("embed"), &models.APIConfig{}, 128) -} - -func TestEmbedderResolver_ExplicitEmbeddingModelWins(t *testing.T) { - var gotTenantID, gotEmbdID string - resolver := newEmbedderResolver( - func(tenantID, embdID string) (*models.EmbeddingModel, error) { - gotTenantID, gotEmbdID = tenantID, embdID - return makeEmbeddingModelForResolver(), nil - }, - func(string) (*entity.Knowledgebase, error) { - t.Fatal("kb lookup should not run when embedding_model is set") - return nil, nil - }, - ) - emb, err := resolver("tenant-1", "kb-1", "explicit-embd") - if err != nil { - t.Fatalf("resolver: %v", err) - } - if emb == nil { - t.Fatal("expected embedder") - } - if gotTenantID != "tenant-1" || gotEmbdID != "explicit-embd" { - t.Fatalf("resolver args = (%q, %q), want (tenant-1, explicit-embd)", gotTenantID, gotEmbdID) - } -} - -func TestEmbedderResolver_FallsBackToDatasetEmbedding(t *testing.T) { - var gotEmbdID string - resolver := newEmbedderResolver( - func(_ string, embdID string) (*models.EmbeddingModel, error) { - gotEmbdID = embdID - return makeEmbeddingModelForResolver(), nil - }, - func(kbID string) (*entity.Knowledgebase, error) { - if kbID != "kb-1" { - t.Fatalf("kb lookup id = %q, want kb-1", kbID) - } - return &entity.Knowledgebase{ID: "kb-1", EmbdID: "lookup-embd"}, nil - }, - ) - if _, err := resolver("tenant-1", "kb-1", ""); err != nil { - t.Fatalf("resolver: %v", err) - } - if gotEmbdID != "lookup-embd" { - t.Fatalf("got embd id %q, want lookup-embd", gotEmbdID) - } -} - -func TestEmbedderResolver_MissingDatasetEmbeddingReturnsError(t *testing.T) { - resolver := newEmbedderResolver( - func(string, string) (*models.EmbeddingModel, error) { - t.Fatal("model resolver should not be called") - return nil, nil - }, - func(string) (*entity.Knowledgebase, error) { - return &entity.Knowledgebase{ID: "kb-1", EmbdID: ""}, nil - }, - ) - _, err := resolver("tenant-1", "kb-1", "") - if err == nil { - t.Fatal("expected error when dataset embd_id is missing, got nil") - } - if !strings.Contains(err.Error(), "dataset has no embd_id configured") { - t.Fatalf("err = %v, want dataset has no embd_id configured", err) - } -} - -func TestEmbedderResolver_MissingEmbeddingModelAndKBReturnsError(t *testing.T) { - resolver := newEmbedderResolver( - func(string, string) (*models.EmbeddingModel, error) { - t.Fatal("model resolver should not be called") - return nil, nil - }, - func(string) (*entity.Knowledgebase, error) { - t.Fatal("kb lookup should not be called without a kb_id") - return nil, nil - }, - ) - _, err := resolver("tenant-1", "", "") - if err == nil { - t.Fatal("expected error when neither embedding_model nor kb_id provided") - } - if !strings.Contains(err.Error(), "neither embedding_model nor kb_id") { - t.Fatalf("err = %v, want neither embedding_model nor kb_id", err) - } -} diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go new file mode 100644 index 0000000000..be8a552d4b --- /dev/null +++ b/internal/ingestion/task/embedder.go @@ -0,0 +1,95 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "fmt" + "strings" + + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/entity/models" + componentpkg "ragflow/internal/ingestion/component" + "ragflow/internal/service" +) + +type embedder struct { + model *models.EmbeddingModel +} + +func (e *embedder) MaxTokens() int { + if e == nil || e.model == nil { + return 0 + } + return e.model.MaxTokens +} + +func (e *embedder) Encode(texts []string) ([]componentpkg.EmbeddingResult, error) { + config := &models.EmbeddingConfig{Dimension: 0} + embeds, err := e.model.ModelDriver.Embed(e.model.ModelName, texts, e.model.APIConfig, config) + if err != nil { + return nil, err + } + vecs := make([]componentpkg.EmbeddingResult, len(embeds)) + for i, v := range embeds { + vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: v.TokenCount} + } + return vecs, nil +} + +// newEmbedderResolver builds the production embedder resolver used by the +// Tokenizer component. It honors an explicit embedding-model id (from the +// Tokenizer's setups) and falls back to the dataset's configured embd_id when +// none is given. Kept as a constructor over injectable deps so the resolution +// logic stays unit-testable without a live model provider / DB. +func newEmbedderResolver( + getEmbeddingModel func(tenantID, embdID string) (*models.EmbeddingModel, error), + getKnowledgebaseByID func(kbID string) (*entity.Knowledgebase, error), +) componentpkg.EmbedderResolver { + return func(tenantID, kbID, embeddingModel string) (componentpkg.Embedder, error) { + embdID := strings.TrimSpace(embeddingModel) + if embdID == "" { + if strings.TrimSpace(kbID) == "" { + return nil, fmt.Errorf("embedding requested but neither embedding_model nor kb_id provided") + } + kb, err := getKnowledgebaseByID(kbID) + if err != nil { + return nil, err + } + if kb == nil || strings.TrimSpace(kb.EmbdID) == "" { + return nil, fmt.Errorf("embedding requested but dataset has no embd_id configured") + } + embdID = kb.EmbdID + } + model, err := getEmbeddingModel(tenantID, embdID) + if err != nil { + return nil, err + } + return &embedder{model: model}, nil + } +} + +// init wires the production embedder resolver into the component package. The +// component package must not import internal/service (dependency direction), +// so the concrete resolver is injected here - the task package is the +// composition root for ingestion runs. +func init() { + componentpkg.DefaultEmbedderResolver = newEmbedderResolver( + service.NewModelProviderService().GetEmbeddingModel, + dao.NewKnowledgebaseDAO().GetByID, + ) +} diff --git a/internal/ingestion/task/embedder_test.go b/internal/ingestion/task/embedder_test.go new file mode 100644 index 0000000000..0a21c295bb --- /dev/null +++ b/internal/ingestion/task/embedder_test.go @@ -0,0 +1,177 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "strings" + "testing" + + "ragflow/internal/entity" + "ragflow/internal/entity/models" +) + +func makeEmbeddingModelForResolver() *models.EmbeddingModel { + return models.NewEmbeddingModel(&stubDriver{}, strPtr("embed"), &models.APIConfig{}, 128) +} + +func TestEmbedderResolver_ExplicitEmbeddingModelWins(t *testing.T) { + var gotTenantID, gotEmbdID string + resolver := newEmbedderResolver( + func(tenantID, embdID string) (*models.EmbeddingModel, error) { + gotTenantID, gotEmbdID = tenantID, embdID + return makeEmbeddingModelForResolver(), nil + }, + func(string) (*entity.Knowledgebase, error) { + t.Fatal("kb lookup should not run when embedding_model is set") + return nil, nil + }, + ) + emb, err := resolver("tenant-1", "kb-1", "explicit-embd") + if err != nil { + t.Fatalf("resolver: %v", err) + } + if emb == nil { + t.Fatal("expected embedder") + } + if gotTenantID != "tenant-1" || gotEmbdID != "explicit-embd" { + t.Fatalf("resolver args = (%q, %q), want (tenant-1, explicit-embd)", gotTenantID, gotEmbdID) + } +} + +func TestEmbedderResolver_FallsBackToDatasetEmbedding(t *testing.T) { + var gotEmbdID string + resolver := newEmbedderResolver( + func(_ string, embdID string) (*models.EmbeddingModel, error) { + gotEmbdID = embdID + return makeEmbeddingModelForResolver(), nil + }, + func(kbID string) (*entity.Knowledgebase, error) { + if kbID != "kb-1" { + t.Fatalf("kb lookup id = %q, want kb-1", kbID) + } + return &entity.Knowledgebase{ID: "kb-1", EmbdID: "lookup-embd"}, nil + }, + ) + if _, err := resolver("tenant-1", "kb-1", ""); err != nil { + t.Fatalf("resolver: %v", err) + } + if gotEmbdID != "lookup-embd" { + t.Fatalf("got embd id %q, want lookup-embd", gotEmbdID) + } +} + +func TestEmbedderResolver_MissingDatasetEmbeddingReturnsError(t *testing.T) { + resolver := newEmbedderResolver( + func(string, string) (*models.EmbeddingModel, error) { + t.Fatal("model resolver should not be called") + return nil, nil + }, + func(string) (*entity.Knowledgebase, error) { + return &entity.Knowledgebase{ID: "kb-1", EmbdID: ""}, nil + }, + ) + _, err := resolver("tenant-1", "kb-1", "") + if err == nil { + t.Fatal("expected error when dataset embd_id is missing, got nil") + } + if !strings.Contains(err.Error(), "dataset has no embd_id configured") { + t.Fatalf("err = %v, want dataset has no embd_id configured", err) + } +} + +func TestEmbedderResolver_MissingEmbeddingModelAndKBReturnsError(t *testing.T) { + resolver := newEmbedderResolver( + func(string, string) (*models.EmbeddingModel, error) { + t.Fatal("model resolver should not be called") + return nil, nil + }, + func(string) (*entity.Knowledgebase, error) { + t.Fatal("kb lookup should not be called without a kb_id") + return nil, nil + }, + ) + _, err := resolver("tenant-1", "", "") + if err == nil { + t.Fatal("expected error when neither embedding_model nor kb_id provided") + } + if !strings.Contains(err.Error(), "neither embedding_model nor kb_id") { + t.Fatalf("err = %v, want neither embedding_model nor kb_id", err) + } +} + +// stubDriver records the texts passed to Embed for verification. +type stubDriver struct { + capturedTexts []string +} + +func (d *stubDriver) Embed(modelName *string, texts []string, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig) ([]models.EmbeddingData, error) { + d.capturedTexts = texts + result := make([]models.EmbeddingData, len(texts)) + for i := range texts { + result[i] = models.EmbeddingData{ + Embedding: []float64{float64(i), 0.1}, + Index: i, + TokenCount: len(texts[i]), + } + } + return result, nil +} +func (d *stubDriver) NewInstance(baseURL map[string]string) models.ModelDriver { return d } +func (d *stubDriver) Name() string { return "stub" } +func (d *stubDriver) ChatWithMessages(modelName string, messages []models.Message, apiConfig *models.APIConfig, chatModelConfig *models.ChatConfig) (*models.ChatResponse, error) { + return nil, nil +} +func (d *stubDriver) ChatStreamlyWithSender(modelName string, messages []models.Message, apiConfig *models.APIConfig, modelConfig *models.ChatConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) Rerank(modelName *string, query string, documents []string, apiConfig *models.APIConfig, rerankConfig *models.RerankConfig) (*models.RerankResponse, error) { + return nil, nil +} +func (d *stubDriver) TranscribeAudio(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig) (*models.ASRResponse, error) { + return nil, nil +} +func (d *stubDriver) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) AudioSpeech(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig) (*models.TTSResponse, error) { + return nil, nil +} +func (d *stubDriver) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) OCRFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, ocrConfig *models.OCRConfig) (*models.OCRFileResponse, error) { + return nil, nil +} +func (d *stubDriver) ParseFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, parseFileConfig *models.ParseFileConfig) (*models.ParseFileResponse, error) { + return nil, nil +} +func (d *stubDriver) ListModels(apiConfig *models.APIConfig) ([]models.ListModelResponse, error) { + return nil, nil +} +func (d *stubDriver) Balance(apiConfig *models.APIConfig) (map[string]interface{}, error) { + return nil, nil +} +func (d *stubDriver) CheckConnection(apiConfig *models.APIConfig) error { return nil } +func (d *stubDriver) ListTasks(apiConfig *models.APIConfig) ([]models.ListTaskStatus, error) { + return nil, nil +} +func (d *stubDriver) ShowTask(taskID string, apiConfig *models.APIConfig) (*models.TaskResponse, error) { + return nil, nil +} +func (d *stubDriver) ToolCall(name string, arguments map[string]interface{}) (string, error) { + return "", nil +} diff --git a/internal/ingestion/task/embedding_test.go b/internal/ingestion/task/embedding_test.go deleted file mode 100644 index 89c362e9bf..0000000000 --- a/internal/ingestion/task/embedding_test.go +++ /dev/null @@ -1,77 +0,0 @@ -package task - -import ( - "ragflow/internal/entity/models" -) - -// stubDriver records the texts passed to Embed for verification. -type stubDriver struct { - capturedTexts []string -} - -func (d *stubDriver) Embed(modelName *string, texts []string, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig) ([]models.EmbeddingData, error) { - d.capturedTexts = texts - result := make([]models.EmbeddingData, len(texts)) - for i := range texts { - result[i] = models.EmbeddingData{ - Embedding: []float64{float64(i), 0.1}, - Index: i, - TokenCount: len(texts[i]), - } - } - return result, nil -} -func (d *stubDriver) NewInstance(baseURL map[string]string) models.ModelDriver { return d } -func (d *stubDriver) Name() string { return "stub" } -func (d *stubDriver) ChatWithMessages(modelName string, messages []models.Message, apiConfig *models.APIConfig, chatModelConfig *models.ChatConfig) (*models.ChatResponse, error) { - return nil, nil -} -func (d *stubDriver) ChatStreamlyWithSender(modelName string, messages []models.Message, apiConfig *models.APIConfig, modelConfig *models.ChatConfig, sender func(*string, *string) error) error { - return nil -} -func (d *stubDriver) Rerank(modelName *string, query string, documents []string, apiConfig *models.APIConfig, rerankConfig *models.RerankConfig) (*models.RerankResponse, error) { - return nil, nil -} -func (d *stubDriver) TranscribeAudio(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig) (*models.ASRResponse, error) { - return nil, nil -} -func (d *stubDriver) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig, sender func(*string, *string) error) error { - return nil -} -func (d *stubDriver) AudioSpeech(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig) (*models.TTSResponse, error) { - return nil, nil -} -func (d *stubDriver) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig, sender func(*string, *string) error) error { - return nil -} -func (d *stubDriver) OCRFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, ocrConfig *models.OCRConfig) (*models.OCRFileResponse, error) { - return nil, nil -} -func (d *stubDriver) ParseFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, parseFileConfig *models.ParseFileConfig) (*models.ParseFileResponse, error) { - return nil, nil -} -func (d *stubDriver) ListModels(apiConfig *models.APIConfig) ([]models.ListModelResponse, error) { - return nil, nil -} -func (d *stubDriver) Balance(apiConfig *models.APIConfig) (map[string]interface{}, error) { - return nil, nil -} -func (d *stubDriver) CheckConnection(apiConfig *models.APIConfig) error { return nil } -func (d *stubDriver) ListTasks(apiConfig *models.APIConfig) ([]models.ListTaskStatus, error) { - return nil, nil -} -func (d *stubDriver) ShowTask(taskID string, apiConfig *models.APIConfig) (*models.TaskResponse, error) { - return nil, nil -} -func (d *stubDriver) ToolCall(name string, arguments map[string]interface{}) (string, error) { - return "", nil -} - -func makeTestEmbeddingModel(stub *stubDriver, maxTokens int) *models.EmbeddingModel { - return &models.EmbeddingModel{ - ModelDriver: stub, - ModelName: strPtr("test-model"), - APIConfig: &models.APIConfig{}, - MaxTokens: maxTokens, - } -} diff --git a/internal/ingestion/task/embedding_token_test.go b/internal/ingestion/task/embedding_token_test.go deleted file mode 100644 index d7ea315628..0000000000 --- a/internal/ingestion/task/embedding_token_test.go +++ /dev/null @@ -1,55 +0,0 @@ -package task - -import ( - "testing" -) - -// ============================================================================= -// GetEmbeddingTokenConsumption -// ============================================================================= - -func TestGetEmbeddingTokenConsumption_Int(t *testing.T) { - input := map[string]any{EmbeddingTokenConsumptionKey: 42} - result := GetEmbeddingTokenConsumption(input) - if result != 42 { - t.Errorf("got %d, want 42", result) - } -} - -func TestGetEmbeddingTokenConsumption_Float64(t *testing.T) { - input := map[string]any{EmbeddingTokenConsumptionKey: float64(42)} - result := GetEmbeddingTokenConsumption(input) - if result != 42 { - t.Errorf("got %d, want 42", result) - } -} - -func TestGetEmbeddingTokenConsumption_MissingKey(t *testing.T) { - result := GetEmbeddingTokenConsumption(map[string]any{}) - if result != 0 { - t.Errorf("got %d, want 0", result) - } -} - -func TestGetEmbeddingTokenConsumption_NilMap(t *testing.T) { - result := GetEmbeddingTokenConsumption(nil) - if result != 0 { - t.Errorf("got %d, want 0", result) - } -} - -func TestGetEmbeddingTokenConsumption_WrongType(t *testing.T) { - input := map[string]any{EmbeddingTokenConsumptionKey: "not a number"} - result := GetEmbeddingTokenConsumption(input) - if result != 0 { - t.Errorf("got %d, want 0", result) - } -} - -func TestGetEmbeddingTokenConsumption_Zero(t *testing.T) { - input := map[string]any{EmbeddingTokenConsumptionKey: 0} - result := GetEmbeddingTokenConsumption(input) - if result != 0 { - t.Errorf("got %d, want 0", result) - } -} diff --git a/internal/ingestion/task/golden_compare.go b/internal/ingestion/task/golden_compare.go index ab61168e9c..106ade14a0 100644 --- a/internal/ingestion/task/golden_compare.go +++ b/internal/ingestion/task/golden_compare.go @@ -20,35 +20,35 @@ import ( "time" ) -// GoldenDataflowResult is the structured output used by the local golden tools. -type GoldenDataflowResult struct { +// GoldenCompareResult is the structured output used by the local golden tools. +type GoldenCompareResult struct { NormalizedChunks []map[string]any `json:"normalized_chunks"` ProcessedChunks []map[string]any `json:"processed_chunks"` MergedMetadata map[string]any `json:"merged_metadata"` } -// ProcessPipelineOutputForGolden replays the deterministic dataflow post-processing +// ProcessPipelineOutputForGolden replays the deterministic pipeline post-processing // steps from a pipeline.run()-style output without embedding or external writes. func ProcessPipelineOutputForGolden( pipelineOutput map[string]any, docID string, kbID string, docName string, -) GoldenDataflowResult { +) GoldenCompareResult { normalized := NormalizeChunks(pipelineOutput) if normalized == nil { normalized = []map[string]any{} } processed := deepCopyChunks(normalized) - metadata := ProcessChunksForDataflow(processed, docID, kbID, docName, time.Now()) + metadata := ProcessChunksForPipeline(processed, docID, kbID, docName, time.Now()) if metadata == nil { metadata = map[string]any{} } - return GoldenDataflowResult{ + return GoldenCompareResult{ NormalizedChunks: normalized, ProcessedChunks: processed, MergedMetadata: metadata, } -} \ No newline at end of file +} diff --git a/internal/ingestion/task/golden_compare_test.go b/internal/ingestion/task/golden_compare_test.go index 23005732ae..0764078aa2 100644 --- a/internal/ingestion/task/golden_compare_test.go +++ b/internal/ingestion/task/golden_compare_test.go @@ -29,7 +29,7 @@ func TestProcessPipelineOutputForGolden_Markdown(t *testing.T) { } } -func TestProcessChunksForDataflow_StableFields(t *testing.T) { +func TestProcessChunksForPipeline_StableFields(t *testing.T) { now := time.Date(2026, 7, 3, 20, 0, 0, 0, time.FixedZone("CST", 8*3600)) chunks := []map[string]any{ { @@ -41,7 +41,7 @@ func TestProcessChunksForDataflow_StableFields(t *testing.T) { }, } - meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "sample.md", now) + meta := ProcessChunksForPipeline(chunks, "doc-1", "kb-1", "sample.md", now) chunk := chunks[0] if chunk["doc_id"] != "doc-1" { diff --git a/internal/ingestion/task/dataflow_e2e_test.go b/internal/ingestion/task/pipeline_e2e_test.go similarity index 88% rename from internal/ingestion/task/dataflow_e2e_test.go rename to internal/ingestion/task/pipeline_e2e_test.go index a980f0f2de..3b20fe00d4 100644 --- a/internal/ingestion/task/dataflow_e2e_test.go +++ b/internal/ingestion/task/pipeline_e2e_test.go @@ -161,7 +161,7 @@ func isPortOpen(host string, port int) bool { // Main E2E Test with Subtests for Elasticsearch and Infinity // ============================================================================= -func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { +func TestPipelineE2E_TaskHandlerToPipelineExecutor(t *testing.T) { testCases := []struct { name string engineType engine.EngineType @@ -178,7 +178,7 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - t.Logf("Running Dataflow E2E test with engine: %s", tc.name) + t.Logf("Running Pipeline E2E test with engine: %s", tc.name) // Setup test database db := testutil.SetupTestDB(t) @@ -209,6 +209,7 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { if err != nil { t.Fatalf("LoadFromIngestionTask failed: %v", err) } + taskCtx.Ctx = context.Background() // Track what was called var ( @@ -220,13 +221,13 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { // Create TaskHandler with mocked DataflowService factory handler := NewTaskHandler(taskCtx) - handler.WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0) + handler.WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + svc := mustNewPipelineExecutor(t, ctx, canvasID, 0) // Mock loadDSLFunc - svc.WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { + svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { loadDSLCalled = true - return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, dataflowID, nil + return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil }) // Mock runPipelineFunc - returns test chunks (with vectors to skip embedding) @@ -250,7 +251,7 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { }) // Use the injected DocEngine for insertChunks! - svc.WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { insertChunksCalled = true t.Logf("DocEngine InsertChunks called! baseName=%s datasetID=%s len(chunks)=%d", baseName, datasetID, len(chunks)) ids, err := docEngine.InsertChunks(ctx, chunks, baseName, datasetID) @@ -261,26 +262,12 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { return ids, err }) - svc.WithChunkCounter(&stubChunkCounter{}) return svc, nil }) - // Also set progress callback - var progressEvents []struct { - prog float64 - msg string - } - taskCtx.ProgressFunc = func(prog float64, msg string) { - t.Logf("PROGRESS: %.2f %s", prog, msg) - progressEvents = append(progressEvents, struct { - prog float64 - msg string - }{prog, msg}) - } - // Execute the task handler! t.Logf("Calling TaskHandler.Handle()...") - err = handler.Handle() + _, err = handler.Handle() if err != nil { t.Fatalf("TaskHandler.Handle failed: %v", err) } @@ -330,20 +317,6 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { } } - // Verify progress reported - if len(progressEvents) == 0 { - t.Fatal("Expected at least one progress event") - } - foundDone := false - for _, ev := range progressEvents { - if ev.prog == 1.0 { - foundDone = true - } - } - if !foundDone { - t.Fatal("Expected progress to reach 1.0") - } - // Verify final task status can be marked completed ingestSvc := service.NewIngestionTaskService() if err := ingestSvc.MarkCompleted(taskID); err != nil { @@ -358,7 +331,7 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { t.Errorf("Final task status = %q, want %q", finalTask.Status, common.COMPLETED) } - t.Logf("SUCCESS: Dataflow E2E test passed with %s engine!", tc.name) + t.Logf("SUCCESS: Pipeline E2E test passed with %s engine!", tc.name) }) } } diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go new file mode 100644 index 0000000000..9878982a9d --- /dev/null +++ b/internal/ingestion/task/pipeline_executor.go @@ -0,0 +1,237 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "context" + "encoding/json" + "fmt" + "ragflow/internal/utility" + "strings" + "time" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine" + "ragflow/internal/entity" + pipelinepkg "ragflow/internal/ingestion/pipeline" +) + +// PipelineResult is the outcome of a pipeline run: chunks have been +// indexed, and these bookkeeping inputs remain for the caller to apply to +// document state (metadata merge + chunk/token counter bumps). +type PipelineResult struct { + DocID string + KbID string + Metadata map[string]any + ChunkCount int + TokenConsumption int +} + +type PipelineExecutor struct { + taskCtx *TaskContext + canvasID string + docBulkSize int + + indexWriter *chunkIndexWriter + logCreateFunc func(log *entity.PipelineOperationLog) error + loadDSLFunc func(ctx context.Context, canvasID string) (string, string, error) + runPipelineFunc func(ctx context.Context, dsl string) (map[string]any, string, error) + progressSink pipelinepkg.ProgressSink +} + +func validateTaskContext(taskCtx *TaskContext) error { + if taskCtx == nil { + return fmt.Errorf("pipeline executor: nil task context") + } + if taskCtx.Doc.ID == "" { + return fmt.Errorf("pipeline executor: empty document id") + } + if taskCtx.Doc.KbID == "" { + return fmt.Errorf("pipeline executor: empty document knowledgebase id") + } + if taskCtx.Doc.Name == nil || *taskCtx.Doc.Name == "" { + return fmt.Errorf("pipeline executor: empty document name") + } + if taskCtx.KB.ID == "" { + return fmt.Errorf("pipeline executor: empty knowledgebase id") + } + if taskCtx.Tenant.ID == "" { + return fmt.Errorf("pipeline executor: empty tenant id") + } + return nil +} + +func NewPipelineExecutor( + taskCtx *TaskContext, + canvasID string, + docBulkSize int, +) (*PipelineExecutor, error) { + if err := validateTaskContext(taskCtx); err != nil { + return nil, err + } + if strings.TrimSpace(canvasID) == "" { + return nil, fmt.Errorf("pipeline executor: empty canvas id") + } + svc := &PipelineExecutor{ + taskCtx: taskCtx, + canvasID: canvasID, + docBulkSize: docBulkSize, + indexWriter: newChunkIndexWriter( + func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + return engine.Get().InsertChunks(ctx, chunks, baseName, datasetID) + }, + fmt.Sprintf("ragflow_%s", taskCtx.Tenant.ID), + taskCtx.Doc.KbID, + docBulkSize, + ), + logCreateFunc: dao.NewPipelineOperationLogDAO().Create, + } + svc.loadDSLFunc = svc.defaultLoadDSL + svc.runPipelineFunc = svc.defaultRunPipeline + return svc, nil +} + +func (s *PipelineExecutor) WithInsertFunc(f InsertFunc) *PipelineExecutor { + s.indexWriter.insertFunc = f + return s +} + +func (s *PipelineExecutor) WithLogCreateFunc(f func(log *entity.PipelineOperationLog) error) *PipelineExecutor { + s.logCreateFunc = f + return s +} + +func (s *PipelineExecutor) WithLoadDSLFunc(f func(ctx context.Context, canvasID string) (string, string, error)) *PipelineExecutor { + s.loadDSLFunc = f + return s +} + +func (s *PipelineExecutor) WithRunPipelineFunc(f func(ctx context.Context, dsl string) (map[string]any, string, error)) *PipelineExecutor { + s.runPipelineFunc = f + return s +} + +// WithProgressSink injects a sink that receives pipeline component progress +// events. The sink owns all document/ingestion_task_log persistence; when +// unset, the pipeline runs DB-independent (progress events are dropped). +func (s *PipelineExecutor) WithProgressSink(sink pipelinepkg.ProgressSink) *PipelineExecutor { + s.progressSink = sink + return s +} + +func (s *PipelineExecutor) KB() *entity.Knowledgebase { return &s.taskCtx.KB } +func (s *PipelineExecutor) Doc() *entity.Document { return &s.taskCtx.Doc } +func (s *PipelineExecutor) Tenant() *entity.Tenant { return &s.taskCtx.Tenant } + +func (s *PipelineExecutor) Execute(ctx context.Context) (*PipelineResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + dsl, correctedID, err := s.loadDSLFunc(ctx, s.canvasID) + if err != nil { + return nil, err + } + if correctedID != "" { + s.canvasID = correctedID + } + + pipelineOutput, pipelineDSL, err := s.runPipelineFunc(ctx, dsl) + if err != nil { + return nil, err + } + + if s.taskCtx.Doc.ID == CANVAS_DEBUG_DOC_ID { + s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") + return nil, nil + } + + result, err := s.processOutput(ctx, pipelineOutput) + if err != nil { + return nil, err + } + + if pipelineDSL != "" { + s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") + } + return result, nil +} + +func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map[string]any) (*PipelineResult, error) { + taskStart := time.Now() + if pipelineOutput == nil { + return nil, nil + } + if err := ctx.Err(); err != nil { + return nil, err + } + + chunks := NormalizeChunks(pipelineOutput) + if chunks == nil { + return nil, nil + } + + embeddingTokenConsumption := GetEmbeddingTokenConsumption(pipelineOutput) + metadata := ProcessChunksForPipeline( + chunks, + s.taskCtx.Doc.ID, + s.taskCtx.Doc.KbID, + *s.taskCtx.Doc.Name, + time.Now(), + ) + + indexStart := time.Now() + if err := s.indexWriter.Write(ctx, chunks); err != nil { + return nil, err + } + _ = time.Since(indexStart) + _ = time.Since(taskStart) + + return &PipelineResult{ + DocID: s.taskCtx.Doc.ID, + KbID: s.taskCtx.Doc.KbID, + Metadata: metadata, + ChunkCount: len(chunks), + TokenConsumption: embeddingTokenConsumption, + }, nil +} + +func (s *PipelineExecutor) recordPipelineLog(docID, dsl, status string) { + var dslMap entity.JSONMap + if err := json.Unmarshal([]byte(dsl), &dslMap); err != nil { + dslMap = entity.JSONMap{"raw": dsl} + } + log := &entity.PipelineOperationLog{ + ID: utility.GenerateUUID(), + TenantID: s.Tenant().ID, + KbID: s.KB().ID, + DocumentID: docID, + PipelineID: &s.canvasID, + TaskType: string(entity.PipelineTaskTypeParse), + DSL: dslMap, + ParserID: s.taskCtx.Doc.ParserID, + DocumentName: *s.Doc().Name, + DocumentSuffix: s.taskCtx.Doc.Suffix, + DocumentType: s.taskCtx.Doc.Type, + SourceFrom: s.taskCtx.Doc.SourceType, + OperationStatus: status, + } + if err := s.logCreateFunc(log); err != nil { + common.Warn(fmt.Sprintf("failed to record pipeline log: %v", err)) + } +} diff --git a/internal/ingestion/task/pipeline_executor_defaults.go b/internal/ingestion/task/pipeline_executor_defaults.go new file mode 100644 index 0000000000..0b41062991 --- /dev/null +++ b/internal/ingestion/task/pipeline_executor_defaults.go @@ -0,0 +1,89 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "context" + "encoding/json" + "fmt" + + "ragflow/internal/common" + "ragflow/internal/dao" + pipelinepkg "ragflow/internal/ingestion/pipeline" +) + +func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, canvasID string) (string, string, error) { + if s == nil || s.taskCtx == nil { + return "", "", fmt.Errorf("pipeline executor: nil task context") + } + if canvasID == "" { + return "", "", fmt.Errorf("pipeline executor: empty canvas id") + } + canvas, err := dao.NewUserCanvasDAO().GetByID(canvasID) + if err != nil { + return "", "", fmt.Errorf("load canvas %s: %w", canvasID, err) + } + + canvasTitle := "" + if canvas.Title != nil { + canvasTitle = *canvas.Title + } + common.Info(fmt.Sprintf("load canvas %s, name %s", canvasID, canvasTitle)) + + raw, err := json.Marshal(canvas.DSL) + if err != nil { + return "", "", fmt.Errorf("marshal canvas dsl %s: %w", canvasID, err) + } + return string(raw), canvasID, nil +} + +func (s *PipelineExecutor) defaultRunPipeline(ctx context.Context, dsl string) (map[string]any, string, error) { + if s == nil || s.taskCtx == nil { + return nil, dsl, fmt.Errorf("pipeline executor: nil task context") + } + + // Use doc ID as pipeline ID if available, otherwise a placeholder + pipelineID := "pipeline_" + s.taskCtx.Doc.ID + if s.taskCtx.IngestionTask != nil && s.taskCtx.IngestionTask.ID != "" { + pipelineID = s.taskCtx.IngestionTask.ID + } + pipe, err := pipelinepkg.NewPipelineFromDSL([]byte(dsl), pipelineID, + pipelinepkg.WithProgressSink(s.progressSink), + pipelinepkg.WithDocumentID(s.taskCtx.Doc.ID)) + if err != nil { + return nil, dsl, fmt.Errorf("compile pipeline dsl: %w", err) + } + inputs := map[string]any{} + if s.taskCtx.Doc.ID != "" { + inputs["doc_id"] = s.taskCtx.Doc.ID + } + if s.taskCtx.File != nil { + inputs["file"] = s.taskCtx.File + } + inputs["tenant_id"] = s.taskCtx.Tenant.ID + inputs["kb_id"] = s.taskCtx.KB.ID + + output, err := pipe.Run(ctx, inputs) + if err != nil { + return nil, dsl, err + } + payload, err := pipelinepkg.ExtractPayload(dsl, output) + if err != nil { + return nil, dsl, err + } + return payload, dsl, nil +} diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go new file mode 100644 index 0000000000..d430f80510 --- /dev/null +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -0,0 +1,442 @@ +package task + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "ragflow/internal/agent/runtime" + "ragflow/internal/dao" + "ragflow/internal/entity" + pipelinepkg "ragflow/internal/ingestion/pipeline" +) + +// ============================================================================= +// Test helpers +// ============================================================================= + +func strPtr(s string) *string { return &s } + +func makeTaskCtx() *TaskContext { + return &TaskContext{ + IngestionTask: &entity.IngestionTask{ + ID: "task-1", + DocumentID: "doc-1", + }, + Doc: entity.Document{ + ID: "doc-1", + KbID: "kb-1", + Name: strPtr("test-doc.pdf"), + Suffix: ".pdf", + Type: "pdf", + }, + KB: entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ + ID: "tenant-1", + }, + } +} + +func setupPipelineExecutorTestDB(t *testing.T) func() { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.UserCanvas{}, &entity.PipelineOperationLog{}); err != nil { + t.Fatalf("auto-migrate sqlite: %v", err) + } + origDB := dao.DB + dao.DB = db + return func() { dao.DB = origDB } +} + +func mustNewPipelineExecutor(t *testing.T, taskCtx *TaskContext, canvasID string, docBulkSize int) *PipelineExecutor { + t.Helper() + svc, err := NewPipelineExecutor(taskCtx, canvasID, docBulkSize) + if err != nil { + t.Fatalf("NewPipelineExecutor: %v", err) + } + return svc +} + +// ============================================================================= +// NewPipelineExecutor — constructor +// ============================================================================= + +func TestNewPipelineExecutor_Basic(t *testing.T) { + svc, err := NewPipelineExecutor(makeTaskCtx(), "flow-1", 0) + if err != nil { + t.Fatalf("NewPipelineExecutor: %v", err) + } + if svc == nil { + t.Fatal("NewPipelineExecutor returned nil") + } + if svc.taskCtx == nil { + t.Error("taskCtx should not be nil") + } + if svc.indexWriter == nil || svc.logCreateFunc == nil || svc.loadDSLFunc == nil || svc.runPipelineFunc == nil { + t.Fatal("expected production dependencies to be fully initialized") + } +} + +func TestNewPipelineExecutor_RejectsNilTaskContext(t *testing.T) { + _, err := NewPipelineExecutor(nil, "flow-1", 0) + if err == nil { + t.Fatal("expected error for nil task context") + } +} + +func TestNewPipelineExecutor_RejectsEmptyCanvasID(t *testing.T) { + _, err := NewPipelineExecutor(makeTaskCtx(), "", 0) + if err == nil { + t.Fatal("expected error for empty canvas id") + } +} + +func TestNewPipelineExecutor_RejectsIncompleteTaskContext(t *testing.T) { + tests := []struct { + name string + mutate func(*TaskContext) + }{ + {name: "missing doc id", mutate: func(ctx *TaskContext) { ctx.Doc.ID = "" }}, + {name: "missing kb id", mutate: func(ctx *TaskContext) { ctx.Doc.KbID = "" }}, + {name: "missing doc name", mutate: func(ctx *TaskContext) { ctx.Doc.Name = nil }}, + {name: "missing knowledgebase id", mutate: func(ctx *TaskContext) { ctx.KB.ID = "" }}, + {name: "missing tenant id", mutate: func(ctx *TaskContext) { ctx.Tenant.ID = "" }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := makeTaskCtx() + tt.mutate(ctx) + _, err := NewPipelineExecutor(ctx, "flow-1", 0) + if err == nil { + t.Fatal("expected validation error") + } + }) + } +} + +func TestNewPipelineExecutor_DocBulkSize(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 128) + if svc.docBulkSize != 128 { + t.Errorf("docBulkSize = %d, want 128", svc.docBulkSize) + } +} + +func TestNewPipelineExecutor_CanvasID(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "my-flow-id", 0) + if svc.canvasID != "my-flow-id" { + t.Errorf("canvasID = %q, want %q", svc.canvasID, "my-flow-id") + } +} + +func TestKB_Doc_Tenant_Accessors(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0) + if svc.KB().ID != "kb-1" { + t.Errorf("KB().ID = %q, want \"kb-1\"", svc.KB().ID) + } + if svc.Doc().ID != "doc-1" { + t.Errorf("Doc().ID = %q, want \"doc-1\"", svc.Doc().ID) + } + if svc.Tenant().ID != "tenant-1" { + t.Errorf("Tenant().ID = %q, want \"tenant-1\"", svc.Tenant().ID) + } +} + +// ============================================================================= +// processChunks +// ============================================================================= + +func TestPipelineExecutor_ProcessChunks_WrapsProcessChunksForPipeline(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0) + chunks := []map[string]any{{"text": "hello world"}} + meta := ProcessChunksForPipeline(chunks, svc.taskCtx.Doc.ID, svc.taskCtx.Doc.KbID, *svc.taskCtx.Doc.Name, time.Now()) + + // Verify the wrapper method works correctly and chunks are processed + if chunks[0]["doc_id"] != "doc-1" { + t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"]) + } + if meta != nil { + // No need to verify the detailed content of meta as ProcessChunksForPipeline already has comprehensive tests + } +} + +// ============================================================================= +// insertChunks +// ============================================================================= + +func TestInsertChunks_EmptyChunks(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithInsertFunc( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }, + ) + err := svc.indexWriter.Write(context.Background(), nil) + if err != nil { + t.Errorf("expected no error for nil chunks, got %v", err) + } +} + +func TestInsertChunks_BaseNameAndDatasetID(t *testing.T) { + var capturedBaseName, capturedDatasetID string + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithInsertFunc( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + capturedBaseName = baseName + capturedDatasetID = datasetID + return nil, nil + }, + ) + chunks := []map[string]any{{"text": "hello"}} + err := svc.indexWriter.Write(context.Background(), chunks) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if capturedBaseName != "ragflow_tenant-1" { + t.Errorf("baseName = %q, want \"ragflow_tenant-1\"", capturedBaseName) + } + if capturedDatasetID != "kb-1" { + t.Errorf("datasetID = %q, want \"kb-1\"", capturedDatasetID) + } +} + +func TestRecordPipelineLog(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + svc.recordPipelineLog("doc-1", `{"components": {}}`, "done") +} + +func TestRecordPipelineLog_InvalidJSONFallback(t *testing.T) { + var captured *entity.PipelineOperationLog + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { + captured = log + return nil + }, + ) + svc.recordPipelineLog("doc-1", "not-valid-json", "done") + if captured == nil { + t.Fatal("logCreateFunc was not called") + } + raw, ok := captured.DSL["raw"].(string) + if !ok || raw != "not-valid-json" { + t.Fatalf("DSL = %v, want {\"raw\": \"not-valid-json\"}", captured.DSL) + } +} + +func TestRecordPipelineLog_ValidJSONParsed(t *testing.T) { + var captured *entity.PipelineOperationLog + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { + captured = log + return nil + }, + ) + svc.recordPipelineLog("doc-1", `{"components": {"a": {"obj": {"component_name": "Parser", "params": {}}}}}`, "done") + if captured == nil { + t.Fatal("logCreateFunc was not called") + } + if captured.DSL["raw"] != nil { + t.Fatalf("DSL should be parsed JSON, not fallback raw; got %v", captured.DSL) + } +} + +// ============================================================================= +// updateDocumentMetadata +// ============================================================================= + +func TestRunPipeline_NilOutput(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0) + _, err := svc.processOutput(context.Background(), nil) + if err != nil { + t.Errorf("expected nil error for nil output, got %v", err) + } +} + +func TestRunPipeline_EmptyOutput(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + _, err := svc.processOutput(context.Background(), map[string]any{}) + if err != nil { + t.Errorf("expected nil error for empty output, got %v", err) + } +} + +func TestRunPipeline_NormalizedEmpty(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + _, err := svc.processOutput(context.Background(), map[string]any{"markdown": ""}) + if err != nil { + t.Errorf("expected nil error for empty normalized output, got %v", err) + } +} + +func TestRunPipeline_FullFlow(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }) + output := map[string]any{ + "chunks": []map[string]any{ + {"text": "hello"}, + {"text": "world"}, + }, + } + _, err := svc.processOutput(context.Background(), output) + if err != nil { + t.Errorf("unexpected error: %v", err) + } +} + +func TestRunPipeline_AlreadyHasVectors(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }) + + output := map[string]any{ + "chunks": []map[string]any{ + {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, + }, + } + _, err := svc.processOutput(context.Background(), output) + if err != nil { + t.Errorf("unexpected error: %v", err) + } +} + +func TestRunPipeline_ContextCanceled(t *testing.T) { + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := svc.processOutput(ctx, map[string]any{ + "chunks": []map[string]any{{"text": "hello"}}, + }) + if err == nil { + t.Error("expected context canceled error") + } +} + +func TestPipelineExecutor_Run_MainFlowWithStubs(t *testing.T) { + logged := false + inserted := false + + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0). + WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"n1"}],"edges":[]}`, "flow-corrected", nil + }). + WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + return map[string]any{ + "chunks": []map[string]any{ + {"text": "hello world"}, + }, + }, dsl, nil + }). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + inserted = true + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { + logged = true + if log.PipelineID == nil || *log.PipelineID != "flow-corrected" { + t.Fatalf("PipelineID = %v, want flow-corrected", log.PipelineID) + } + return nil + }) + + _, err := svc.Execute(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !inserted { + t.Fatal("expected insertChunks to be called") + } + if !logged { + t.Fatal("expected pipeline log to be created") + } +} + +// ============================================================================= +// Stub implementations for testing +// ============================================================================= + +// recordingProgressSink captures progress events for asserting the executor +// forwards its sink through defaultRunPipeline into the pipeline. +type recordingProgressSink struct { + mu sync.Mutex + total int + totalSet bool + events []pipelinepkg.ProgressEvent +} + +func (r *recordingProgressSink) OnComponentTotal(taskID string, total int) { + r.mu.Lock() + defer r.mu.Unlock() + r.total = total + r.totalSet = true +} + +func (r *recordingProgressSink) OnComponentProgress(ev pipelinepkg.ProgressEvent) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, ev) +} + +type sinkPassthroughStage struct{} + +func (sinkPassthroughStage) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) { + return inputs, nil +} + +// TestPipelineExecutorDefaultRunPipelineForwardsSink verifies the sink set via +// WithProgressSink is threaded through defaultRunPipeline into the pipeline, +// which reports the component total and lifecycle events back to the sink. +func TestPipelineExecutorDefaultRunPipelineForwardsSink(t *testing.T) { + const nameA = "task.SinkPassthroughA" + runtime.MustRegister(nameA, runtime.CategoryIngestion, + func(_ string, _ map[string]any) (runtime.Component, error) { return sinkPassthroughStage{}, nil }, + runtime.Metadata{Version: "1.0.0"}) + + sink := &recordingProgressSink{} + svc := mustNewPipelineExecutor(t, makeTaskCtx(), "flow-1", 0) + svc.WithProgressSink(sink) + + dsl := `{"dsl":{"components":{"begin":{"obj":{"component_name":"Begin","params":{}},"downstream":["a"]},"a":{"obj":{"component_name":"` + nameA + `","params":{}},"upstream":["begin"]}},"path":["begin","a"],"graph":{"nodes":[]}}}` + + if _, _, err := svc.defaultRunPipeline(context.Background(), dsl); err != nil { + t.Fatalf("defaultRunPipeline: %v", err) + } + + sink.mu.Lock() + defer sink.mu.Unlock() + if !sink.totalSet || sink.total != 2 { + t.Fatalf("OnComponentTotal = (%d, set=%v), want 2", sink.total, sink.totalSet) + } + if len(sink.events) == 0 { + t.Fatal("expected progress events forwarded to sink, got none") + } + for _, ev := range sink.events { + if ev.TaskID != "task-1" { + t.Fatalf("event TaskID = %q, want task-1", ev.TaskID) + } + if ev.DocumentID != "doc-1" { + t.Fatalf("event DocumentID = %q, want doc-1", ev.DocumentID) + } + } +} diff --git a/internal/ingestion/task/dataflow_real_pipeline_integration_test.go b/internal/ingestion/task/pipeline_real_integration_test.go similarity index 89% rename from internal/ingestion/task/dataflow_real_pipeline_integration_test.go rename to internal/ingestion/task/pipeline_real_integration_test.go index 3aedaa55ad..a31082e42b 100644 --- a/internal/ingestion/task/dataflow_real_pipeline_integration_test.go +++ b/internal/ingestion/task/pipeline_real_integration_test.go @@ -35,8 +35,7 @@ import ( gormlogger "gorm.io/gorm/logger" ) -func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { - prepareTokenizerResourceForTaskIntegration(t) +func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { requireTokenizerPool(t) cfg := mustLoadTaskRealIntegrationConfig(t) @@ -123,19 +122,17 @@ func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { TenantID: tenantID, EmbdID: "embd-1", }, - Tenant: entity.Tenant{ID: tenantID}, - ProgressFunc: func(prog float64, msg string) {}, + Tenant: entity.Tenant{ID: tenantID}, } var inserted [][]map[string]any - svc := mustNewDataflowService(t, taskCtx, canvasID, 0, 0). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { inserted = append(inserted, deepCopyTaskChunks(chunks)) return nil, nil - }). - WithChunkCounter(&stubChunkCounter{}) + }) - if err := svc.Execute(context.Background()); err != nil { + if _, err := svc.Execute(context.Background()); err != nil { t.Fatalf("Run: %v", err) } if len(inserted) != 1 { @@ -154,8 +151,7 @@ func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { } } -func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *testing.T) { - prepareTokenizerResourceForTaskIntegration(t) +func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *testing.T) { requireTokenizerPool(t) cfg := mustLoadTaskRealIntegrationConfig(t) @@ -261,14 +257,12 @@ func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *test TenantID: tenantID, EmbdID: "embd-1", }, - Tenant: entity.Tenant{ID: tenantID}, - ProgressFunc: func(prog float64, msg string) {}, + Tenant: entity.Tenant{ID: tenantID}, } - svc := mustNewDataflowService(t, taskCtx, canvasID, 0, 0). - WithChunkCounter(&stubChunkCounter{}) + svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0) - if err := svc.Execute(context.Background()); err != nil { + if _, err := svc.Execute(context.Background()); err != nil { t.Fatalf("Run: %v", err) } @@ -302,8 +296,7 @@ func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *test } } -func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { - prepareTokenizerResourceForTaskIntegration(t) +func TestRunPipeline_RealPipelineOutput_ProducesIndexFields(t *testing.T) { requireTokenizerPool(t) cfg := mustLoadTaskRealIntegrationConfig(t) @@ -339,7 +332,7 @@ func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { } templateBytes = disableTokenizerEmbeddingForTaskTemplate(t, templateBytes) - pipe, err := pipelinepkg.NewPipelineFromDSL(templateBytes, "task-dataflow-real-pipeline") + pipe, err := pipelinepkg.NewPipelineFromDSL(templateBytes, "task-real-pipeline") if err != nil { t.Fatalf("NewPipelineFromDSL: %v", err) } @@ -383,20 +376,18 @@ func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { TenantID: tenantID, EmbdID: "embd-1", }, - Tenant: entity.Tenant{ID: tenantID}, - ProgressFunc: func(prog float64, msg string) {}, + Tenant: entity.Tenant{ID: tenantID}, } var inserted [][]map[string]any - svc := mustNewDataflowService(t, taskCtx, "flow-real-1", 0, 0). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + svc := mustNewPipelineExecutor(t, taskCtx, "flow-real-1", 0). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { inserted = append(inserted, deepCopyTaskChunks(chunks)) return nil, nil - }). - WithChunkCounter(&stubChunkCounter{}) + }) - if err := svc.processOutput(context.Background(), pipelineOut); err != nil { - t.Fatalf("RunDataflow: %v", err) + if _, err := svc.processOutput(context.Background(), pipelineOut); err != nil { + t.Fatalf("RunPipeline: %v", err) } if len(inserted) != 1 { @@ -423,7 +414,7 @@ func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { t.Fatalf("chunks[%d].content_sm_ltks = %T/%v, want non-empty string", i, ck["content_sm_ltks"], ck["content_sm_ltks"]) } if _, hasText := ck["text"]; hasText { - t.Fatalf("chunks[%d] should not keep raw text field after RunDataflow: %v", i, ck["text"]) + t.Fatalf("chunks[%d] should not keep raw text field after RunPipeline: %v", i, ck["text"]) } } } @@ -514,23 +505,6 @@ func disableTokenizerEmbeddingForTaskTemplate(t *testing.T, raw []byte) []byte { return out } -func prepareTokenizerResourceForTaskIntegration(t *testing.T) { - t.Helper() - if common.GetEnv(common.EnvRAGFlowDictPath) != "" { - return - } - const systemDictPath = "/usr/share/infinity/resource" - if _, err := os.Stat(filepath.Join(systemDictPath, "rag", "huqie.txt")); err != nil { - t.Skipf("system tokenizer resource not found at %s: %v", systemDictPath, err) - } - if err := os.Setenv(common.EnvRAGFlowDictPath, systemDictPath); err != nil { - t.Fatalf("set RAGFLOW_DICT_PATH=%s: %v", systemDictPath, err) - } - t.Cleanup(func() { - _ = os.Unsetenv(common.EnvRAGFlowDictPath) - }) -} - func taskMustSymlink(t *testing.T, src, dst string) { t.Helper() if err := os.Symlink(src, dst); err != nil { @@ -619,20 +593,13 @@ func taskMustPrepareTokenizerOpenCC(t *testing.T, root string) { func requireTokenizerPool(t *testing.T) { t.Helper() - if tokenizer.IsInitialized() { - return - } - cfg := &tokenizer.PoolConfig{ - DictPath: common.GetEnv(common.EnvRAGFlowDictPath), + if err := tokenizer.Init(&tokenizer.PoolConfig{ + DictPath: "/usr/share/infinity/resource", MinSize: 1, MaxSize: 2, IdleTimeout: 30 * time.Second, AcquireTimeout: 5 * time.Second, - } - if cfg.DictPath == "" { - cfg.DictPath = "/usr/share/infinity/resource" - } - if err := tokenizer.Init(cfg); err != nil { + }); err != nil { t.Skipf("tokenizer pool init failed: %v", err) } } diff --git a/internal/ingestion/task/pipeline_run_test.go b/internal/ingestion/task/pipeline_run_test.go new file mode 100644 index 0000000000..8ae1ae202a --- /dev/null +++ b/internal/ingestion/task/pipeline_run_test.go @@ -0,0 +1,96 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "ragflow/internal/dao" + "ragflow/internal/entity" + pipelinepkg "ragflow/internal/ingestion/pipeline" +) + +func TestPipelineExecutor_DefaultLoadDSL_UsesUserCanvas(t *testing.T) { + cleanup := setupPipelineExecutorTestDB(t) + defer cleanup() + + dslMap := entity.JSONMap{"dsl": map[string]any{"graph": map[string]any{"nodes": []any{}, "edges": []any{}}}} + title := "title 1" + if err := dao.NewUserCanvasDAO().Create(&entity.UserCanvas{Title: &title, ID: "canvas-1", UserID: "u1", Permission: "me", CanvasCategory: "agent_canvas", DSL: dslMap}); err != nil { + t.Fatalf("create user canvas: %v", err) + } + + ctx := makeTaskCtx() + svc := mustNewPipelineExecutor(t, ctx, "canvas-1", 0) + gotDSL, correctedID, err := svc.loadDSLFunc(context.Background(), "canvas-1") + if err != nil { + t.Fatalf("loadDSLFunc: %v", err) + } + if correctedID != "canvas-1" { + t.Fatalf("correctedID = %q, want canvas-1", correctedID) + } + var decoded map[string]any + if err := json.Unmarshal([]byte(gotDSL), &decoded); err != nil { + t.Fatalf("unmarshal dsl: %v", err) + } + if _, ok := decoded["dsl"].(map[string]any); !ok { + t.Fatalf("decoded dsl = %v, want top-level dsl map", decoded) + } +} + +func TestExtractPipelinePayload_UnwrapsSingleTerminalOutput(t *testing.T) { + dsl := `{"dsl":{"components":{"Parser:A":{"downstream":["Tokenizer:B"]},"Tokenizer:B":{"downstream":[]}}}}` + out := map[string]any{ + "Tokenizer:B": map[string]any{ + "output_format": "chunks", + "chunks": []map[string]any{{"text": "hello world"}}, + }, + "state": map[string]any{"Tokenizer:B": map[string]any{"output_format": "chunks"}}, + } + + payload, err := pipelinepkg.ExtractPayload(dsl, out) + if err != nil { + t.Fatalf("ExtractPayload: %v", err) + } + if got := payload["output_format"]; got != "chunks" { + t.Fatalf("output_format = %v, want chunks", got) + } + chunks, ok := payload["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks = %T, want []map[string]any", payload["chunks"]) + } + if len(chunks) != 1 || chunks[0]["text"] != "hello world" { + t.Fatalf("chunks = %v, want single hello world chunk", chunks) + } +} + +func TestExtractPipelinePayload_ErrorsOnMultipleTerminals(t *testing.T) { + dsl := `{"dsl":{"components":{"Tokenizer:A":{"downstream":[]},"Tokenizer:B":{"downstream":[]}}}}` + _, err := pipelinepkg.ExtractPayload(dsl, map[string]any{ + "Tokenizer:A": map[string]any{"output_format": "chunks"}, + "Tokenizer:B": map[string]any{"output_format": "chunks"}, + }) + if err == nil { + t.Fatal("expected error for multiple terminals") + } + if !strings.Contains(err.Error(), "exactly 1 terminal") { + t.Fatalf("err = %v, want exactly 1 terminal", err) + } +} diff --git a/internal/ingestion/task/real_consumer_test.go b/internal/ingestion/task/real_consumer_test.go index f556880a1f..0a9003a79e 100644 --- a/internal/ingestion/task/real_consumer_test.go +++ b/internal/ingestion/task/real_consumer_test.go @@ -1,5 +1,4 @@ -//go:build manual -// +build manual +//go:build integration // // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. @@ -20,19 +19,14 @@ package task import ( + "context" "encoding/json" - "errors" - "fmt" - "strings" "testing" "ragflow/internal/common" "ragflow/internal/dao" - "ragflow/internal/engine/nats" "ragflow/internal/entity" - - "github.com/glebarez/sqlite" - "gorm.io/gorm" + "ragflow/internal/ingestion/testutil" ) // TestRealProducerConsumer exercises the project's real producer and consumer code paths: @@ -40,11 +34,8 @@ import ( // Producer: document.go pattern — Create(IngestionTask) → PublishTask(NATS) // Consumer: Ingestor.Start() core logic — calls each actual function in sequence func TestRealProducerConsumer(t *testing.T) { - // ── 1. NATS ── - natsEngine := nats.NewNatsEngine("localhost", 4222) - if err := natsEngine.Init(); err != nil { - t.Fatalf("NATS Init: %v", err) - } + // ── 1. NATS (embedded in-process server) ── + natsEngine := testutil.SetupNatsEngine(t) if err := natsEngine.InitConsumer("tasks.>"); err != nil { t.Fatalf("InitConsumer: %v", err) } @@ -59,27 +50,14 @@ func TestRealProducerConsumer(t *testing.T) { } // ── 2. SQLite DB ── - dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())) - db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{TranslateError: true}) - if err != nil { - t.Fatalf("failed to open sqlite: %v", err) - } - sqlDB, err := db.DB() - if err != nil { - t.Fatalf("failed to get sql DB: %v", err) - } - sqlDB.SetMaxOpenConns(1) - db.AutoMigrate( - &entity.IngestionTask{}, &entity.Task{}, - &entity.Document{}, &entity.Knowledgebase{}, &entity.Tenant{}, - ) - origDB := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = origDB }) + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() - db.Create(&entity.Tenant{ID: "t1", LLMID: "gpt-4", Status: ptr("1")}) - db.Create(&entity.Knowledgebase{ID: "kb1", TenantID: "t1", EmbdID: "e1", Status: ptr("1"), ParserConfig: entity.JSONMap{}}) - db.Create(&entity.Document{ID: "doc-real", KbID: "kb1", ParserID: "naive", ParserConfig: entity.JSONMap{}}) + db.Create(&entity.Tenant{ID: "t1", LLMID: "gpt-4", Status: testutil.StrPtr("1")}) + db.Create(&entity.Knowledgebase{ID: "kb1", TenantID: "t1", EmbdID: "e1", Status: testutil.StrPtr("1"), ParserConfig: entity.JSONMap{}}) + docName := "doc-real.pdf" + db.Create(&entity.Document{ID: "doc-real", KbID: "kb1", ParserID: "naive", ParserConfig: entity.JSONMap{}, Name: &docName}) // ── 3. Producer: Mirrors document.go:1062-1085 exactly ── ingestionTask := &entity.IngestionTask{ @@ -125,15 +103,19 @@ func TestRealProducerConsumer(t *testing.T) { // Mirrors Start():142-143 — SetRunningByIngestor ingestionTaskDAO := dao.NewIngestionTaskDAO() - task, err := ingestionTaskDAO.SetRunningByIngestor(taskMsg.TaskID) + _, err = ingestionTaskDAO.UpdateStatusIfCurrent(taskMsg.TaskID, common.CREATED, common.RUNNING) if err != nil { - if errors.Is(err, common.ErrTaskNotFound) { - t.Logf("Consumer: task %s not found in ingestion_task table — skipped", taskMsg.TaskID) - taskHandle.Ack() - return - } t.Fatalf("SetRunningByIngestor: %v", err) } + task, err := ingestionTaskDAO.GetByID(taskMsg.TaskID) + if err != nil { + t.Fatalf("GetByID: %v", err) + } + if task == nil { + t.Logf("Consumer: task %s not found in ingestion_task table — skipped", taskMsg.TaskID) + taskHandle.Ack() + return + } t.Logf("Consumer: SetRunningByIngestor status=%s", task.Status) // Mirrors Start():167-180 — status check @@ -148,21 +130,43 @@ func TestRealProducerConsumer(t *testing.T) { } // ── 5. executeTask (our modified version) ── + // Set a pipeline ID so the handler can resolve the canvas. + if err := db.Model(&entity.Document{}).Where("id = ?", "doc-real").Update("pipeline_id", "pipeline-real").Error; err != nil { + t.Fatalf("set pipeline_id: %v", err) + } + tc, err := LoadFromIngestionTask(task) if err != nil { t.Fatalf("LoadFromIngestionTask: %v", err) } + tc.Ctx = context.Background() t.Logf("Consumer: Loaded Doc=%s Parser=%s KB=%s Tenant=%s", tc.Doc.ID, tc.Doc.ParserID, tc.KB.ID, tc.Tenant.ID) handler := NewTaskHandler(tc) - if err := handler.Handle(); err != nil { + handler.WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + svc, err := NewPipelineExecutor(ctx, canvasID, 0) + if err != nil { + return nil, err + } + svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil + }) + svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + return nil, "", nil + }) + svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }) + return svc, nil + }) + if _, err := handler.Handle(); err != nil { t.Fatalf("Handle: %v", err) } t.Log("Consumer: TaskHandler.Handle() — OK") // Mirrors executeTask — mark as completed - if err := ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED); err != nil { + if _, err := ingestionTaskDAO.UpdateStatusIfCurrent(task.ID, common.RUNNING, common.COMPLETED); err != nil { t.Fatalf("UpdateStatus: %v", err) } diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go index 228ea6d950..e51b98ded7 100644 --- a/internal/ingestion/task/task_context.go +++ b/internal/ingestion/task/task_context.go @@ -21,6 +21,7 @@ import ( "fmt" "strings" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" ) @@ -38,7 +39,11 @@ type TaskContext struct { PipelineID string File any - ProgressFunc ProgressFunc + // Handle is the message-queue ack handle for the task message that scheduled + // this context. The scheduler sets it before queueing; the worker acks on a + // durably-persisted terminal status and nacks otherwise (e.g. shutdown + // mid-task) so the message is redelivered and resumed after restart. + Handle common.TaskHandle } // NewTaskContextForScheduling creates a lightweight TaskContext for queue scheduling. diff --git a/internal/ingestion/task/task_context_test.go b/internal/ingestion/task/task_context_test.go index 23cd7bbdb3..dc2226dc36 100644 --- a/internal/ingestion/task/task_context_test.go +++ b/internal/ingestion/task/task_context_test.go @@ -17,60 +17,23 @@ package task import ( - "fmt" - "strings" "testing" - "ragflow/internal/dao" "ragflow/internal/entity" - - "github.com/glebarez/sqlite" - "gorm.io/gorm" + "ragflow/internal/ingestion/testutil" ) -func setupTestDB(t *testing.T) *gorm.DB { - t.Helper() - dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())) - db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{TranslateError: true}) - if err != nil { - t.Fatalf("open sqlite: %v", err) - } - sqlDB, err := db.DB() - if err != nil { - t.Fatalf("failed to get sql DB: %v", err) - } - sqlDB.SetMaxOpenConns(1) - if err := db.AutoMigrate( - &entity.Task{}, - &entity.Document{}, - &entity.Knowledgebase{}, - &entity.Tenant{}, - ); err != nil { - t.Fatalf("migrate: %v", err) - } - return db -} - -func pushDB(t *testing.T, db *gorm.DB) { - t.Helper() - orig := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = orig }) -} - -func ptr(s string) *string { return &s } - func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) { - db := setupTestDB(t) - pushDB(t, db) + db := testutil.SetupTestDB(t) + testutil.ReplaceDBForTest(t, db) if err := db.Create(&entity.Document{ ID: "doc-1", KbID: "kb-1", ParserID: "naive", ParserConfig: entity.JSONMap{}, - Name: ptr("doc.pdf"), - Status: ptr("1"), + Name: testutil.StrPtr("doc.pdf"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create document: %v", err) } @@ -81,13 +44,13 @@ func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) EmbdID: "embd-1", PipelineID: &kbPipelineID, ParserConfig: entity.JSONMap{}, - Status: ptr("1"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create knowledgebase: %v", err) } if err := db.Create(&entity.Tenant{ ID: "tenant-1", - Status: ptr("1"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create tenant: %v", err) } @@ -106,8 +69,8 @@ func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) } func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { - db := setupTestDB(t) - pushDB(t, db) + db := testutil.SetupTestDB(t) + testutil.ReplaceDBForTest(t, db) docPipelineID := "doc-flow-1" if err := db.Create(&entity.Document{ @@ -116,8 +79,8 @@ func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { ParserID: "naive", ParserConfig: entity.JSONMap{}, PipelineID: &docPipelineID, - Name: ptr("doc.pdf"), - Status: ptr("1"), + Name: testutil.StrPtr("doc.pdf"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create document: %v", err) } @@ -128,13 +91,13 @@ func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { EmbdID: "embd-1", PipelineID: &kbPipelineID, ParserConfig: entity.JSONMap{}, - Status: ptr("1"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create knowledgebase: %v", err) } if err := db.Create(&entity.Tenant{ ID: "tenant-1", - Status: ptr("1"), + Status: testutil.StrPtr("1"), }).Error; err != nil { t.Fatalf("create tenant: %v", err) } diff --git a/internal/ingestion/task/task_handler.go b/internal/ingestion/task/task_handler.go index 6f85bedae9..5cd7833445 100644 --- a/internal/ingestion/task/task_handler.go +++ b/internal/ingestion/task/task_handler.go @@ -17,7 +17,6 @@ package task import ( - "context" "fmt" "strings" ) @@ -25,66 +24,38 @@ import ( // TaskHandler dispatches document processing tasks by task_type. // Mirrors Python task_handler.py:handle(). type TaskHandler struct { - ctx *TaskContext - newDataflowService func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) + ctx *TaskContext + newPipelineExecutor func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) } // NewTaskHandler creates a TaskHandler for the given task context. func NewTaskHandler(ctx *TaskContext) *TaskHandler { return &TaskHandler{ ctx: ctx, - newDataflowService: func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - return NewDataflowService(ctx, dataflowID, 0, 0) + newPipelineExecutor: func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + return NewPipelineExecutor(ctx, canvasID, 0) }, } } -func (h *TaskHandler) WithDataflowServiceFactory(factory func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error)) *TaskHandler { - h.newDataflowService = factory +func (h *TaskHandler) WithPipelineExecutorFactory(factory func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error)) *TaskHandler { + h.newPipelineExecutor = factory return h } // Handle routes the task by type and executes the appropriate handler. -func (h *TaskHandler) Handle() error { +func (h *TaskHandler) Handle() (*PipelineResult, error) { if h.ctx == nil { - return fmt.Errorf("task handler: nil context") + return nil, fmt.Errorf("task handler: nil context") } - return h.handleDataflow() + return h.handlePipeline() } -func (h *TaskHandler) handleMemory() error { - return nil // stub -} - -func (h *TaskHandler) handleDataflow() error { - dataflowID := "" - if strings.Trim(h.ctx.PipelineID, " ") != "" { - dataflowID = h.ctx.PipelineID - } - svc, err := h.newDataflowService(h.ctx, dataflowID) +func (h *TaskHandler) handlePipeline() (*PipelineResult, error) { + svc, err := h.newPipelineExecutor(h.ctx, strings.TrimSpace(h.ctx.PipelineID)) if err != nil { - return err + return nil, err } - runCtx := h.ctx.Ctx - if runCtx == nil { - runCtx = context.Background() - } - return svc.Execute(runCtx) -} - -func (h *TaskHandler) handleRaptor() error { - return nil // stub -} - -func (h *TaskHandler) handleGraphRAG() error { - return nil // stub -} - -func (h *TaskHandler) handleStub(name string) error { - return nil -} - -func (h *TaskHandler) handleStandard() error { - return nil // stub + return svc.Execute(h.ctx.Ctx) } diff --git a/internal/ingestion/task/task_handler_test.go b/internal/ingestion/task/task_handler_test.go index 7f2138ee9d..ed687ffbcd 100644 --- a/internal/ingestion/task/task_handler_test.go +++ b/internal/ingestion/task/task_handler_test.go @@ -33,47 +33,44 @@ func makeTaskHandlerTestContext(pipelineID string) *TaskContext { ID: "tenant-1", LLMID: "gpt-4", }, - ProgressFunc: func(prog float64, msg string) {}, } } -func newNoopDataflowService(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - if strings.TrimSpace(dataflowID) == "" { - dataflowID = "flow-1" +func newNoopPipelineExecutor(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + if strings.TrimSpace(canvasID) == "" { + canvasID = "flow-1" } - svc, err := NewDataflowService(ctx, dataflowID, 0, 0) + svc, err := NewPipelineExecutor(ctx, canvasID, 0) if err != nil { return nil, err } svc = svc. - WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil + WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil }). WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { return map[string]any{ "chunks": []map[string]any{{ - "text": "stub dataflow chunk", + "text": "stub pipeline chunk", "q_2_vec": []float64{0.1, 0.2}, }}, }, dsl, nil }). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { return nil, nil }). WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil - }). - WithDocService(&stubDocService{}). - WithChunkCounter(&stubChunkCounter{}) + }) return svc, nil } func newNoopTaskHandler(ctx *TaskContext) *TaskHandler { - return NewTaskHandler(ctx).WithDataflowServiceFactory(newNoopDataflowService) + return NewTaskHandler(ctx).WithPipelineExecutorFactory(newNoopPipelineExecutor) } func TestTaskHandler_HandleRejectsNilContext(t *testing.T) { - if err := NewTaskHandler(nil).Handle(); err == nil { + if _, err := NewTaskHandler(nil).Handle(); err == nil { t.Fatal("expected error for nil context") } } @@ -81,38 +78,32 @@ func TestTaskHandler_HandleRejectsNilContext(t *testing.T) { func TestTaskHandler_HandleRequiresPipelineID(t *testing.T) { ctx := makeTaskHandlerTestContext("") handler := NewTaskHandler(ctx) - if err := handler.Handle(); err == nil { + if _, err := handler.Handle(); err == nil { t.Fatal("expected error for empty pipeline id") } } -func TestTaskHandler_DefaultDataflowServiceInjectsProgress(t *testing.T) { +func TestTaskHandler_HandleRunWithFactory(t *testing.T) { ctx := makeTaskHandlerTestContext("flow-1") - handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - svc, err := NewDataflowService(ctx, dataflowID, 0, 0) - if err != nil { - t.Fatalf("NewDataflowService: %v", err) - } - if svc.progressFunc == nil { - t.Fatal("expected default progress func to be injected") - } - return newNoopDataflowService(ctx, dataflowID) + ctx.Ctx = context.Background() + handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + return newNoopPipelineExecutor(ctx, canvasID) }) - if err := handler.Handle(); err != nil { + if _, err := handler.Handle(); err != nil { t.Fatalf("unexpected error: %v", err) } } -func TestTaskHandler_Dataflow_UsesTaskContext(t *testing.T) { +func TestTaskHandler_Pipeline_UsesTaskContext(t *testing.T) { ctx := makeTaskHandlerTestContext("flow-1") type ctxKey string const key ctxKey = "trace" ctx.Ctx = context.WithValue(context.Background(), key, "task-ctx") - handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - return mustNewDataflowService(t, ctx, dataflowID, 0, 0). - WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil + handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + return mustNewPipelineExecutor(t, ctx, canvasID, 0). + WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil }). WithRunPipelineFunc(func(runCtx context.Context, dsl string) (map[string]any, string, error) { if got := runCtx.Value(key); got != "task-ctx" { @@ -120,76 +111,42 @@ func TestTaskHandler_Dataflow_UsesTaskContext(t *testing.T) { } return map[string]any{"chunks": []map[string]any{{"text": "stub", "q_2_vec": []float64{0.1, 0.2}}}}, dsl, nil }). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { return nil, nil }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). - WithDocService(&stubDocService{}). - WithChunkCounter(&stubChunkCounter{}), nil + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }), nil }) - if err := handler.Handle(); err != nil { + if _, err := handler.Handle(); err != nil { t.Fatalf("unexpected error: %v", err) } } -func TestTaskHandler_Dataflow_UsesBackgroundContextWhenMissing(t *testing.T) { - ctx := makeTaskHandlerTestContext("flow-1") - - handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - return mustNewDataflowService(t, ctx, dataflowID, 0, 0). - WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil - }). - WithRunPipelineFunc(func(runCtx context.Context, dsl string) (map[string]any, string, error) { - if runCtx == nil { - t.Fatal("runCtx is nil") - } - return map[string]any{"chunks": []map[string]any{{"text": "stub", "q_2_vec": []float64{0.1, 0.2}}}}, dsl, nil - }). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). - WithDocService(&stubDocService{}). - WithChunkCounter(&stubChunkCounter{}), nil - }) - if err := handler.Handle(); err != nil { - t.Fatalf("unexpected error: %v", err) - } -} - -func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { +func TestTaskHandler_Pipeline_ShowsProgressAndPipelineLog(t *testing.T) { ctx := makeTaskHandlerTestContext("flow-1") + ctx.Ctx = context.Background() ctx.Doc.PipelineID = testStrPtr("flow-1") - ctx.Doc.Name = testStrPtr("verify-dataflow.pdf") + ctx.Doc.Name = testStrPtr("verify-pipeline.pdf") var pipelineCalled bool var insertCalled bool var logCreateCalls int var insertedChunkCount int - var progressProgs []float64 - var progressMsgs []string - ctx.ProgressFunc = func(prog float64, msg string) { - progressProgs = append(progressProgs, prog) - progressMsgs = append(progressMsgs, msg) - } - - handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { - svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0). - WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil + handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { + svc := mustNewPipelineExecutor(t, ctx, canvasID, 0). + WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil }). WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { pipelineCalled = true return map[string]any{ "chunks": []map[string]any{{ - "text": "stub dataflow chunk", + "text": "stub pipeline chunk", "q_2_vec": []float64{0.1, 0.2}, }}, }, dsl, nil }). - WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { insertCalled = true insertedChunkCount = len(chunks) return nil, nil @@ -197,13 +154,11 @@ func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { logCreateCalls++ return nil - }). - WithDocService(&stubDocService{}). - WithChunkCounter(&stubChunkCounter{}) + }) return svc, nil }) - if err := handler.Handle(); err != nil { + if _, err := handler.Handle(); err != nil { t.Fatalf("handler.Handle() error: %v", err) } @@ -216,29 +171,6 @@ func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { if insertedChunkCount != 1 { t.Fatalf("insertedChunkCount = %d, want 1", insertedChunkCount) } - if len(progressProgs) == 0 { - t.Fatal("expected progress callbacks, got none") - } - - foundStartIndex := false - for _, msg := range progressMsgs { - if strings.Contains(msg, "Start to index") { - foundStartIndex = true - break - } - } - if !foundStartIndex { - t.Fatalf("expected progress message containing %q, got %v", "Start to index", progressMsgs) - } - - if got := progressProgs[len(progressProgs)-1]; got != 1.0 { - t.Fatalf("final progress = %v, want 1.0", got) - } - - lastMsg := progressMsgs[len(progressMsgs)-1] - if !strings.Contains(lastMsg, "Indexing done") { - t.Fatalf("final progress msg = %q, want substring %q", lastMsg, "Indexing done") - } if logCreateCalls != 1 { t.Fatalf("logCreateCalls = %d, want 1", logCreateCalls) diff --git a/internal/ingestion/task/tool/compare_dataflow_golden.go b/internal/ingestion/task/tool/compare_pipeline_golden.go similarity index 99% rename from internal/ingestion/task/tool/compare_dataflow_golden.go rename to internal/ingestion/task/tool/compare_pipeline_golden.go index a65124f8c4..9dbc6fd487 100644 --- a/internal/ingestion/task/tool/compare_dataflow_golden.go +++ b/internal/ingestion/task/tool/compare_pipeline_golden.go @@ -58,7 +58,7 @@ func main() { os.Exit(1) } -func runActual(input map[string]any, docID string, kbID string, docName string) (result task.GoldenDataflowResult, err error) { +func runActual(input map[string]any, docID string, kbID string, docName string) (result task.GoldenCompareResult, err error) { defer func() { if r := recover(); r != nil { err = fmt.Errorf("panic: %v", r) diff --git a/internal/ingestion/testutil/nats_helper.go b/internal/ingestion/testutil/nats_helper.go new file mode 100644 index 0000000000..59701dddd7 --- /dev/null +++ b/internal/ingestion/testutil/nats_helper.go @@ -0,0 +1,70 @@ +// +// 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 testutil + +import ( + "net" + "testing" + "time" + + natsengine "ragflow/internal/engine/nats" + + "github.com/nats-io/nats-server/v2/server" +) + +// SetupNatsEngine starts an in-process NATS server with JetStream enabled +// on a random port, creates a NatsEngine connected to it, and calls Init() +// to set up the RAGFLOW_TASKS stream. The server is automatically shut down +// when the test completes. +// +// The caller is responsible for calling InitConsumer if a consumer is needed. +func SetupNatsEngine(t *testing.T) *natsengine.NatsEngine { + t.Helper() + + opts := &server.Options{ + Port: -1, + JetStream: true, + StoreDir: t.TempDir(), + NoLog: true, + NoSigs: true, + } + + ns, err := server.NewServer(opts) + if err != nil { + t.Fatalf("create embedded NATS server: %v", err) + } + + ns.Start() + + if !ns.ReadyForConnections(10 * time.Second) { + ns.Shutdown() + t.Fatal("embedded NATS server did not become ready within 10s") + } + + t.Cleanup(func() { + ns.Shutdown() + ns.WaitForShutdown() + }) + + addr := ns.Addr().(*net.TCPAddr) + engine := natsengine.NewNatsEngine("127.0.0.1", addr.Port) + if err := engine.Init(); err != nil { + t.Fatalf("NatsEngine.Init against embedded server: %v", err) + } + + return engine +} diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index a13dea8ad6..833b692494 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -37,7 +37,6 @@ import ( "ragflow/internal/engine/redis" "ragflow/internal/entity" "ragflow/internal/entity/models" - "ragflow/internal/server" "regexp" "sort" "strconv" @@ -91,7 +90,6 @@ func searchConfigMap(value interface{}) (map[string]interface{}, bool) { // ChunkService chunk service type ChunkService struct { docEngine engine.DocEngine - engineType server.EngineType embeddingCache *utility.EmbeddingLRU kbDAO *dao.KnowledgebaseDAO userTenantDAO *dao.UserTenantDAO @@ -117,10 +115,8 @@ type ChunkService struct { // NewChunkService creates chunk service func NewChunkService() *ChunkService { - cfg := server.GetConfig() return &ChunkService{ docEngine: engine.Get(), - engineType: cfg.DocEngine.Type, embeddingCache: utility.NewEmbeddingLRU(1000), // default capacity kbDAO: dao.NewKnowledgebaseDAO(), userTenantDAO: dao.NewUserTenantDAO(), diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index 49a6f21148..a74987ae82 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -1331,6 +1331,7 @@ func (e *parseTestDocEngine) Close() error { func (e *parseTestDocEngine) GetType() string { return "test" } +func (e *parseTestDocEngine) SupportsPageRank() bool { return false } func (e *parseTestDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { return nil @@ -1504,6 +1505,7 @@ func (m *switchChunksEngineMock) GetScores(map[string]interface{}) map[string]fl func (m *switchChunksEngineMock) Ping(context.Context) error { return nil } func (m *switchChunksEngineMock) Close() error { return nil } func (m *switchChunksEngineMock) GetType() string { return "elasticsearch" } +func (m *switchChunksEngineMock) SupportsPageRank() bool { return true } func (m *switchChunksEngineMock) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { return nil } diff --git a/internal/service/chunk_types.go b/internal/service/chunk_types.go index 21fdf9592e..606265de93 100644 --- a/internal/service/chunk_types.go +++ b/internal/service/chunk_types.go @@ -24,7 +24,6 @@ import ( "ragflow/internal/common" "ragflow/internal/engine/redis" "ragflow/internal/entity" - "ragflow/internal/server" "strings" "ragflow/internal/dao" @@ -48,7 +47,6 @@ var ( // ChunkService chunk service type ChunkService struct { docEngine engine.DocEngine - engineType server.EngineType embeddingCache *utility.EmbeddingLRU kbDAO *dao.KnowledgebaseDAO taskDAO *dao.TaskDAO diff --git a/internal/service/dataset.go b/internal/service/dataset.go index 2f08e51861..967cc43b34 100644 --- a/internal/service/dataset.go +++ b/internal/service/dataset.go @@ -37,7 +37,6 @@ import ( enginetypes "ragflow/internal/engine/types" "ragflow/internal/entity" "ragflow/internal/entity/models" - "ragflow/internal/server" "ragflow/internal/service/nlp" "ragflow/internal/storage" "ragflow/internal/utility" @@ -156,16 +155,10 @@ type DatasetService struct { searchService *SearchService docEngine engine.DocEngine embeddingCache *utility.EmbeddingLRU - engineType server.EngineType } // NewDatasetService creates a new datasets service. func NewDatasetService() *DatasetService { - cfg := server.GetConfig() - engineType := server.EngineType("") - if cfg != nil { - engineType = cfg.DocEngine.Type - } return &DatasetService{ kbDAO: dao.NewKnowledgebaseDAO(), documentDAO: dao.NewDocumentDAO(), @@ -178,7 +171,6 @@ func NewDatasetService() *DatasetService { searchService: NewSearchService(), docEngine: engine.Get(), embeddingCache: utility.NewEmbeddingLRU(1000), - engineType: engineType, } } @@ -2833,7 +2825,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req UpdateDat if *req.Pagerank < 0 || *req.Pagerank > 100 { return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") } - if d.engineType == server.EngineInfinity { + if !d.docEngine.SupportsPageRank() { return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch") } indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) diff --git a/internal/service/document.go b/internal/service/document.go index 69a0908d22..27c3f31a0a 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -45,7 +45,6 @@ import ( "ragflow/internal/engine/redis" enginetypes "ragflow/internal/engine/types" "ragflow/internal/entity" - "ragflow/internal/server" "ragflow/internal/storage" "ragflow/internal/tokenizer" "ragflow/internal/utility" @@ -64,7 +63,6 @@ type DocumentService struct { ingestionTaskLogDAO *dao.IngestionTaskLogDAO ingestionTaskSvc *IngestionTaskService docEngine engine.DocEngine - engineType server.EngineType metadataSvc *MetadataService taskDAO *dao.TaskDAO file2DocumentDAO *dao.File2DocumentDAO @@ -75,7 +73,6 @@ type DocumentService struct { // NewDocumentService create document service func NewDocumentService() *DocumentService { - cfg := server.GetConfig() publisher := NewMessageQueueTaskPublisher() ingestionTaskSvc := NewIngestionTaskService() ingestionTaskSvc.SetTaskPublisher(publisher) @@ -86,7 +83,6 @@ func NewDocumentService() *DocumentService { ingestionTaskSvc: ingestionTaskSvc, kbDAO: dao.NewKnowledgebaseDAO(), docEngine: engine.Get(), - engineType: cfg.DocEngine.Type, metadataSvc: NewMetadataService(), taskDAO: dao.NewTaskDAO(), file2DocumentDAO: dao.NewFile2DocumentDAO(), @@ -651,6 +647,18 @@ func (s *DocumentService) IncrementChunkNum(docID, kbID string, chunkNum, tokenN }) } +// UpdateRunProgress mirrors a pipeline run's live progress into the document +// row so the document-list endpoint (which reads document.progress/run/ +// progress_msg) reflects in-flight Go pipeline progress. Best-effort by +// design; callers log and continue on error. +func (s *DocumentService) UpdateRunProgress(docID string, progress float64, run, progressMsg string) error { + return s.documentDAO.UpdateByID(docID, map[string]interface{}{ + "progress": progress, + "run": run, + "progress_msg": progressMsg, + }) +} + // DeleteDocument delete document — delegates to full cleanup logic. func (s *DocumentService) DeleteDocument(id string) error { return s.deleteDocumentFull(id) @@ -1945,43 +1953,20 @@ func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID strin return docs, nil } -// cancelDocParse sets Redis cancel signals for the document's active tasks and -// marks the document run status as CANCEL. Returns an error if the document is -// not in a cancellable state or the status update fails. +// cancelDocParse stops the ingestion task for the document by calling +// RequestStop (which sets Redis {taskID}-cancel + DB STOPPING), then marks +// the document run status as CANCEL. func (s *DocumentService) cancelDocParse(doc *entity.Document) error { - tasks, taskErr := s.taskDAO.GetByDocID(doc.ID) - if taskErr != nil { - common.Error(fmt.Sprintf("error when load task %s", doc.ID), taskErr) - return fmt.Errorf("failed to get tasks for %s: %v", doc.ID, taskErr) + task, err := s.ingestionTaskDAO.GetByDocumentID(doc.ID) + if err != nil { + return fmt.Errorf("failed to get ingestion task for %s: %v", doc.ID, err) + } + if task == nil { + return fmt.Errorf("no ingestion task found for document %s", doc.ID) } - hasUnfinishedTask := false - for _, t := range tasks { - if t.Progress < 1 { - hasUnfinishedTask = true - break - } - } - - canCancel := false - if doc.Run != nil { - if *doc.Run == string(entity.TaskStatusRunning) || *doc.Run == string(entity.TaskStatusCancel) { - canCancel = true - } - } - if hasUnfinishedTask { - canCancel = true - } - if !canCancel { - return fmt.Errorf("can't stop parsing document that has not started or already completed") - } - - // Set Redis cancel signal for each task (best-effort) - redisClient := redis.Get() - for _, t := range tasks { - if redisClient != nil { - redisClient.Set(fmt.Sprintf("%s-cancel", t.ID), "x", 0) - } + if _, err := s.ingestionTaskSvc.RequestStop(task.ID); err != nil { + return fmt.Errorf("failed to stop ingestion task %s: %v", task.ID, err) } if upErr := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"run": string(entity.TaskStatusCancel)}); upErr != nil { diff --git a/internal/service/document_test.go b/internal/service/document_test.go index d1d7b9670f..e5042d9210 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -184,6 +184,7 @@ func (fakeChatDocEngine) Close() error { func (fakeChatDocEngine) GetType() string { return "fake" } +func (fakeChatDocEngine) SupportsPageRank() bool { return false } func (fakeChatDocEngine) FilterDocIdsByMetaPushdown(context.Context, []string, []map[string]interface{}, string) []string { return nil } @@ -1147,7 +1148,7 @@ func TestStopParseDocuments_Success(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusRunning), 10, 5) - insertTestTask(t, "task-1", "doc-1") + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) @@ -1181,7 +1182,7 @@ func TestStopParseDocuments_CancelStatus(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) // Doc is already in CANCEL state — should still be accepted insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusCancel), 10, 5) - insertTestTask(t, "task-1", "doc-1") + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) @@ -1226,9 +1227,9 @@ func TestStopParseDocuments_UnfinishedTask(t *testing.T) { pushServiceDB(t, db) insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) - // Doc with Run="0" but has an unfinished task (progress < 1) → can cancel + // Doc with Run="0" → cancelDocParse calls RequestStop on the ingestion task. insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusUnstart), 10, 5) - insertTestTaskWithProgress(t, "task-1", "doc-1", 0.0) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) @@ -1293,7 +1294,7 @@ func TestStopParseDocuments_Deduplicate(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusRunning), 10, 5) - insertTestTask(t, "task-1", "doc-1") + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 586a94dc56..9d1ea29102 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -3,14 +3,23 @@ package service import ( "errors" "fmt" + "time" "ragflow/internal/common" "ragflow/internal/dao" + redis2 "ragflow/internal/engine/redis" "ragflow/internal/entity" "gorm.io/gorm" ) +// Run-count key for IngestionTaskLog.Checkpoint, consumed by +// ListAllForAdmin and IncrementRunCount to track how many times +// the task has been picked up by a worker. +const ( + stepKeyRunCount = "run_count" +) + type InvalidTaskTransitionError struct { TaskID string From string @@ -116,7 +125,7 @@ func (s *IngestionTaskService) RequestStopMany(tasks []string, ownerUserID *stri taskResponses := make([]*entity.IngestionTask, 0, len(tasks)) for _, taskID := range tasks { if ownerUserID != nil { - task, err := s.getTask(taskID) + task, err := s.GetTask(taskID) if err != nil { return nil, err } @@ -168,11 +177,20 @@ func (s *IngestionTaskService) ListAllForAdmin() ([]map[string]interface{}, erro "status": task.Status, } - latestLog, err := s.ingestionTaskLogDAO.LatestLogByTaskID(task.ID) - if err == nil && latestLog != nil && latestLog.Checkpoint != nil { - if step, ok := latestLog.Checkpoint["current_step"].(float64); ok { - showTask["step"] = int(step) + if count, ok := s.lastRunCount(task.ID); ok { + showTask["run_count"] = count + } + + showTask["component_total"] = task.ComponentTotal + if task.ComponentTotal > 0 { + progress, err := s.ingestionTaskLogDAO.AggregateProgress(task.ID, task.ComponentTotal) + if err == nil { + showTask["component_done"] = progress.Done + } else { + showTask["component_done"] = 0 } + } else { + showTask["component_done"] = 0 } showTasks = append(showTasks, showTask) @@ -181,7 +199,7 @@ func (s *IngestionTaskService) ListAllForAdmin() ([]map[string]interface{}, erro } func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTask, error) { - task, err := s.getTask(taskID) + task, err := s.GetTask(taskID) if err != nil { return nil, err } @@ -198,7 +216,7 @@ func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTas } func (s *IngestionTaskService) RequestStop(taskID string) (*entity.IngestionTask, error) { - task, err := s.getTask(taskID) + task, err := s.GetTask(taskID) if err != nil { return nil, err } @@ -206,7 +224,17 @@ func (s *IngestionTaskService) RequestStop(taskID string) (*entity.IngestionTask case common.CREATED: return s.transition(taskID, common.STOPPED) case common.RUNNING: - return s.transition(taskID, common.STOPPING) + task, err := s.transition(taskID, common.STOPPING) + if err != nil { + return nil, err + } + // Mirror Python's cancel_all_task_of: set Redis cancel flag so the + // running worker's pollCancel detects the stop immediately rather + // than waiting for the next DB poll (up to 3s). + if rc := redis2.Get(); rc != nil { + rc.Set(fmt.Sprintf("%s-cancel", taskID), "x", 1*time.Hour) + } + return task, nil default: return task, nil } @@ -222,11 +250,26 @@ func (s *IngestionTaskService) MarkFailed(taskID string) error { return err } +// MarkStopped transitions the task from STOPPING to STOPPED (terminal). +// Idempotent: returns nil if the task is already in a terminal state +// (STOPPED, COMPLETED, or FAILED). +func (s *IngestionTaskService) MarkStopped(taskID string) error { + task, err := s.GetTask(taskID) + if err != nil { + return err + } + if task.Status == common.STOPPED || task.Status == common.COMPLETED || task.Status == common.FAILED { + return nil + } + _, err = s.transition(taskID, common.STOPPED) + return err +} + func (s *IngestionTaskService) Remove(taskID string, userID *string) (*dao.TaskInfo, error) { return s.ingestionTaskDAO.Delete(taskID, userID) } -func (s *IngestionTaskService) getTask(taskID string) (*entity.IngestionTask, error) { +func (s *IngestionTaskService) GetTask(taskID string) (*entity.IngestionTask, error) { task, err := s.ingestionTaskDAO.GetByID(taskID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { @@ -260,7 +303,7 @@ func validateTransition(from, to string) error { } func (s *IngestionTaskService) newTaskStatusConflictError(taskID, expectedFrom, attemptedTo string) error { - current, err := s.getTask(taskID) + current, err := s.GetTask(taskID) if err != nil { return err } @@ -273,7 +316,7 @@ func (s *IngestionTaskService) newTaskStatusConflictError(taskID, expectedFrom, } func (s *IngestionTaskService) transition(taskID string, to string) (*entity.IngestionTask, error) { - task, err := s.getTask(taskID) + task, err := s.GetTask(taskID) if err != nil { return nil, err } @@ -355,3 +398,70 @@ func (s *IngestionTaskService) enqueueTask(taskID string) error { } return s.taskPublisher.PublishTaskMessage("tasks.RAGFLOW", taskMessage) } + +// UpdateComponentTotal records the number of components in the task's DSL +// graph - the authoritative denominator for progress percentage. +func (s *IngestionTaskService) UpdateComponentTotal(taskID string, total int) error { + return s.ingestionTaskDAO.UpdateComponentTotal(taskID, total) +} + +// RecordComponentProgress appends a component lifecycle row to +// ingestion_task_log (phase: 0 started / 1 done / 2 errored). The row's +// Checkpoint is empty; component progress and step checkpoints are distinct +// row models sharing the same table. +func (s *IngestionTaskService) RecordComponentProgress(taskID, component string, index, phase int, message string) error { + entry := &entity.IngestionTaskLog{ + TaskID: taskID, + Checkpoint: entity.JSONMap{}, + ComponentIndex: index, + Phase: phase, + Component: component, + Message: message, + } + return s.ingestionTaskLogDAO.Create(entry) +} + +// AggregateTaskProgress returns the SQL-aggregated component progress for a +// task (done/failed/running/percent against the given total denominator). +func (s *IngestionTaskService) AggregateTaskProgress(taskID string, total int) (*dao.TaskProgress, error) { + return s.ingestionTaskLogDAO.AggregateProgress(taskID, total) +} + +// lastRunCount scans all task logs (newest first) for a run_count entry, +// skipping component-progress rows whose Checkpoint is empty. It returns +// the counter and whether one was found. +func (s *IngestionTaskService) lastRunCount(taskID string) (int, bool) { + logs, err := s.ingestionTaskLogDAO.ListLogsByTaskID(taskID) + if err != nil { + return 0, false + } + for i := len(logs) - 1; i >= 0; i-- { + if count, ok := common.GetInt(logs[i].Checkpoint[stepKeyRunCount]); ok { + return count, true + } + } + return 0, false +} + +// IncrementRunCount scans existing task logs for the previous run_count +// (skipping component-progress rows that have no run_count), then INSERTS a +// new row with the bumped counter. This avoids the race where the latest log +// is a component-progress row whose empty Checkpoint would cause a parse +// failure. ListAllForAdmin reads run_count back to render the attempt number. +// +// A corrupted run_count value in an existing row is skipped (the row is +// ignored). A failure to persist the new row is best-effort (logged) and +// does not return an error — matching the legacy semantics that the run +// proceeds even if the counter write fails. +func (s *IngestionTaskService) IncrementRunCount(taskID string) error { + prevCount, _ := s.lastRunCount(taskID) + + entry := &entity.IngestionTaskLog{ + TaskID: taskID, + Checkpoint: entity.JSONMap{stepKeyRunCount: prevCount + 1}, + } + if err := s.ingestionTaskLogDAO.Create(entry); err != nil { + common.Error(fmt.Sprintf("Failed to persist run_count for task %s", taskID), err) + } + return nil +} diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index 498e2da3d5..2001dcef4b 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -153,7 +153,7 @@ func TestIngestionTaskServiceRemoveManyRemovesOwnedTasks(t *testing.T) { } } -func TestIngestionTaskServiceListAllForAdminIncludesStepAndUserEmail(t *testing.T) { +func TestIngestionTaskServiceListAllForAdminIncludesRunAndUserEmail(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) status := "1" @@ -171,7 +171,7 @@ func TestIngestionTaskServiceListAllForAdminIncludesStepAndUserEmail(t *testing. insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") if err := dao.DB.Create(&entity.IngestionTaskLog{ TaskID: "task-1", - Checkpoint: entity.JSONMap{"current_step": 3}, + Checkpoint: entity.JSONMap{"run_count": 3}, }).Error; err != nil { t.Fatalf("insert task log: %v", err) } @@ -187,8 +187,14 @@ func TestIngestionTaskServiceListAllForAdminIncludesStepAndUserEmail(t *testing. if tasks[0]["user"] != "user-1@test.com" { t.Fatalf("user = %v, want user-1@test.com", tasks[0]["user"]) } - if tasks[0]["step"] != 3 { - t.Fatalf("step = %v, want 3", tasks[0]["step"]) + if tasks[0]["run_count"] != 3 { + t.Fatalf("run_count = %v, want 3", tasks[0]["run_count"]) + } + if tasks[0]["component_total"] != 0 { + t.Fatalf("component_total = %v, want 0", tasks[0]["component_total"]) + } + if tasks[0]["component_done"] != 0 { + t.Fatalf("component_done = %v, want 0", tasks[0]["component_done"]) } } @@ -463,3 +469,282 @@ func TestIngestionTaskServiceRemoveDeletesOwnedTask(t *testing.T) { t.Fatal("task should be removed") } } + +func TestIngestionTaskServiceUpdateComponentTotalPersistsDenominator(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + if err := svc.UpdateComponentTotal("task-1", 4); err != nil { + t.Fatalf("UpdateComponentTotal failed: %v", err) + } + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.ComponentTotal != 4 { + t.Fatalf("component_total = %d, want 4", task.ComponentTotal) + } +} + +func TestIngestionTaskServiceRecordComponentProgressAppendsRow(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + t.Fatalf("RecordComponentProgress failed: %v", err) + } + logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID("task-1") + if err != nil { + t.Fatalf("list logs: %v", err) + } + if len(logs) != 1 { + t.Fatalf("expected 1 log row, got %d", len(logs)) + } + row := logs[0] + if row.Component != "Parser" || row.ComponentIndex != 0 || row.Phase != 1 || row.Message != "Parser Done" { + t.Fatalf("unexpected log row: %+v", row) + } + if len(row.Checkpoint) != 0 { + t.Fatalf("component progress row must have empty checkpoint, got %v", row.Checkpoint) + } +} + +func TestIngestionTaskServiceAggregateTaskProgressClassifiesByPhase(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + t.Fatalf("record Parser: %v", err) + } + if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 0, "Chunker Started"); err != nil { + t.Fatalf("record Chunker: %v", err) + } + agg, err := svc.AggregateTaskProgress("task-1", 2) + if err != nil { + t.Fatalf("AggregateTaskProgress failed: %v", err) + } + if agg.Done != 1 || agg.Running != 1 || agg.Failed != 0 { + t.Fatalf("aggregate = %+v, want Done=1 Running=1 Failed=0", agg) + } + if agg.Percent != 50 { + t.Fatalf("percent = %v, want 50", agg.Percent) + } +} + +func TestIngestionTaskServiceIncrementRunCountInitializesAndBumps(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("IncrementRunCount (first call) failed: %v", err) + } + run, ok := findLastRunCount(t, "task-1") + if !ok || run != 1 { + t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) + } + + // Second call bumps the existing counter to 2. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("IncrementRunCount (second call) failed: %v", err) + } + run, _ = findLastRunCount(t, "task-1") + if run != 2 { + t.Fatalf("run_count after second bump = %v, want 2", run) + } +} + +func TestIngestionTaskServiceIncrementRunCountSkippedCorruptedRunCount(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Create(&entity.IngestionTaskLog{ + TaskID: "task-1", + Checkpoint: entity.JSONMap{"run_count": "not-a-number"}, + }).Error; err != nil { + t.Fatalf("insert bad task log: %v", err) + } + + svc := NewIngestionTaskService() + // Corrupted value is skipped; a fresh run_count=1 row is created. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("IncrementRunCount should skip corrupted value, got: %v", err) + } + run, ok := findLastRunCount(t, "task-1") + if !ok || run != 1 { + t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) + } +} + +func TestIngestionTaskServiceIncrementRunCountRecoversFromComponentProgressLog(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + // Simulate a previous run that created some component-progress logs + // but died before recording a run_count row. The latest log has no run_count. + svc := NewIngestionTaskService() + if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + t.Fatalf("record Parser: %v", err) + } + if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 1, "Chunker Done"); err != nil { + t.Fatalf("record Chunker: %v", err) + } + // Verify latest log has empty checkpoint (no run_count). + latest, err := dao.NewIngestionTaskLogDAO().LatestLogByTaskID("task-1") + if err != nil { + t.Fatalf("load latest: %v", err) + } + if len(latest.Checkpoint) != 0 { + t.Fatalf("component-progress row should have empty checkpoint, got %v", latest.Checkpoint) + } + + // IncrementRunCount should create a new row with run_count=1. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("IncrementRunCount failed: %v", err) + } + run, ok := findLastRunCount(t, "task-1") + if !ok || run != 1 { + t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) + } + + // AggregateProgress should still work (run_count row with component="" + // has phase=0, which doesn't affect counts). + agg, err := svc.AggregateTaskProgress("task-1", 2) + if err != nil { + t.Fatalf("AggregateTaskProgress: %v", err) + } + if agg.Done != 2 { + t.Fatalf("Done = %d, want 2 (run_count row didn't interfere)", agg.Done) + } +} + +func TestIngestionTaskServiceIncrementRunCountAccumulatesAcrossRetries(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + + svc := NewIngestionTaskService() + + // First attempt: IncrementRunCount creates run_count=1. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("first IncrementRunCount: %v", err) + } + // Simulate first run: some components progress, then failure. + if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + t.Fatalf("record Parser: %v", err) + } + + // Second attempt (retry): should find previous run_count=1 and create row with run_count=2. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("second IncrementRunCount: %v", err) + } + // More progress, then failure. + if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 1, "Chunker Done"); err != nil { + t.Fatalf("record Chunker: %v", err) + } + + // Third attempt (retry): should find previous run_count=2 and create row with run_count=3. + if err := svc.IncrementRunCount("task-1"); err != nil { + t.Fatalf("third IncrementRunCount: %v", err) + } + + run, ok := findLastRunCount(t, "task-1") + if !ok || run != 3 { + t.Fatalf("run_count = %v (ok=%v), want 3", run, ok) + } + + // ListAllForAdmin should still pick up the correct run_count. + allLogs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID("task-1") + if err != nil { + t.Fatalf("list all logs: %v", err) + } + t.Logf("total log rows: %d", len(allLogs)) +} + +func TestIngestionTaskServiceMarkStoppedTransitionsStoppingTask(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + // Override to STOPPING. + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1"). + Update("status", common.STOPPING).Error; err != nil { + t.Fatalf("set task STOPPING: %v", err) + } + + svc := NewIngestionTaskService() + if err := svc.MarkStopped("task-1"); err != nil { + t.Fatalf("MarkStopped failed: %v", err) + } + + task, err := dao.NewIngestionTaskDAO().GetByID("task-1") + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.STOPPED { + t.Fatalf("task status = %s, want STOPPED", task.Status) + } +} + +func TestIngestionTaskServiceMarkStoppedIdempotentOnAlreadyStopped(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + if err := dao.DB.Model(&entity.IngestionTask{}).Where("id = ?", "task-1"). + Update("status", common.STOPPED).Error; err != nil { + t.Fatalf("set task STOPPED: %v", err) + } + + svc := NewIngestionTaskService() + if err := svc.MarkStopped("task-1"); err != nil { + t.Fatalf("MarkStopped on already STOPPED task should be idempotent, got: %v", err) + } +} + +func TestDocumentServiceUpdateRunProgressMirrorsFields(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestDoc(t, "doc-1", "kb-1", 0, 0) + + svc := testDocumentService(t) + if err := svc.UpdateRunProgress("doc-1", 0.5, "1", "halfway"); err != nil { + t.Fatalf("UpdateRunProgress failed: %v", err) + } + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("load document: %v", err) + } + if doc.Progress != 0.5 { + t.Fatalf("progress = %v, want 0.5", doc.Progress) + } + if doc.Run == nil || *doc.Run != "1" { + t.Fatalf("run = %v, want 1", doc.Run) + } + if doc.ProgressMsg == nil || *doc.ProgressMsg != "halfway" { + t.Fatalf("progress_msg = %v, want halfway", doc.ProgressMsg) + } +} + +// findLastRunCount scans all logs for a task and returns the latest run_count. +// It avoids non-deterministic LatestLogByTaskID ordering when multiple rows +// share the same create_time. +func findLastRunCount(t *testing.T, taskID string) (int, bool) { + t.Helper() + logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) + if err != nil { + t.Fatalf("list logs for %s: %v", taskID, err) + } + for i := len(logs) - 1; i >= 0; i-- { + if count, ok := common.GetInt(logs[i].Checkpoint[stepKeyRunCount]); ok { + return count, true + } + } + return 0, false +}