mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 22:25:41 +08:00
Refactor: drop dead canvas runtime selector and tokenizer embedding wiring (#16809)
Two refactors on the Go port (agent-go-port): - Remove the dead per-tenant canvas runtime selector (write-only Redis scaffolding with no runtime callers) and its dependent metrics/admin code. - Move the tokenizer embedding-model id from the shared ingestion globals into a Tokenizer-scoped setup, and wire the production embedder resolver in the ingestion task package. 32 files changed, 861 insertions, 1228 deletions.
This commit is contained in:
@@ -26,7 +26,6 @@ import (
|
||||
"ragflow/internal/admin"
|
||||
"ragflow/internal/agent/audio"
|
||||
"ragflow/internal/agent/canvas"
|
||||
"ragflow/internal/agent/runtime"
|
||||
agenttool "ragflow/internal/agent/tool"
|
||||
"ragflow/internal/handler"
|
||||
ingestion "ragflow/internal/ingestion/service"
|
||||
@@ -806,23 +805,6 @@ func startServer(config *server.Config) {
|
||||
docDAO,
|
||||
docEngine,
|
||||
)
|
||||
// Per-tenant canvas-runtime override selector, backed by the
|
||||
// existing Redis client and the global logger. The handler is
|
||||
// ALWAYS constructed, even when Redis is briefly unavailable at
|
||||
// startup, so the POST /api/v1/admin/canvas-runtime/:tenant_id
|
||||
// endpoint stays registered and returns the explicit
|
||||
// ErrSelectorNotConfigured (HTTP 500) path until Redis recovers.
|
||||
// Skipping handler construction when rdb == nil silently removed
|
||||
// the route until the next process restart, so a transient
|
||||
// Redis blip at boot stranded canary operators with a 404 they
|
||||
// could not diagnose from the client side. Keep the route hot.
|
||||
var adminRuntimeSelector *runtime.Selector
|
||||
if redisClient := redis.Get(); redisClient != nil {
|
||||
if rdb := redisClient.GetClient(); rdb != nil {
|
||||
adminRuntimeSelector = runtime.NewSelector(rdb, common.Logger)
|
||||
}
|
||||
}
|
||||
adminRuntimeHandler := handler.NewAdminRuntimeHandler(adminRuntimeSelector)
|
||||
componentsSvc := service.NewComponentsService()
|
||||
componentsHandler := handler.NewComponentsHandler(componentsSvc)
|
||||
|
||||
@@ -853,7 +835,6 @@ func startServer(config *server.Config) {
|
||||
pluginHandler,
|
||||
modelHandler,
|
||||
fileCommitHandler,
|
||||
adminRuntimeHandler,
|
||||
openaiChatHandler,
|
||||
botHandler,
|
||||
componentsHandler)
|
||||
|
||||
@@ -75,6 +75,15 @@ type CompileOptions struct {
|
||||
// graph does not pause on completion and force an extra, needless
|
||||
// ResumeWithData round.
|
||||
InterruptAfterNonTerminal bool
|
||||
// SetupOverrides is a run-level override map keyed by cpnID. Each
|
||||
// component's `params["setups"]` is merged only with its own entry
|
||||
// (an arbitrary string-keyed map); the override wins on top-level key
|
||||
// collision (see node_body.go mergeSetups). Components absent from the
|
||||
// map are left untouched. Used by the ingestion pipeline so a single
|
||||
// Pipeline.Run can override the DSL-baked component setups without
|
||||
// mutating the shared *Canvas (see node_body.go applySetupOverrides /
|
||||
// mergeSetups).
|
||||
SetupOverrides map[string]any
|
||||
}
|
||||
|
||||
// CompileOption mutates a CompileOptions before the compile runs.
|
||||
@@ -121,6 +130,14 @@ func WithInterruptAfterNonTerminalCpn() CompileOption {
|
||||
return func(o *CompileOptions) { o.InterruptAfterNonTerminal = true }
|
||||
}
|
||||
|
||||
// WithSetupOverrides attaches a run-level setups override map (keyed by
|
||||
// cpnID) to the compile. Each component's `params["setups"]` is merged with
|
||||
// its own entry at compile time (run-level wins on key collision, see
|
||||
// node_body.go mergeSetups). Passing nil is a no-op.
|
||||
func WithSetupOverrides(m map[string]any) CompileOption {
|
||||
return func(o *CompileOptions) { o.SetupOverrides = m }
|
||||
}
|
||||
|
||||
// Compile builds the eino Workflow from the Canvas and returns the
|
||||
// compiled Runnable. State pre/post handlers are wired inside BuildWorkflow
|
||||
// (see scheduler.go). Checkpoint store + serializer are wired here as
|
||||
@@ -196,6 +213,14 @@ func Compile(ctx context.Context, c *Canvas, opts ...CompileOption) (*CompiledCa
|
||||
}
|
||||
}
|
||||
|
||||
// Thread the run-level setups override (if any) into ctx so each
|
||||
// component's `params["setups"]` is merged with its own entry inside
|
||||
// buildNodeBody. The override is keyed by cpnID; the canvas package
|
||||
// never imports ingestion.
|
||||
if cfg.SetupOverrides != nil {
|
||||
ctx = withSetupOverrides(ctx, cfg.SetupOverrides)
|
||||
}
|
||||
|
||||
wf, err := BuildWorkflow(ctx, c)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("canvas: build workflow: %w", err)
|
||||
|
||||
128
internal/agent/canvas/compile_setup_override_test.go
Normal file
128
internal/agent/canvas/compile_setup_override_test.go
Normal file
@@ -0,0 +1,128 @@
|
||||
//
|
||||
// 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 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 canvas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// TestCompile_SetupOverrides exercises the full canvas-level wiring:
|
||||
//
|
||||
// canvas.Compile(ctx, dsl, WithSetupOverrides(override))
|
||||
//
|
||||
// threads the cpnID-keyed override through ctx into
|
||||
// BuildWorkflow → buildNodeBody → applySetupOverrides → mergeSetups, so
|
||||
// each component's factory receives its own merged params["setups"]. Only
|
||||
// the entry for a component's own cpnID applies; components absent from the
|
||||
// override map keep their base params (no spurious "setups" injected).
|
||||
func TestCompile_SetupOverrides(t *testing.T) {
|
||||
captured := map[string]map[string]any{} // component_name -> params received by factory
|
||||
factory := func(name string, params map[string]any) (runtime.Component, error) {
|
||||
// Deep-shallow copy so later mutations don't hide what the
|
||||
// factory actually received.
|
||||
cp := make(map[string]any, len(params))
|
||||
for k, v := range params {
|
||||
cp[k] = v
|
||||
}
|
||||
captured[name] = cp
|
||||
return &stubComponent{params: cp}, nil
|
||||
}
|
||||
|
||||
dsl := &Canvas{
|
||||
Components: map[string]CanvasComponent{
|
||||
"parser_0": {
|
||||
Obj: CanvasComponentObj{
|
||||
ComponentName: "Parser",
|
||||
Params: map[string]any{
|
||||
"setups": map[string]any{
|
||||
"pdf": map[string]any{
|
||||
"output_format": "one",
|
||||
"parse_method": "naive",
|
||||
},
|
||||
"doc": map[string]any{
|
||||
"output_format": "one",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Upstream: []string{},
|
||||
Downstream: []string{"sink_0"},
|
||||
},
|
||||
"sink_0": {
|
||||
Obj: CanvasComponentObj{
|
||||
ComponentName: "Sink",
|
||||
Params: map[string]any{"name": "sink"},
|
||||
},
|
||||
Upstream: []string{"parser_0"},
|
||||
Downstream: []string{},
|
||||
},
|
||||
},
|
||||
Path: []string{"parser_0", "sink_0"},
|
||||
}
|
||||
|
||||
// Override is keyed by cpnID. Only "parser_0" is present, so its
|
||||
// "pdf" entry is fully replaced (parse_method dropped) and "docx" is
|
||||
// injected; "doc" survives from the base. "sink_0" is absent and must
|
||||
// keep its base params untouched.
|
||||
override := map[string]any{
|
||||
"parser_0": map[string]any{
|
||||
"pdf": map[string]any{
|
||||
"output_format": "detailed",
|
||||
},
|
||||
"docx": map[string]any{
|
||||
"output_format": "one",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ctx := WithComponentFactory(context.Background(), factory)
|
||||
if _, err := Compile(ctx, dsl, WithSetupOverrides(override)); err != nil {
|
||||
t.Fatalf("Compile: %v", err)
|
||||
}
|
||||
|
||||
// parser_0: override merged into params["setups"] before the factory.
|
||||
got, ok := captured["Parser"]["setups"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("Parser factory did not receive a setups map: %#v", captured["Parser"])
|
||||
}
|
||||
// doc preserved from base.
|
||||
if v, _ := got["doc"].(map[string]any); v == nil || v["output_format"] != "one" {
|
||||
t.Errorf("doc should be preserved from base: %#v", got["doc"])
|
||||
}
|
||||
// docx injected by the override (was absent in base).
|
||||
if v, _ := got["docx"].(map[string]any); v == nil || v["output_format"] != "one" {
|
||||
t.Errorf("docx injection missing: %#v", got["docx"])
|
||||
}
|
||||
// pdf: override fully replaces the base entry (shallow merge).
|
||||
pdf, _ := got["pdf"].(map[string]any)
|
||||
if pdf == nil {
|
||||
t.Fatalf("pdf setup missing: %#v", got)
|
||||
}
|
||||
if pdf["output_format"] != "detailed" {
|
||||
t.Errorf("pdf.output_format not overridden: %#v", pdf)
|
||||
}
|
||||
if _, ok := pdf["parse_method"]; ok {
|
||||
t.Errorf("pdf.parse_method should be dropped by shallow merge: %#v", pdf)
|
||||
}
|
||||
|
||||
// sink_0: absent from the override → no setups key injected.
|
||||
if _, ok := captured["Sink"]["setups"]; ok {
|
||||
t.Errorf("Sink should not receive a setups key: %#v", captured["Sink"])
|
||||
}
|
||||
}
|
||||
@@ -74,7 +74,75 @@ type nodeBodyFn = func(ctx context.Context, in map[string]any) (map[string]any,
|
||||
// Outputs bucket. UserFillUpNodeBody tags its output itself so the
|
||||
// interrupt-driven branch still attributes the resume payload to the
|
||||
// right cpn.
|
||||
// ctxKeySetupOverrides carries the run-level setups override map into
|
||||
// BuildWorkflow so a component's `params["setups"]` can be merged with it
|
||||
// at compile time. The map is keyed by cpnID; each component only sees the
|
||||
// entry for its own id (an arbitrary string-keyed map). It mirrors the ctx
|
||||
// plumbing used for the per-run component factory
|
||||
// (componentFactoryFromContext): the override is threaded through
|
||||
// canvas.Compile → BuildWorkflow → buildNodeBody without the canvas
|
||||
// package ever importing the ingestion layer.
|
||||
const ctxKeySetupOverrides ctxKey = "canvas_setup_overrides"
|
||||
|
||||
// withSetupOverrides attaches a run-level setups override map to ctx. It is
|
||||
// a no-op when m is nil so callers can pass a possibly-nil run parameter
|
||||
// straight through.
|
||||
func withSetupOverrides(ctx context.Context, m map[string]any) context.Context {
|
||||
if m == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, ctxKeySetupOverrides, m)
|
||||
}
|
||||
|
||||
func setupOverridesFromContext(ctx context.Context) map[string]any {
|
||||
m, _ := ctx.Value(ctxKeySetupOverrides).(map[string]any)
|
||||
return m
|
||||
}
|
||||
|
||||
// applySetupOverrides returns a clone of params with the per-component
|
||||
// setups override (already resolved for this cpnID by the caller) merged
|
||||
// into params["setups"]. The override wins on top-level key collisions. The
|
||||
// original params map is never mutated — the merge result is a fresh map —
|
||||
// because the params come from the shared *Canvas and a per-run override
|
||||
// must not leak into the next Run on the same Pipeline.
|
||||
func applySetupOverrides(params, cpnOverride map[string]any) map[string]any {
|
||||
if len(cpnOverride) == 0 {
|
||||
return params
|
||||
}
|
||||
out := make(map[string]any, len(params)+1)
|
||||
for k, v := range params {
|
||||
out[k] = v
|
||||
}
|
||||
base, _ := out["setups"].(map[string]any)
|
||||
out["setups"] = mergeSetups(base, cpnOverride)
|
||||
return out
|
||||
}
|
||||
|
||||
// mergeSetups merges a component-level setups map (base) with a run-level
|
||||
// override map. The maps are arbitrary string-keyed maps; when the same
|
||||
// top-level key exists in both, the override value wins (a full replacement
|
||||
// of that entry). The merge is shallow: only the top-level key-value pairs
|
||||
// are considered. (The Parser component happens to use file-type keys such
|
||||
// as "pdf"/"docx" as one example, but that is not required by this merge.)
|
||||
func mergeSetups(base, override map[string]any) map[string]any {
|
||||
merged := make(map[string]any, len(base)+len(override))
|
||||
for k, v := range base {
|
||||
merged[k] = v
|
||||
}
|
||||
for k, ov := range override {
|
||||
merged[k] = ov
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func buildNodeBody(ctx context.Context, cpnID, name string, params map[string]any) (nodeBodyFn, error) {
|
||||
if overrides := setupOverridesFromContext(ctx); len(overrides) > 0 {
|
||||
// overrides is keyed by cpnID; a component only sees its own
|
||||
// entry. Components absent from the map are left untouched.
|
||||
if cpnOverride, ok := overrides[cpnID].(map[string]any); ok && len(cpnOverride) > 0 {
|
||||
params = applySetupOverrides(params, cpnOverride)
|
||||
}
|
||||
}
|
||||
if isLegacyNoOp(name) {
|
||||
return legacyNoOpBody(cpnID), nil
|
||||
}
|
||||
|
||||
152
internal/agent/canvas/node_body_setup_override_test.go
Normal file
152
internal/agent/canvas/node_body_setup_override_test.go
Normal file
@@ -0,0 +1,152 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in 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 canvas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// stubComponent records the params it was constructed with so a test can
|
||||
// assert that buildNodeBody forwarded the merged setups. Its Invoke is a
|
||||
// no-op echo.
|
||||
type stubComponent struct {
|
||||
params map[string]any
|
||||
}
|
||||
|
||||
func (s *stubComponent) Invoke(_ context.Context, in map[string]any) (map[string]any, error) {
|
||||
return in, nil
|
||||
}
|
||||
|
||||
// TestBuildNodeBody_SetupOverrides asserts that a run-level setups override
|
||||
// (threaded via ctx, keyed by cpnID) is merged into the component's
|
||||
// `params["setups"]` before the factory is called. Only the entry for the
|
||||
// component's own cpnID applies. The merge is shallow: a top-level key
|
||||
// present in the override fully replaces the base entry for that key
|
||||
// (no inner deep-merge), while base keys absent from the override survive.
|
||||
func TestBuildNodeBody_SetupOverrides(t *testing.T) {
|
||||
captured := map[string]any{}
|
||||
factory := func(name string, params map[string]any) (runtime.Component, error) {
|
||||
// Deep-shallow copy so later mutations by the builder don't
|
||||
// hide what the factory actually received.
|
||||
cp := make(map[string]any, len(params))
|
||||
for k, v := range params {
|
||||
cp[k] = v
|
||||
}
|
||||
captured = cp
|
||||
return &stubComponent{params: cp}, nil
|
||||
}
|
||||
|
||||
baseParams := map[string]any{
|
||||
"setups": map[string]any{
|
||||
"pdf": map[string]any{
|
||||
"output_format": "one",
|
||||
"parse_method": "naive",
|
||||
},
|
||||
"doc": map[string]any{
|
||||
"output_format": "one",
|
||||
},
|
||||
},
|
||||
}
|
||||
// Run-level override is keyed by cpnID. For "cpn-parser": the whole
|
||||
// "pdf" entry is replaced (parse_method is dropped), and a new "docx"
|
||||
// entry is injected. "doc" is untouched because it is absent from the
|
||||
// override. The entry for a different cpnID must not leak in.
|
||||
override := map[string]any{
|
||||
"cpn-parser": map[string]any{
|
||||
"pdf": map[string]any{
|
||||
"output_format": "detailed",
|
||||
},
|
||||
"docx": map[string]any{
|
||||
"output_format": "one",
|
||||
},
|
||||
},
|
||||
"cpn-other": map[string]any{
|
||||
"pdf": map[string]any{
|
||||
"output_format": "should-not-apply",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ctx := WithComponentFactory(context.Background(), factory)
|
||||
ctx = withSetupOverrides(ctx, override)
|
||||
|
||||
body, err := buildNodeBody(ctx, "cpn-parser", "Parser", baseParams)
|
||||
if err != nil {
|
||||
t.Fatalf("buildNodeBody: %v", err)
|
||||
}
|
||||
if _, err := body(context.Background(), map[string]any{"x": 1}); err != nil {
|
||||
t.Fatalf("body: %v", err)
|
||||
}
|
||||
|
||||
got, ok := captured["setups"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("factory did not receive a setups map: %#v", captured["setups"])
|
||||
}
|
||||
// doc is untouched (absent from the override).
|
||||
if v, _ := got["doc"].(map[string]any); v == nil || v["output_format"] != "one" {
|
||||
t.Errorf("doc should be preserved from base: %#v", got["doc"])
|
||||
}
|
||||
// docx injected by the override (was absent in base).
|
||||
if v, _ := got["docx"].(map[string]any); v == nil || v["output_format"] != "one" {
|
||||
t.Errorf("docx injection missing: %#v", got["docx"])
|
||||
}
|
||||
// pdf: the override fully replaces the base entry, so output_format is
|
||||
// overridden and the base-only parse_method is dropped (shallow merge).
|
||||
pdf, _ := got["pdf"].(map[string]any)
|
||||
if pdf == nil {
|
||||
t.Fatalf("pdf setup missing: %#v", got)
|
||||
}
|
||||
if pdf["output_format"] != "detailed" {
|
||||
t.Errorf("pdf.output_format not overridden: %#v", pdf)
|
||||
}
|
||||
if _, ok := pdf["parse_method"]; ok {
|
||||
t.Errorf("pdf.parse_method should be dropped by shallow merge: %#v", pdf)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildNodeBody_SetupOverridesNilIsNoOp asserts that with no override in
|
||||
// ctx the component receives exactly its base params (no spurious setups key
|
||||
// injected, and the original map is untouched).
|
||||
func TestBuildNodeBody_SetupOverridesNilIsNoOp(t *testing.T) {
|
||||
captured := map[string]any{}
|
||||
factory := func(name string, params map[string]any) (runtime.Component, error) {
|
||||
cp := make(map[string]any, len(params))
|
||||
for k, v := range params {
|
||||
cp[k] = v
|
||||
}
|
||||
captured = cp
|
||||
return &stubComponent{params: cp}, nil
|
||||
}
|
||||
baseParams := map[string]any{"name": "x"}
|
||||
|
||||
ctx := WithComponentFactory(context.Background(), factory)
|
||||
body, err := buildNodeBody(ctx, "cpn", "Parser", baseParams)
|
||||
if err != nil {
|
||||
t.Fatalf("buildNodeBody: %v", err)
|
||||
}
|
||||
if _, err := body(context.Background(), map[string]any{}); err != nil {
|
||||
t.Fatalf("body: %v", err)
|
||||
}
|
||||
if _, ok := captured["setups"]; ok {
|
||||
t.Errorf("setups key should not be injected when no override is present: %#v", captured)
|
||||
}
|
||||
if _, ok := baseParams["setups"]; ok {
|
||||
t.Errorf("base params map must not be mutated: %#v", baseParams)
|
||||
}
|
||||
}
|
||||
@@ -1,102 +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 runtime
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
)
|
||||
|
||||
// outcomeSuccess / outcomeError / outcomeCancelled are the only outcome
|
||||
// label values the canvas-run metric emits. Keeping the set closed lets
|
||||
// downstream alerts reason about the cardinality.
|
||||
const (
|
||||
OutcomeSuccess = "success"
|
||||
OutcomeError = "error"
|
||||
OutcomeCancelled = "cancelled"
|
||||
)
|
||||
|
||||
var (
|
||||
canvasRunsTotal = prometheus.NewCounterVec(
|
||||
prometheus.CounterOpts{
|
||||
Name: "ragflow_canvas_runs_total",
|
||||
Help: "Total canvas runs by runtime mode and outcome.",
|
||||
},
|
||||
[]string{"runtime", "outcome"},
|
||||
)
|
||||
|
||||
canvasRunDuration = prometheus.NewHistogramVec(
|
||||
prometheus.HistogramOpts{
|
||||
Name: "ragflow_canvas_run_duration_seconds",
|
||||
Help: "Canvas run latency in seconds.",
|
||||
Buckets: prometheus.DefBuckets,
|
||||
},
|
||||
[]string{"runtime"},
|
||||
)
|
||||
|
||||
registerOnce sync.Once
|
||||
)
|
||||
|
||||
// init registers the package's metrics with the default Prometheus
|
||||
// registry. The sync.Once guard means tests that re-import this package
|
||||
// (or any future entry point that also imports it) won't trigger
|
||||
// "duplicate metrics collector registration" panics. If a test wants a
|
||||
// clean registry, call ResetMetrics().
|
||||
func init() {
|
||||
registerOnce.Do(func() {
|
||||
prometheus.MustRegister(canvasRunsTotal, canvasRunDuration)
|
||||
})
|
||||
}
|
||||
|
||||
// ResetMetricsForTesting unregisters the package metrics from the default
|
||||
// Prometheus registry, clears any recorded samples, and re-registers them
|
||||
// so the next ObserveRun call sees a clean slate. Intended for unit tests
|
||||
// that need to assert on freshly-registered metrics.
|
||||
func ResetMetricsForTesting() {
|
||||
prometheus.DefaultRegisterer.Unregister(canvasRunsTotal)
|
||||
prometheus.DefaultRegisterer.Unregister(canvasRunDuration)
|
||||
canvasRunsTotal.Reset()
|
||||
canvasRunDuration.Reset()
|
||||
registerOnce = sync.Once{}
|
||||
registerOnce.Do(func() {
|
||||
prometheus.MustRegister(canvasRunsTotal, canvasRunDuration)
|
||||
})
|
||||
}
|
||||
|
||||
// ObserveRun emits a counter + histogram observation for one canvas run.
|
||||
// The runtime label is the string form of RuntimeMode; the outcome label
|
||||
// must be one of the Outcome* constants. Negative or zero durations are
|
||||
// dropped from the histogram to avoid skewing percentile math.
|
||||
//
|
||||
// An empty runtime defaults to the process-wide Default() so the metric
|
||||
// tag matches what Selector.Select would have returned for the same
|
||||
// tenant (review follow-up M4). Callers that need to record the
|
||||
// Python-routed-fallback case explicitly should pass RuntimePython.
|
||||
func ObserveRun(runtime RuntimeMode, outcome string, duration time.Duration) {
|
||||
if runtime == "" {
|
||||
runtime = Default()
|
||||
}
|
||||
if outcome == "" {
|
||||
outcome = OutcomeError
|
||||
}
|
||||
canvasRunsTotal.WithLabelValues(string(runtime), outcome).Inc()
|
||||
if duration > 0 {
|
||||
canvasRunDuration.WithLabelValues(string(runtime)).Observe(duration.Seconds())
|
||||
}
|
||||
}
|
||||
@@ -1,116 +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
|
||||
//
|
||||
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/testutil"
|
||||
dto "github.com/prometheus/client_model/go"
|
||||
)
|
||||
|
||||
// labelsMatch reports whether the metric carries a label named "runtime"
|
||||
// with the supplied value.
|
||||
func labelsMatch(pairs []*dto.LabelPair, wantRuntime string) bool {
|
||||
for _, p := range pairs {
|
||||
if p.GetName() == "runtime" && p.GetValue() == wantRuntime {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestObserveRun_IncrementsCounter(t *testing.T) {
|
||||
// ResetMetricsForTesting re-registers the metrics on the default
|
||||
// registry so testutil can read them.
|
||||
ResetMetricsForTesting()
|
||||
|
||||
ObserveRun(RuntimeGo, OutcomeSuccess, 250*time.Millisecond)
|
||||
ObserveRun(RuntimeGo, OutcomeSuccess, 500*time.Millisecond)
|
||||
ObserveRun(RuntimePython, OutcomeError, 750*time.Millisecond)
|
||||
|
||||
if got := testutil.ToFloat64(canvasRunsTotal.WithLabelValues("go", "success")); got != 2 {
|
||||
t.Errorf("go/success counter = %v, want 2", got)
|
||||
}
|
||||
if got := testutil.ToFloat64(canvasRunsTotal.WithLabelValues("python", "error")); got != 1 {
|
||||
t.Errorf("python/error counter = %v, want 1", got)
|
||||
}
|
||||
if got := testutil.ToFloat64(canvasRunsTotal.WithLabelValues("python", "success")); got != 0 {
|
||||
t.Errorf("python/success counter = %v, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserveRun_RecordsDuration(t *testing.T) {
|
||||
ResetMetricsForTesting()
|
||||
|
||||
ObserveRun(RuntimeGo, OutcomeSuccess, time.Second)
|
||||
ObserveRun(RuntimeGo, OutcomeSuccess, 2*time.Second)
|
||||
|
||||
gathered, err := prometheus.DefaultGatherer.Gather()
|
||||
if err != nil {
|
||||
t.Fatalf("Gather: %v", err)
|
||||
}
|
||||
|
||||
var found bool
|
||||
for _, mf := range gathered {
|
||||
if !strings.Contains(mf.GetName(), "canvas_run_duration_seconds") {
|
||||
continue
|
||||
}
|
||||
for _, m := range mf.Metric {
|
||||
if labelsMatch(m.Label, "go") {
|
||||
found = true
|
||||
if m.Histogram == nil || m.Histogram.GetSampleCount() != 2 {
|
||||
t.Errorf("go histogram count = %v, want 2", m.Histogram.GetSampleCount())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("did not find canvas_run_duration_seconds metric for runtime=go")
|
||||
}
|
||||
}
|
||||
|
||||
func TestObserveRun_NormalisesEmptyArgs(t *testing.T) {
|
||||
ResetMetricsForTesting()
|
||||
|
||||
// Pin the env-driven default to Python so the test is hermetic
|
||||
// regardless of the host environment. The assertion is "the
|
||||
// empty-runtime label equals Default()", so we set the env to
|
||||
// make the expected value explicit.
|
||||
t.Setenv("RAGFLOW_CANVAS_DEFAULT_RUNTIME", string(RuntimePython))
|
||||
ResetDefaultCache()
|
||||
|
||||
ObserveRun("", "", 0) // all empty — should default to python/error and not observe histogram
|
||||
|
||||
if got := testutil.ToFloat64(canvasRunsTotal.WithLabelValues("python", "error")); got != 1 {
|
||||
t.Errorf("defaulted counter = %v, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestObserveRun_DefaultFallsToGoAfterPhase7 documents the
|
||||
// empty-runtime fallback: the metric's fallback now follows
|
||||
// selector.Default() (which is RuntimeGo), so the runtime label
|
||||
// is consistent with what Selector.Select would have returned
|
||||
// for the same tenant.
|
||||
func TestObserveRun_DefaultFallsToGoAfterPhase7(t *testing.T) {
|
||||
ResetMetricsForTesting()
|
||||
t.Setenv("RAGFLOW_CANVAS_DEFAULT_RUNTIME", "")
|
||||
ResetDefaultCache()
|
||||
|
||||
ObserveRun("", OutcomeError, 0)
|
||||
|
||||
if got := testutil.ToFloat64(canvasRunsTotal.WithLabelValues(string(Default()), "error")); got != 1 {
|
||||
t.Errorf("defaulted counter = %v, want 1 for runtime=%s outcome=error", got, Default())
|
||||
}
|
||||
}
|
||||
@@ -1,168 +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 runtime implements per-tenant runtime selection for the
|
||||
// agent canvas port.
|
||||
//
|
||||
// Two pieces live in this package:
|
||||
//
|
||||
// - Selector (this file): reads/writes the per-tenant runtime
|
||||
// override in Redis. The default is RuntimeGo; per-tenant
|
||||
// overrides still let operators force a tenant back to Python
|
||||
// during the agent_api.py deprecation window.
|
||||
// - Metrics (metrics.go): Prometheus counter + histogram for
|
||||
// per-run observation, keyed by runtime mode.
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"ragflow/internal/common"
|
||||
"sync"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// RuntimeMode identifies which agent-canvas runtime implementation
|
||||
// serves a given tenant. Supports "go" and "python"; "auto" is
|
||||
// reserved for future adaptive policies.
|
||||
type RuntimeMode string
|
||||
|
||||
const (
|
||||
// RuntimeGo routes the tenant to the Go-side eino
|
||||
// implementation. This is the process-wide default.
|
||||
RuntimeGo RuntimeMode = "go"
|
||||
// RuntimePython routes the tenant to the legacy Python
|
||||
// agent_api.py implementation. Retained for the 1-release
|
||||
// deprecation window; per-tenant overrides via Selector.Set
|
||||
// can still force a tenant to this mode.
|
||||
RuntimePython RuntimeMode = "python"
|
||||
// RuntimeAuto defers to the per-tenant override, then to the
|
||||
// process-wide Default(). It exists as a sentinel for clients that
|
||||
// want explicit "I don't care, pick for me" semantics.
|
||||
RuntimeAuto RuntimeMode = "auto"
|
||||
)
|
||||
|
||||
// defaultEnvKey is the environment variable consulted by Default() when no
|
||||
// override is registered for a tenant.
|
||||
const defaultEnvKey = "RAGFLOW_CANVAS_DEFAULT_RUNTIME"
|
||||
|
||||
// overrideKeyPrefix is the Redis key namespace for per-tenant runtime
|
||||
// overrides. Final keys look like "tenant_canvas_runtime:<tenantID>".
|
||||
const overrideKeyPrefix = "tenant_canvas_runtime:"
|
||||
|
||||
var (
|
||||
defaultOnce sync.Once
|
||||
defaultMode RuntimeMode
|
||||
)
|
||||
|
||||
// Default returns the process-wide default runtime mode.
|
||||
//
|
||||
// The default is Go. The per-tenant override (via Selector.Set)
|
||||
// can still force a tenant back to Python for the 1-release
|
||||
// deprecation window of agent_api.py.
|
||||
//
|
||||
// The value is read once from the RAGFLOW_CANVAS_DEFAULT_RUNTIME env var;
|
||||
// subsequent calls return the cached result. Unknown env values fall back
|
||||
// to RuntimeGo (the new default) so a misconfig still lands on the Go path.
|
||||
func Default() RuntimeMode {
|
||||
defaultOnce.Do(func() {
|
||||
raw := common.GetEnv(defaultEnvKey)
|
||||
switch RuntimeMode(raw) {
|
||||
case RuntimeGo, RuntimePython, RuntimeAuto:
|
||||
defaultMode = RuntimeMode(raw)
|
||||
default:
|
||||
defaultMode = RuntimeGo
|
||||
}
|
||||
})
|
||||
return defaultMode
|
||||
}
|
||||
|
||||
// ResetDefaultCache clears the cached default-mode value. Test-only helper.
|
||||
func ResetDefaultCache() {
|
||||
defaultOnce = sync.Once{}
|
||||
defaultMode = ""
|
||||
}
|
||||
|
||||
// Selector resolves the runtime mode for a tenant at request time. It is
|
||||
// safe for concurrent use.
|
||||
type Selector struct {
|
||||
redis *redis.Client
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewSelector constructs a Selector backed by the supplied Redis client. A
|
||||
// nil logger is replaced with zap.NewNop() so callers in tests can omit it.
|
||||
func NewSelector(rdb *redis.Client, logger *zap.Logger) *Selector {
|
||||
if logger == nil {
|
||||
logger = zap.NewNop()
|
||||
}
|
||||
return &Selector{redis: rdb, logger: logger}
|
||||
}
|
||||
|
||||
// overrideKey returns the Redis key for a tenant's runtime override.
|
||||
func overrideKey(tenantID string) string {
|
||||
return overrideKeyPrefix + tenantID
|
||||
}
|
||||
|
||||
// Select returns the runtime mode registered for tenantID. The lookup
|
||||
// order is:
|
||||
//
|
||||
// 1. The Redis key "tenant_canvas_runtime:<tenantID>" if present.
|
||||
// 2. The process-wide Default() (env RAGFLOW_CANVAS_DEFAULT_RUNTIME,
|
||||
// falling back to RuntimeGo).
|
||||
//
|
||||
// A nil Redis client short-circuits to the default and never errors.
|
||||
func (s *Selector) Select(ctx context.Context, tenantID string) (RuntimeMode, error) {
|
||||
if s == nil || s.redis == nil {
|
||||
return Default(), nil
|
||||
}
|
||||
raw, err := s.redis.Get(ctx, overrideKey(tenantID)).Result()
|
||||
if err == redis.Nil {
|
||||
return Default(), nil
|
||||
}
|
||||
if err != nil {
|
||||
s.logger.Warn("runtime selector: redis get failed, falling back to default",
|
||||
zap.String("tenant_id", tenantID), zap.Error(err))
|
||||
return Default(), err
|
||||
}
|
||||
mode := RuntimeMode(raw)
|
||||
switch mode {
|
||||
case RuntimeGo, RuntimePython, RuntimeAuto:
|
||||
return mode, nil
|
||||
default:
|
||||
s.logger.Warn("runtime selector: unrecognized value, falling back to default",
|
||||
zap.String("tenant_id", tenantID), zap.String("value", raw))
|
||||
return Default(), fmt.Errorf("unrecognized runtime mode %q for tenant %q", raw, tenantID)
|
||||
}
|
||||
}
|
||||
|
||||
// Set overrides the runtime mode for a tenant. The override has no TTL
|
||||
// (it is permanent until explicitly changed) so the operator does not have
|
||||
// to remember to re-set it after a Redis flush of short-lived keys. Used
|
||||
// by the admin runtime endpoint and tests.
|
||||
func (s *Selector) Set(ctx context.Context, tenantID string, mode RuntimeMode) error {
|
||||
if s == nil || s.redis == nil {
|
||||
return fmt.Errorf("runtime selector: no redis client configured")
|
||||
}
|
||||
switch mode {
|
||||
case RuntimeGo, RuntimePython, RuntimeAuto:
|
||||
default:
|
||||
return fmt.Errorf("runtime selector: refusing to set invalid mode %q", mode)
|
||||
}
|
||||
return s.redis.Set(ctx, overrideKey(tenantID), string(mode), 0).Err()
|
||||
}
|
||||
@@ -1,187 +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
|
||||
//
|
||||
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// newTestRedis spins up an in-process miniredis and returns a connected
|
||||
// client plus a teardown. The miniredis instance is reachable only from
|
||||
// the test goroutine that called this helper.
|
||||
func newTestRedis(t *testing.T) (*redis.Client, *miniredis.Miniredis) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() {
|
||||
_ = rdb.Close()
|
||||
})
|
||||
return rdb, mr
|
||||
}
|
||||
|
||||
// TestDefault_ReturnsGo is the default-runtime acceptance
|
||||
// assertion: with no env var override, Default() must return
|
||||
// RuntimeGo. The env var is explicitly cleared so the test is
|
||||
// hermetic.
|
||||
func TestDefault_ReturnsGo(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, "")
|
||||
ResetDefaultCache()
|
||||
if got := Default(); got != RuntimeGo {
|
||||
t.Fatalf("Default() = %q, want %q (Go default)", got, RuntimeGo)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefault_RespectsEnv exercises the env-var override path through
|
||||
// Default(). The Go default holds for unset/unknown values; explicit
|
||||
// "python" / "auto" still win.
|
||||
func TestDefault_RespectsEnv(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, string(RuntimePython))
|
||||
ResetDefaultCache()
|
||||
if got := Default(); got != RuntimePython {
|
||||
t.Fatalf("Default() with python env = %q, want %q", got, RuntimePython)
|
||||
}
|
||||
|
||||
t.Setenv(defaultEnvKey, string(RuntimeAuto))
|
||||
ResetDefaultCache()
|
||||
if got := Default(); got != RuntimeAuto {
|
||||
t.Fatalf("Default() with auto env = %q, want %q", got, RuntimeAuto)
|
||||
}
|
||||
|
||||
t.Setenv(defaultEnvKey, "bogus")
|
||||
ResetDefaultCache()
|
||||
if got := Default(); got != RuntimeGo {
|
||||
t.Fatalf("Default() with invalid env = %q, want %q (Go fallback)", got, RuntimeGo)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelector_PerTenantOverride verifies the per-tenant override
|
||||
// mechanism still works with the Go default: even when Default()
|
||||
// would return Go, a tenant explicitly Set to Python must be
|
||||
// routed there.
|
||||
func TestSelector_PerTenantOverride(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, "")
|
||||
ResetDefaultCache()
|
||||
|
||||
rdb, _ := newTestRedis(t)
|
||||
s := NewSelector(rdb, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
// Sanity: the global default is Go now.
|
||||
if got := Default(); got != RuntimeGo {
|
||||
t.Fatalf("setup: Default() = %q, want %q", got, RuntimeGo)
|
||||
}
|
||||
|
||||
// Tenant with no override -> Go.
|
||||
got, err := s.Select(ctx, "tenant_unset")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() unset tenant: %v", err)
|
||||
}
|
||||
if got != RuntimeGo {
|
||||
t.Fatalf("Select() unset tenant = %q, want %q", got, RuntimeGo)
|
||||
}
|
||||
|
||||
// Force tenant_force_python to Python even though default is Go.
|
||||
if err := s.Set(ctx, "tenant_force_python", RuntimePython); err != nil {
|
||||
t.Fatalf("Set() unexpected error: %v", err)
|
||||
}
|
||||
got, err = s.Select(ctx, "tenant_force_python")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() force-python tenant: %v", err)
|
||||
}
|
||||
if got != RuntimePython {
|
||||
t.Fatalf("Select() force-python tenant = %q, want %q (per-tenant override)", got, RuntimePython)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelector_Select_FallbackToDefault covers the "no Redis key" path:
|
||||
// Select must return the process-wide default with no error.
|
||||
func TestSelector_Select_FallbackToDefault(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, string(RuntimeGo))
|
||||
ResetDefaultCache()
|
||||
|
||||
rdb, _ := newTestRedis(t)
|
||||
s := NewSelector(rdb, nil)
|
||||
|
||||
got, err := s.Select(context.Background(), "tenant_42")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() unexpected error: %v", err)
|
||||
}
|
||||
if got != RuntimeGo {
|
||||
t.Fatalf("Select() = %q, want %q", got, RuntimeGo)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelector_SetAndSelectOverride proves that Set makes the override
|
||||
// visible to subsequent Select calls, even if the env default would have
|
||||
// been different.
|
||||
func TestSelector_SetAndSelectOverride(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, string(RuntimeGo))
|
||||
ResetDefaultCache()
|
||||
|
||||
rdb, _ := newTestRedis(t)
|
||||
s := NewSelector(rdb, nil)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := s.Set(ctx, "tenant_a", RuntimePython); err != nil {
|
||||
t.Fatalf("Set() unexpected error: %v", err)
|
||||
}
|
||||
got, err := s.Select(ctx, "tenant_a")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() unexpected error: %v", err)
|
||||
}
|
||||
if got != RuntimePython {
|
||||
t.Fatalf("Select() after Set = %q, want %q", got, RuntimePython)
|
||||
}
|
||||
|
||||
// Different tenant still falls back to the Go default.
|
||||
got, err = s.Select(ctx, "tenant_b")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() for other tenant: %v", err)
|
||||
}
|
||||
if got != RuntimeGo {
|
||||
t.Fatalf("Select() for other tenant = %q, want %q", got, RuntimeGo)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelector_SetRejectsInvalidMode makes sure a bad mode cannot poison
|
||||
// the override key.
|
||||
func TestSelector_SetRejectsInvalidMode(t *testing.T) {
|
||||
rdb, _ := newTestRedis(t)
|
||||
s := NewSelector(rdb, nil)
|
||||
if err := s.Set(context.Background(), "tenant_x", RuntimeMode("rust")); err == nil {
|
||||
t.Fatal("Set() with invalid mode should error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelector_NilRedis_NoError covers the "Redis not initialised" path:
|
||||
// Select returns the default and never errors; Set errors so the caller
|
||||
// knows the override wasn't applied.
|
||||
func TestSelector_NilRedis_NoError(t *testing.T) {
|
||||
t.Setenv(defaultEnvKey, string(RuntimeGo))
|
||||
ResetDefaultCache()
|
||||
|
||||
s := NewSelector(nil, nil)
|
||||
got, err := s.Select(context.Background(), "tenant_1")
|
||||
if err != nil {
|
||||
t.Fatalf("Select() with nil redis errored: %v", err)
|
||||
}
|
||||
if got != RuntimeGo {
|
||||
t.Fatalf("Select() with nil redis = %q, want default %q", got, RuntimeGo)
|
||||
}
|
||||
|
||||
if err := s.Set(context.Background(), "tenant_1", RuntimePython); err == nil {
|
||||
t.Fatal("Set() with nil redis should error")
|
||||
}
|
||||
}
|
||||
@@ -304,6 +304,35 @@ func (s *CanvasState) RecordOutput(cpnID, bucket string, payload any) {
|
||||
b[bucket] = payload
|
||||
}
|
||||
|
||||
// GetGlobal returns a value from the workflow-wide Globals bag. Globals is a
|
||||
// generic, cross-component scratch space owned by CanvasState; the set of
|
||||
// keys an ingestion pipeline elects to store there is ingestion-specific and
|
||||
// therefore lives in the ingestion component package, not here.
|
||||
func (s *CanvasState) GetGlobal(key string) (any, bool) {
|
||||
if s == nil {
|
||||
return nil, false
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
v, ok := s.Globals[key]
|
||||
return v, ok
|
||||
}
|
||||
|
||||
// SetGlobal writes a value into the workflow-wide Globals bag. It is the
|
||||
// single, lock-safe mutation point for Globals so callers never touch the map
|
||||
// field directly.
|
||||
func (s *CanvasState) SetGlobal(key string, val any) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.Globals == nil {
|
||||
s.Globals = make(map[string]any)
|
||||
}
|
||||
s.Globals[key] = val
|
||||
}
|
||||
|
||||
// GetRetrievalChunks returns a snapshot of the chunks recorded in
|
||||
// state.Retrieval["chunks"]. The Retrieval map is the canvas-level
|
||||
// aggregate that the Retrieval tool populates during the ReAct loop;
|
||||
|
||||
@@ -1,107 +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 handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/common"
|
||||
)
|
||||
|
||||
// AdminRuntimeHandler exposes the per-tenant canvas-runtime override API
|
||||
// used by the Phase 6 canary operators. It is intentionally small — the
|
||||
// selector is the only collaborator it needs.
|
||||
type AdminRuntimeHandler struct {
|
||||
selector *runtime.Selector
|
||||
}
|
||||
|
||||
// NewAdminRuntimeHandler constructs an AdminRuntimeHandler backed by the
|
||||
// supplied Selector. A nil selector is treated as a misconfiguration and
|
||||
// the handler refuses every request with HTTP 500.
|
||||
func NewAdminRuntimeHandler(selector *runtime.Selector) *AdminRuntimeHandler {
|
||||
return &AdminRuntimeHandler{selector: selector}
|
||||
}
|
||||
|
||||
// setRuntimeRequest is the wire shape for POST
|
||||
// /api/v1/admin/canvas-runtime/:tenant_id. The mode is required; empty or
|
||||
// unknown values yield 400.
|
||||
type setRuntimeRequest struct {
|
||||
Runtime string `json:"runtime"`
|
||||
}
|
||||
|
||||
// setRuntimeResponse is what the operator sees in the 200 body.
|
||||
type setRuntimeResponse struct {
|
||||
Code common.ErrorCode `json:"code"`
|
||||
TenantID string `json:"tenant_id"`
|
||||
Runtime string `json:"runtime"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// ErrSelectorNotConfigured is returned when the handler was constructed
|
||||
// without a backing Selector. It maps to HTTP 500 in the response path.
|
||||
var ErrSelectorNotConfigured = errors.New("admin runtime: selector not configured")
|
||||
|
||||
// SetTenantRuntime implements POST /api/v1/admin/canvas-runtime/:tenant_id.
|
||||
//
|
||||
// Auth gap: this handler accepts any authenticated request. The dedicated
|
||||
// admin-role middleware is a separate workstream; the Phase 6 PR documents
|
||||
// the gap here so the staging canary operator flips tenants only via a
|
||||
// trusted network. Production rollout MUST wire admin auth before opening
|
||||
// this endpoint publicly.
|
||||
func (h *AdminRuntimeHandler) SetTenantRuntime(c *gin.Context) {
|
||||
if h.selector == nil {
|
||||
common.ResponseWithCodeData(c, common.CodeExceptionError, nil, ErrSelectorNotConfigured.Error())
|
||||
return
|
||||
}
|
||||
|
||||
tenantID := c.Param("tenant_id")
|
||||
if tenantID == "" {
|
||||
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "tenant_id is required")
|
||||
return
|
||||
}
|
||||
|
||||
var req setRuntimeRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "Invalid request body: "+err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
mode := runtime.RuntimeMode(req.Runtime)
|
||||
switch mode {
|
||||
case runtime.RuntimeGo, runtime.RuntimePython, runtime.RuntimeAuto:
|
||||
// allowed
|
||||
default:
|
||||
common.ResponseWithCodeData(c, common.CodeArgumentError, nil, "runtime must be one of: go, python, auto")
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.selector.Set(c.Request.Context(), tenantID, mode); err != nil {
|
||||
common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, setRuntimeResponse{
|
||||
Code: common.CodeSuccess,
|
||||
TenantID: tenantID,
|
||||
Runtime: string(mode),
|
||||
Message: "ok",
|
||||
})
|
||||
}
|
||||
@@ -1,135 +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
|
||||
//
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
// newAdminRuntimeTestRig wires a Selector backed by miniredis and returns
|
||||
// a fully-mounted gin engine with the route registered, so tests can issue
|
||||
// real HTTP requests against it.
|
||||
func newAdminRuntimeTestRig(t *testing.T) (*gin.Engine, *runtime.Selector, *miniredis.Miniredis) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
|
||||
selector := runtime.NewSelector(rdb, nil)
|
||||
h := NewAdminRuntimeHandler(selector)
|
||||
|
||||
eng := gin.New()
|
||||
g := eng.Group("/api/v1/admin")
|
||||
g.POST("/canvas-runtime/:tenant_id", h.SetTenantRuntime)
|
||||
return eng, selector, mr
|
||||
}
|
||||
|
||||
func TestAdminRuntime_SetGo(t *testing.T) {
|
||||
eng, selector, _ := newAdminRuntimeTestRig(t)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"runtime": "go"})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/canvas-runtime/tenant_123", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
var resp setRuntimeResponse
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if resp.Code != 0 || resp.TenantID != "tenant_123" || resp.Runtime != "go" {
|
||||
t.Errorf("unexpected response: %+v", resp)
|
||||
}
|
||||
|
||||
// Round-trip: the selector should now report the override.
|
||||
mode, err := selector.Select(req.Context(), "tenant_123")
|
||||
if err != nil {
|
||||
t.Fatalf("Select(): %v", err)
|
||||
}
|
||||
if mode != runtime.RuntimeGo {
|
||||
t.Errorf("Select() after SetGo = %q, want %q", mode, runtime.RuntimeGo)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminRuntime_SetPython(t *testing.T) {
|
||||
eng, selector, _ := newAdminRuntimeTestRig(t)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"runtime": "python"})
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/canvas-runtime/tenant_xyz", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
mode, err := selector.Select(req.Context(), "tenant_xyz")
|
||||
if err != nil {
|
||||
t.Fatalf("Select(): %v", err)
|
||||
}
|
||||
if mode != runtime.RuntimePython {
|
||||
t.Errorf("Select() after SetPython = %q, want %q", mode, runtime.RuntimePython)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminRuntime_BadRequest(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{"unknown_mode", `{"runtime":"rust"}`},
|
||||
{"empty_mode", `{"runtime":""}`},
|
||||
{"malformed_json", `{"runtime":`},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
eng, _, _ := newAdminRuntimeTestRig(t)
|
||||
req := httptest.NewRequest(http.MethodPost,
|
||||
"/api/v1/admin/canvas-runtime/tenant_1",
|
||||
bytes.NewReader([]byte(tc.body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 envelope", w.Code)
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &env); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
// 101 == CodeArgumentError, the only acceptable error for bad input.
|
||||
if code, _ := env["code"].(float64); code != 101 {
|
||||
t.Errorf("code = %v, want 101 (CodeArgumentError); body=%s", env["code"], w.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -207,9 +207,42 @@ func emptyOutputs() map[string]any {
|
||||
|
||||
func emptyChunkDocs() []schema.ChunkDoc { return []schema.ChunkDoc{} }
|
||||
|
||||
// chunkOutputs builds the canonical chunker output (output_format="chunks" +
|
||||
// chunks). The Go runtime passes only this explicit output to the next node,
|
||||
// so the run-level metadata that downstream components still need (e.g.
|
||||
// `name` for Tokenizer title embedding, or tenant_id/kb_id for embedding
|
||||
// model resolution) is NOT re-emitted here — it lives in the workflow-wide
|
||||
// CanvasState.Globals bag (seeded at pipeline start, published by the File
|
||||
// component) and read directly from ctx. See runtime.CanvasState.Globals.
|
||||
func chunkOutputs(chunks []schema.ChunkDoc) map[string]any {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": schema.ChunkDocsToMaps(chunks),
|
||||
}
|
||||
}
|
||||
|
||||
// withName returns a shallow copy of inputs with name set, so a component can
|
||||
// guarantee `name` is present on the map it forwards to a decode step without
|
||||
// mutating the caller's snapshot.
|
||||
func withName(inputs map[string]any, name string) map[string]any {
|
||||
cp := make(map[string]any, len(inputs)+1)
|
||||
for k, v := range inputs {
|
||||
cp[k] = v
|
||||
}
|
||||
cp["name"] = name
|
||||
return cp
|
||||
}
|
||||
|
||||
// cloneInputs returns a shallow copy of m with room for one extra key. Used to
|
||||
// inject the Globals-resolved `name` into the decode input without mutating
|
||||
// the caller's input snapshot.
|
||||
func cloneInputs(m map[string]any) map[string]any {
|
||||
if m == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
cp := make(map[string]any, len(m)+1)
|
||||
for k, v := range m {
|
||||
cp[k] = v
|
||||
}
|
||||
return cp
|
||||
}
|
||||
|
||||
@@ -39,6 +39,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
@@ -127,10 +128,11 @@ func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam)
|
||||
if len(chunks) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return map[string]any{
|
||||
out := map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": chunks,
|
||||
}, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// groupRecords mirrors `GroupTitleChunker.build_chunks`: merges
|
||||
@@ -312,16 +314,20 @@ func (c *GroupTitleChunkerComponent) Outputs() map[string]string {
|
||||
func (c *GroupTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameGroupTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
inputs = map[string]any{}
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
// `name` is read from the workflow-wide Globals bag (seeded at
|
||||
// pipeline start, published by the File component), not from the
|
||||
// upstream output map.
|
||||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||||
if name == "" {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": "GroupTitleChunker: missing required upstream field \"name\"",
|
||||
}, nil
|
||||
}
|
||||
return invokeGroup(ctx, inputs, &c.param)
|
||||
return invokeGroup(ctx, withName(inputs, name), &c.param)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
)
|
||||
|
||||
const ComponentNameHierarchyTitleChunker = "HierarchyTitleChunker"
|
||||
@@ -265,10 +266,11 @@ func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerPa
|
||||
if len(out2) == 0 {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
return map[string]any{
|
||||
out := map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": out2,
|
||||
}, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// applyRootAsHeadingMaps mirrors the root_chunk_as_heading branch
|
||||
@@ -320,16 +322,20 @@ func (c *HierarchyTitleChunkerComponent) Outputs() map[string]string {
|
||||
func (c *HierarchyTitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameHierarchyTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
inputs = map[string]any{}
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
// `name` is read from the workflow-wide Globals bag (seeded at
|
||||
// pipeline start, published by the File component), not from the
|
||||
// upstream output map.
|
||||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||||
if name == "" {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
"_ERROR": "HierarchyTitleChunker: missing required upstream field \"name\"",
|
||||
}, nil
|
||||
}
|
||||
return invokeHierarchy(ctx, inputs, &c.param)
|
||||
return invokeHierarchy(ctx, withName(inputs, name), &c.param)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -69,5 +69,8 @@ var ChunkerInputs = map[string]string{
|
||||
var ChunkerOutputs = map[string]string{
|
||||
"output_format": "Always \"chunks\" on success.",
|
||||
"chunks": "list[object]: per-chunk map (text + optional meta keys).",
|
||||
"name": "Source document name, carried forward from upstream (pass-through) when present — Tokenizer consumes it for title embedding.",
|
||||
"tenant_id": "Carried forward from upstream (pass-through) when present — Tokenizer consumes it to resolve the embedding model.",
|
||||
"kb_id": "Carried forward from upstream (pass-through) when present — Tokenizer consumes it to resolve the embedding model.",
|
||||
"_ERROR": "Set only on validation failure.",
|
||||
}
|
||||
|
||||
@@ -54,6 +54,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
)
|
||||
|
||||
@@ -359,9 +360,13 @@ func (c *TitleChunkerComponent) Outputs() map[string]string { return ChunkerOutp
|
||||
func (c *TitleChunkerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
return runtime.TrackElapsed(ComponentNameTitleChunker, func() (map[string]any, error) {
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
inputs = map[string]any{}
|
||||
}
|
||||
if _, ok := inputs["name"].(string); !ok {
|
||||
// `name` is read from the workflow-wide Globals bag (seeded at
|
||||
// pipeline start, published by the File component), not from the
|
||||
// upstream output map.
|
||||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||||
if name == "" {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{},
|
||||
|
||||
@@ -66,6 +66,7 @@ import (
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
deepdoctype "ragflow/internal/deepdoc/parser/type"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/parser/chunk"
|
||||
)
|
||||
@@ -163,7 +164,17 @@ func (c *TokenChunkerComponent) invoke(ctx context.Context, inputs map[string]an
|
||||
if inputs == nil {
|
||||
return emptyOutputs(), nil
|
||||
}
|
||||
upstream, err := decodeChunkerFromUpstream(inputs)
|
||||
// `name` lives in the workflow-wide Globals bag (seeded at pipeline
|
||||
// start, published by the File component), not in the upstream output
|
||||
// map. decodeChunkerFromUpstream validates it, so carry the resolved
|
||||
// name into the decode input.
|
||||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||||
decInputs := inputs
|
||||
if name != "" {
|
||||
decInputs = cloneInputs(inputs)
|
||||
decInputs["name"] = name
|
||||
}
|
||||
upstream, err := decodeChunkerFromUpstream(decInputs)
|
||||
if err != nil {
|
||||
return map[string]any{
|
||||
"output_format": "chunks",
|
||||
|
||||
@@ -508,6 +508,9 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, inputs map[string]any)
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("extractor: %w", err)
|
||||
}
|
||||
// Run-level metadata (name, tenant_id, kb_id, ...) is read by
|
||||
// downstream components from the workflow-wide CanvasState.Globals
|
||||
// bag (seeded at pipeline start), so it is not re-emitted here.
|
||||
return map[string]any{
|
||||
"chunks": in.chunks,
|
||||
"output_format": "chunks",
|
||||
|
||||
@@ -43,6 +43,7 @@ import (
|
||||
"fmt"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/storage"
|
||||
)
|
||||
@@ -131,7 +132,6 @@ func (c *FileComponent) Parallelism() int { return 1 }
|
||||
// 2. doc_id is empty — pull the first file descriptor out of
|
||||
// `file` and use its `name`/`id` directly.
|
||||
func (c *FileComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
_ = ctx
|
||||
// Parse the wire input through the schema type so the
|
||||
// validation errors match the package convention.
|
||||
in, err := parseFileInputs(inputs)
|
||||
@@ -152,6 +152,12 @@ func (c *FileComponent) Invoke(ctx context.Context, inputs map[string]any) (map[
|
||||
if in.fileDesc != nil {
|
||||
out["file"] = in.fileDesc
|
||||
}
|
||||
// Publish the resolved run-level metadata into the workflow-wide
|
||||
// CanvasState.Globals bag so downstream components (Tokenizer,
|
||||
// Chunker, ...) read it from ctx instead of relying on this output
|
||||
// re-emitting it. The Go runtime forwards only this explicit output
|
||||
// to the next node, so shared fields must live in Globals.
|
||||
globals.PublishGlobals(ctx, out)
|
||||
return runtime.TrackElapsed("File", func() (map[string]any, error) {
|
||||
return out, nil
|
||||
})
|
||||
|
||||
114
internal/ingestion/component/globals/globals.go
Normal file
114
internal/ingestion/component/globals/globals.go
Normal file
@@ -0,0 +1,114 @@
|
||||
//
|
||||
// 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 globals owns the ingestion-specific run-level metadata contract.
|
||||
//
|
||||
// The generic cross-component scratch space is CanvasState.Globals (in the
|
||||
// agent runtime). Which keys an ingestion pipeline elects to store there, and
|
||||
// how they are seeded / read, is ingestion-specific — so it lives here rather
|
||||
// than in the generic canvas runtime, and in a leaf package that neither
|
||||
// imports the component package nor the pipeline package (so it cannot
|
||||
// participate in an import cycle).
|
||||
package globals
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
)
|
||||
|
||||
// canvasStateFromContext resolves the per-run CanvasState attached by the
|
||||
// pipeline (canvas.WithState). Returns nil when no state is present (e.g.
|
||||
// headless unit tests that don't attach a CanvasState).
|
||||
func canvasStateFromContext(ctx context.Context) *runtime.CanvasState {
|
||||
st, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
// GlobalMetadataKeys enumerates the run-level metadata fields that every
|
||||
// ingestion component may rely on and that the workflow carries for the whole
|
||||
// run, instead of threading through each component's output map.
|
||||
//
|
||||
// The Go ingestion runtime wires components through eino: a component's output
|
||||
// map is the sole input to the next node, and — unlike the Python
|
||||
// ProcessBase.invoke (rag/flow/base.py:42-44) — it does NOT auto-merge every
|
||||
// input kwarg into the output. A narrowing component (File, Parser, Chunker,
|
||||
// Tokenizer, Extractor, ...) would otherwise drop fields the next node still
|
||||
// depends on (e.g. TokenChunker drops `name`, which Tokenizer consumes for
|
||||
// title embedding). Storing the shared fields in CanvasState.Globals restores
|
||||
// the Python behaviour without mutating every component output.
|
||||
// The embedding-model id is intentionally NOT a global: it is a
|
||||
// Tokenizer-scoped setup (params["setups"]["embedding_model"]). Keeping it out
|
||||
// of the shared bag prevents another component (e.g. one expecting a chat
|
||||
// model) from misreading a generic "model_id" global as its own.
|
||||
var GlobalMetadataKeys = []string{
|
||||
"name",
|
||||
"doc_id",
|
||||
"bucket",
|
||||
"path",
|
||||
"file",
|
||||
"tenant_id",
|
||||
"kb_id",
|
||||
}
|
||||
|
||||
// SeedIngestionGlobals copies the whitelisted run-level metadata from `in`
|
||||
// into the CanvasState.Globals bag (last writer wins per key). It is the
|
||||
// single entry point for populating the shared workflow metadata: call it
|
||||
// once at run start (from the pipeline run inputs) and again from components
|
||||
// that derive a field mid-run (e.g. the File component publishing `name`).
|
||||
func SeedIngestionGlobals(ctx context.Context, in map[string]any) {
|
||||
if in == nil {
|
||||
return
|
||||
}
|
||||
if st := canvasStateFromContext(ctx); st != nil {
|
||||
for _, k := range GlobalMetadataKeys {
|
||||
if v, ok := in[k]; ok {
|
||||
st.SetGlobal(k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// PublishGlobals copies the resolved run-level metadata from a component's
|
||||
// output into the workflow-wide CanvasState.Globals bag so downstream
|
||||
// components can read it from ctx. No-op when state is absent.
|
||||
func PublishGlobals(ctx context.Context, out map[string]any) {
|
||||
if out == nil {
|
||||
return
|
||||
}
|
||||
SeedIngestionGlobals(ctx, out)
|
||||
}
|
||||
|
||||
// GlobalOrInput resolves a run-level field from CanvasState.Globals first,
|
||||
// then from the component's own input map, then def. Globals is the canonical
|
||||
// home for shared run metadata; the input fallback keeps headless tests
|
||||
// (which attach no CanvasState) working.
|
||||
func GlobalOrInput(ctx context.Context, inputs map[string]any, key, def string) string {
|
||||
if st := canvasStateFromContext(ctx); st != nil {
|
||||
if v, ok := st.GetGlobal(key); ok {
|
||||
if s, ok := v.(string); ok && s != "" {
|
||||
return s
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := inputs[key].(string); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
return def
|
||||
}
|
||||
@@ -82,6 +82,7 @@ import (
|
||||
"unicode/utf8"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
@@ -337,6 +338,12 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma
|
||||
if path, _ := getString(inputs, "path"); path != "" {
|
||||
out["path"] = path
|
||||
}
|
||||
// Publish the resolved run-level metadata into the workflow-wide
|
||||
// CanvasState.Globals bag so downstream components read it from ctx
|
||||
// instead of relying on this output re-emitting it. The Go runtime
|
||||
// forwards only this explicit output to the next node, so shared
|
||||
// fields must live in Globals.
|
||||
globals.PublishGlobals(ctx, out)
|
||||
// Progress (_created_time / _elapsed_time stamping, start/done
|
||||
// callbacks) is owned by the canvas framework (realComponentBody),
|
||||
// not by this component, so we return the work result directly.
|
||||
|
||||
@@ -136,6 +136,15 @@ func (c *ChunkerFromUpstream) Validate() error {
|
||||
//
|
||||
// self.set_output("output_format", "chunks")
|
||||
// self.set_output("chunks", chunks)
|
||||
//
|
||||
// Unlike the Python runtime (which auto-merges every input kwarg into the
|
||||
// output via ProcessBase.invoke), the Go runtime only forwards the explicit
|
||||
// component output to the next node. The run-level metadata that downstream
|
||||
// consumers still need (e.g. Tokenizer reads `name` for title embedding, and
|
||||
// tenant_id/kb_id for embedding-model resolution) therefore lives in the
|
||||
// workflow-wide CanvasState.Globals bag — seeded at pipeline start and
|
||||
// published by the File component — and is read directly from ctx, not
|
||||
// re-emitted by each chunker. The fields below are the chunker-owned outputs.
|
||||
type ChunkerOutputs struct {
|
||||
// OutputFormat is always "chunks" on success.
|
||||
OutputFormat PayloadFormat `json:"output_format,omitempty"`
|
||||
|
||||
@@ -51,18 +51,19 @@
|
||||
// `LLMBundle(tenant_id, embd_id).encode([...])` from
|
||||
// `rag/flow/tokenizer/tokenizer.py:54-66`; the Go port goes
|
||||
// through `service.ModelProviderService.GetEmbeddingModel`
|
||||
// (callers inject the model bundle, see `EncodeFunc` below).
|
||||
// (callers inject the resolver, see `DefaultEmbedderResolver`).
|
||||
// The component does NOT directly construct a model driver —
|
||||
// the resolution path depends on tenant/DAO context that lives
|
||||
// in `internal/service`, and importing `internal/service` from
|
||||
// `internal/ingestion/component` would invert the dependency
|
||||
// direction (plan §3 import graph: ingestion → agent/runtime
|
||||
// only). The injection point is `EncodeFunc` (package-level
|
||||
// var); production wires it in `main()` (or an analogous
|
||||
// bootstrap step) and tests inject a stub. When `EncodeFunc` is
|
||||
// nil the component short-circuits the embedding branch with
|
||||
// a clear error — the same fail-loud contract the Python side
|
||||
// enforces via `LLMBundle` constructor.
|
||||
// only). The injection point is `DefaultEmbedderResolver`
|
||||
// (package-level var); the ingestion task package wires it in
|
||||
// its init() and tests inject a stub via the test-only
|
||||
// NewTokenizerComponentWithResolver. When no resolver is
|
||||
// available the component short-circuits the embedding branch
|
||||
// with a clear error — the same fail-loud contract the Python
|
||||
// side enforces via `LLMBundle` constructor.
|
||||
//
|
||||
// - BATCHED EMBEDDING (plan §AD-5a): matched. The Python path
|
||||
// chunks calls by `settings.EMBEDDING_BATCH_SIZE` (default 16)
|
||||
@@ -99,6 +100,7 @@ import (
|
||||
"time"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
"ragflow/internal/ingestion/component/schema"
|
||||
"ragflow/internal/tokenizer"
|
||||
)
|
||||
@@ -140,7 +142,18 @@ type Embedder interface {
|
||||
}
|
||||
|
||||
// EmbedderResolver resolves the embedder for one tokenizer invocation.
|
||||
type EmbedderResolver func(tenantID, kbID, modelID string) (Embedder, error)
|
||||
// embeddingModel is the Tokenizer-scoped embedding-model identifier (from the
|
||||
// component's setups); an empty value tells the resolver to fall back to the
|
||||
// dataset's configured model.
|
||||
type EmbedderResolver func(tenantID, kbID, embeddingModel string) (Embedder, error)
|
||||
|
||||
// DefaultEmbedderResolver is the production embedder resolver. It is nil in
|
||||
// this leaf package — which must not import internal/service (see the
|
||||
// EMBEDDING MODEL RESOLUTION note above) — and is injected by the composition
|
||||
// root: the ingestion task package wires a resolver backed by the model
|
||||
// provider in its init(). NewTokenizerComponent falls back to this resolver
|
||||
// when no explicit (test-only) resolver is supplied.
|
||||
var DefaultEmbedderResolver EmbedderResolver
|
||||
|
||||
// TokenizerComponent computes token counts and (optionally) embedding
|
||||
// vectors for an upstream chunk list. Mirrors python
|
||||
@@ -149,8 +162,8 @@ type EmbedderResolver func(tenantID, kbID, modelID string) (Embedder, error)
|
||||
// Inputs:
|
||||
//
|
||||
// tenant_id (string, optional) — used to resolve the embedding model
|
||||
// model_id (string, optional) — explicit override; falls back to
|
||||
// Param.EmbeddingID (future)
|
||||
// kb_id (string, optional) — dataset whose embd_id is used when the
|
||||
// setups embedding_model is unset
|
||||
// output_format (string) — one of json/markdown/text/html/chunks
|
||||
// chunks (list[map]) — chunk list when output_format == "chunks"
|
||||
// json (list[map]) — structured parser payload when output_format == "json" or unset
|
||||
@@ -166,19 +179,31 @@ type EmbedderResolver func(tenantID, kbID, modelID string) (Embedder, error)
|
||||
// output_format — always "chunks" (matches python set_output)
|
||||
// _created_time / _elapsed_time — TrackElapsed bookkeeping
|
||||
type TokenizerComponent struct {
|
||||
param schema.TokenizerParam
|
||||
resolver EmbedderResolver
|
||||
param schema.TokenizerParam
|
||||
resolver EmbedderResolver
|
||||
embeddingModel string
|
||||
}
|
||||
|
||||
// NewTokenizerComponent constructs a TokenizerComponent from DSL
|
||||
// NewTokenizerComponent constructs a production TokenizerComponent from DSL
|
||||
// params. Mirrors python `TokenizerParam` defaults (search_method =
|
||||
// ["full_text","embedding"], filename_embd_weight=0.1, fields=["text"]).
|
||||
// ["full_text","embedding"], filename_embd_weight=0.1, fields=["text"]). The
|
||||
// embedding branch resolves its embedder via the injected
|
||||
// DefaultEmbedderResolver (wired by the ingestion task package).
|
||||
func NewTokenizerComponent(params map[string]any) (runtime.Component, error) {
|
||||
return NewTokenizerComponentWithResolver(params, nil)
|
||||
return newTokenizerComponent(params, nil)
|
||||
}
|
||||
|
||||
// NewTokenizerComponentWithResolver is TEST-ONLY. It injects an explicit
|
||||
// embedder resolver so unit/integration tests can stub the embedding backend
|
||||
// without touching the model provider. Production code MUST use
|
||||
// NewTokenizerComponent and rely on DefaultEmbedderResolver instead.
|
||||
func NewTokenizerComponentWithResolver(params map[string]any, resolver EmbedderResolver) (runtime.Component, error) {
|
||||
return newTokenizerComponent(params, resolver)
|
||||
}
|
||||
|
||||
func newTokenizerComponent(params map[string]any, resolver EmbedderResolver) (runtime.Component, error) {
|
||||
p := schema.TokenizerParam{}.Defaults()
|
||||
embeddingModel := ""
|
||||
if params != nil {
|
||||
if v, ok := params["search_method"]; ok {
|
||||
// Replace (not append) so a caller-supplied
|
||||
@@ -219,19 +244,35 @@ func NewTokenizerComponentWithResolver(params map[string]any, resolver EmbedderR
|
||||
p.Fields = append(p.Fields, t...)
|
||||
}
|
||||
}
|
||||
embeddingModel = embeddingModelFromSetups(params)
|
||||
}
|
||||
if err := p.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("Tokenizer: param check: %w", err)
|
||||
}
|
||||
return &TokenizerComponent{param: p, resolver: resolver}, nil
|
||||
return &TokenizerComponent{param: p, resolver: resolver, embeddingModel: embeddingModel}, nil
|
||||
}
|
||||
|
||||
// embeddingModelFromSetups extracts the embedding-model identifier from the
|
||||
// component's setups map (params["setups"]["embedding_model"]). The embedding
|
||||
// model id is a Tokenizer-scoped setup rather than a run-level global so it is
|
||||
// never mistaken for, e.g., a chat model id shared across components. Empty
|
||||
// when unset — the resolver then falls back to the dataset's configured model.
|
||||
func embeddingModelFromSetups(params map[string]any) string {
|
||||
setups, ok := params["setups"].(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
if v, ok := setups["embedding_model"].(string); ok {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// Inputs returns the parameter metadata.
|
||||
func (c *TokenizerComponent) Inputs() map[string]string {
|
||||
return map[string]string{
|
||||
"tenant_id": "Tenant identifier used to resolve the embedding model (mirrors python self._canvas._tenant_id).",
|
||||
"kb_id": "Optional knowledgebase identifier used to resolve the bound embedding model when model_id is unset.",
|
||||
"model_id": "Optional explicit embedding-model override.",
|
||||
"kb_id": "Optional knowledgebase identifier used to resolve the bound embedding model when the setups embedding_model is unset.",
|
||||
"output_format": "Upstream payload discriminator: json / markdown / text / html / chunks.",
|
||||
"chunks": "List of chunk maps when output_format == \"chunks\".",
|
||||
"json": "Structured parser payload when output_format == \"json\" or unset.",
|
||||
@@ -273,15 +314,30 @@ func (c *TokenizerComponent) Parallelism() int { return 1 }
|
||||
// continue`), but the chunk still carries tokenized fields if
|
||||
// `full_text` is in `search_method`.
|
||||
func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
tenantID := getStringOr(inputs, "tenant_id", "")
|
||||
kbID := getStringOr(inputs, "kb_id", "")
|
||||
modelID := getStringOr(inputs, "model_id", "")
|
||||
upstream, err := decodeTokenizerFromUpstream(inputs)
|
||||
// Run-level metadata lives in the workflow-wide CanvasState.Globals
|
||||
// bag (seeded at pipeline start, published by the File component),
|
||||
// not in the upstream output map — see GlobalOrInput.
|
||||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||||
tenantID := globals.GlobalOrInput(ctx, inputs, "tenant_id", "")
|
||||
kbID := globals.GlobalOrInput(ctx, inputs, "kb_id", "")
|
||||
// The embedding-model id is a Tokenizer-scoped setup (params["setups"]),
|
||||
// resolved at construction, not a run-level global — see
|
||||
// embeddingModelFromSetups.
|
||||
embeddingModel := c.embeddingModel
|
||||
|
||||
// decodeTokenizerFromUpstream validates `name`; carry the resolved
|
||||
// name into the decode input so both a Globals-backed run and a
|
||||
// headless run (no Globals attached) satisfy it.
|
||||
decInputs := inputs
|
||||
if name != "" {
|
||||
decInputs = cloneInputs(inputs)
|
||||
decInputs["name"] = name
|
||||
}
|
||||
upstream, err := decodeTokenizerFromUpstream(decInputs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chunks := chunksFromTokenizerUpstream(upstream)
|
||||
name := upstream.Name
|
||||
titleStem := titleExtRE.ReplaceAllString(name, "")
|
||||
|
||||
return runtime.TrackElapsed("Tokenizer", func() (map[string]any, error) {
|
||||
@@ -299,7 +355,7 @@ func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any)
|
||||
}
|
||||
|
||||
if contains(c.param.SearchMethod, "embedding") {
|
||||
chunks, tokenCount, err := c.embedChunks(ctx, tenantID, kbID, modelID, name, chunks)
|
||||
chunks, tokenCount, err := c.embedChunks(ctx, tenantID, kbID, embeddingModel, name, chunks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -314,14 +370,20 @@ func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any)
|
||||
})
|
||||
}
|
||||
|
||||
func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, modelID, name string, chunks []schema.ChunkDoc) ([]schema.ChunkDoc, int, error) {
|
||||
func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, embeddingModel, name string, chunks []schema.ChunkDoc) ([]schema.ChunkDoc, int, error) {
|
||||
if len(chunks) == 0 {
|
||||
return chunks, 0, nil
|
||||
}
|
||||
if c.resolver == nil {
|
||||
return nil, 0, fmt.Errorf("Tokenizer: embedding requested but resolver is unset")
|
||||
// An explicit (test-only) resolver wins; production wiring leaves it nil
|
||||
// and falls back to the injected DefaultEmbedderResolver.
|
||||
resolver := c.resolver
|
||||
if resolver == nil {
|
||||
resolver = DefaultEmbedderResolver
|
||||
}
|
||||
embedder, err := c.resolver(tenantID, kbID, modelID)
|
||||
if resolver == nil {
|
||||
return nil, 0, fmt.Errorf("Tokenizer: embedding requested but no embedder resolver configured")
|
||||
}
|
||||
embedder, err := resolver(tenantID, kbID, embeddingModel)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("Tokenizer: resolve embedder: %w", err)
|
||||
}
|
||||
@@ -703,28 +765,24 @@ func hasEmbeddingVector(ck schema.ChunkDoc) bool {
|
||||
}
|
||||
|
||||
func getStringOr(m map[string]any, key, def string) string {
|
||||
if v, ok := getStringLocal(m, key); ok && v != "" {
|
||||
if v, ok := m[key].(string); ok && v != "" {
|
||||
return v
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
// getStringLocal mirrors file.go's getString; we keep a local copy
|
||||
// so the tokenizer package does not depend on the file package's
|
||||
// helper signature. Reads either a string or a byte slice (JSON
|
||||
// decoding yields string for string fields by default).
|
||||
func getStringLocal(m map[string]any, key string) (string, bool) {
|
||||
v, ok := m[key]
|
||||
if !ok || v == nil {
|
||||
return "", false
|
||||
// cloneInputs returns a shallow copy of m with room for one extra key.
|
||||
// Used to inject the Globals-resolved `name` into the decode input without
|
||||
// mutating the caller's input snapshot.
|
||||
func cloneInputs(m map[string]any) map[string]any {
|
||||
if m == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
switch s := v.(type) {
|
||||
case string:
|
||||
return s, true
|
||||
case []byte:
|
||||
return string(s), true
|
||||
cp := make(map[string]any, len(m)+1)
|
||||
for k, v := range m {
|
||||
cp[k] = v
|
||||
}
|
||||
return "", false
|
||||
return cp
|
||||
}
|
||||
|
||||
func contains(s []string, v string) bool {
|
||||
|
||||
@@ -499,7 +499,8 @@ func TestTokenizerComponent_Invoke_FullTextAndEmbedding(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestTokenizerComponent_Invoke_EmbedNoResolver covers the
|
||||
// "embedding requested but resolver is unset" branch — must
|
||||
// "embedding requested but no embedder resolver configured" branch
|
||||
// (explicit resolver nil and DefaultEmbedderResolver unset) — must
|
||||
// return a clear error, not panic.
|
||||
func TestTokenizerComponent_Invoke_EmbedNoResolver(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
|
||||
@@ -31,6 +31,7 @@ import (
|
||||
"ragflow/internal/dao"
|
||||
redis2 "ragflow/internal/engine/redis"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/ingestion/component/globals"
|
||||
|
||||
"github.com/cloudwego/eino/compose"
|
||||
)
|
||||
@@ -188,13 +189,23 @@ var defaultCheckpointTTL = 24 * time.Hour
|
||||
// There is no pipeline-layer partial resume entry point: execution always
|
||||
// starts from the graph entry and component-level replay decisions belong to
|
||||
// the components themselves.
|
||||
func (p *Pipeline) Run(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, setups ...map[string]any) (map[string]any, error) {
|
||||
if p == nil {
|
||||
return nil, fmt.Errorf("pipeline: Run on nil pipeline")
|
||||
}
|
||||
if p.canvas == nil {
|
||||
return nil, fmt.Errorf("pipeline: canvas is nil")
|
||||
}
|
||||
// runSetups, when non-nil, overrides components' DSL-baked
|
||||
// `params["setups"]` at compile time. It is keyed by cpnID; each
|
||||
// component is merged only with its own entry, and within that entry a
|
||||
// top-level key fully replaces the base entry for that key (see
|
||||
// canvas.mergeSetups). It is variadic so existing callers that pass
|
||||
// only (ctx, inputs) keep working.
|
||||
var runSetups map[string]any
|
||||
if len(setups) > 0 {
|
||||
runSetups = setups[0]
|
||||
}
|
||||
if runtime.DefaultFactory() == nil {
|
||||
runtime.InstallDefaultRegistryFactory()
|
||||
}
|
||||
@@ -234,6 +245,9 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any) (map[string]a
|
||||
canvas.WithInterruptAfterNonTerminalCpn(),
|
||||
)
|
||||
}
|
||||
// Run-level setups (keyed by cpnID) override the DSL-baked component
|
||||
// setups at compile time (higher priority; see canvas.WithSetupOverrides).
|
||||
compileOpts = append(compileOpts, canvas.WithSetupOverrides(runSetups))
|
||||
compiled, err := canvas.Compile(compileCtx, p.canvas, compileOpts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pipeline: Run: compile canvas: %w", err)
|
||||
@@ -263,6 +277,14 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any) (map[string]a
|
||||
|
||||
current := cloneMapOrEmpty(inputs)
|
||||
|
||||
// Seed the workflow-wide Globals bag with the run-level metadata
|
||||
// (name, tenant_id, kb_id, model_id, doc_id, ...) once, from the
|
||||
// pipeline run inputs. Downstream components read these from ctx
|
||||
// instead of relying on every node re-emitting them. The File
|
||||
// component re-publishes `name` (and storage refs) as it derives
|
||||
// them mid-run.
|
||||
globals.SeedIngestionGlobals(runCtx, current)
|
||||
|
||||
if !resumable {
|
||||
return p.runPlain(runCtx, current, compiled, tracker, runState)
|
||||
}
|
||||
|
||||
@@ -20,7 +20,6 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"ragflow/internal/agent/runtime"
|
||||
componentpkg "ragflow/internal/ingestion/component"
|
||||
"ragflow/internal/utility"
|
||||
"regexp"
|
||||
@@ -115,14 +114,55 @@ type PipelineExecutor struct {
|
||||
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
|
||||
getEmbeddingModelFunc func(tenantID, embdID string) (*models.EmbeddingModel, error)
|
||||
getKnowledgebaseByIDFunc func(kbID string) (*entity.Knowledgebase, error)
|
||||
loadDSLFunc func(ctx context.Context, dataflowID string) (string, string, error)
|
||||
runPipelineFunc func(ctx context.Context, dsl string) (map[string]any, string, error)
|
||||
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 {
|
||||
@@ -163,7 +203,6 @@ func NewDataflowService(
|
||||
if taskCtx != nil && taskCtx.ProgressFunc != nil {
|
||||
progressFn = taskCtx.ProgressFunc
|
||||
}
|
||||
modelProvider := service.NewModelProviderService()
|
||||
svc := &PipelineExecutor{
|
||||
taskCtx: taskCtx,
|
||||
dataflowID: dataflowID,
|
||||
@@ -175,9 +214,7 @@ func NewDataflowService(
|
||||
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,
|
||||
getEmbeddingModelFunc: modelProvider.GetEmbeddingModel,
|
||||
getKnowledgebaseByIDFunc: dao.NewKnowledgebaseDAO().GetByID,
|
||||
logCreateFunc: dao.NewPipelineOperationLogDAO().Create,
|
||||
}
|
||||
svc.loadDSLFunc = svc.defaultLoadDSL
|
||||
svc.runPipelineFunc = svc.defaultRunPipeline
|
||||
@@ -199,11 +236,6 @@ func (s *PipelineExecutor) WithLogCreateFunc(f func(log *entity.PipelineOperatio
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) WithGetEmbeddingModelFunc(f func(tenantID, embdID string) (*models.EmbeddingModel, error)) *PipelineExecutor {
|
||||
s.getEmbeddingModelFunc = f
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) WithDocService(d docService) *PipelineExecutor {
|
||||
s.docSvc = d
|
||||
return s
|
||||
@@ -401,10 +433,6 @@ func (s *PipelineExecutor) progress(prog float64, msg string) {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) getEmbeddingModel(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
return s.getEmbeddingModelFunc(tenantID, embdID)
|
||||
}
|
||||
|
||||
func hasVectors(chunks []map[string]any) bool {
|
||||
for _, ck := range chunks {
|
||||
for k := range ck {
|
||||
@@ -443,36 +471,6 @@ func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, dataflowID string
|
||||
return string(raw), dataflowID, nil
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) tokenizerEmbedderResolver() componentpkg.EmbedderResolver {
|
||||
return func(tenantID, kbID, modelID string) (componentpkg.Embedder, error) {
|
||||
if strings.TrimSpace(tenantID) == "" && s != nil && s.taskCtx != nil {
|
||||
tenantID = s.taskCtx.Tenant.ID
|
||||
}
|
||||
if strings.TrimSpace(kbID) == "" && s != nil && s.taskCtx != nil {
|
||||
kbID = s.taskCtx.KB.ID
|
||||
}
|
||||
model, err := s.resolveEmbeddingModel(tenantID, kbID, modelID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &embedder{model: model}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) taskScopedComponentFactory() runtime.ComponentFactory {
|
||||
resolver := s.tokenizerEmbedderResolver()
|
||||
return func(name string, params map[string]any) (runtime.Component, error) {
|
||||
if strings.EqualFold(name, componentpkg.ComponentNameTokenizer) {
|
||||
return componentpkg.NewTokenizerComponentWithResolver(params, resolver)
|
||||
}
|
||||
factory, _, _, ok := runtime.DefaultRegistry.Lookup(name)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("runtime: unknown component %q", name)
|
||||
}
|
||||
return factory(name, params)
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
@@ -487,7 +485,6 @@ func (s *PipelineExecutor) defaultRunPipeline(ctx context.Context, dsl string) (
|
||||
if err != nil {
|
||||
return nil, dsl, fmt.Errorf("compile pipeline dsl: %w", err)
|
||||
}
|
||||
pipe.WithComponentFactory(s.taskScopedComponentFactory())
|
||||
inputs := map[string]any{}
|
||||
if s.taskCtx.Doc.ID != "" {
|
||||
inputs["doc_id"] = s.taskCtx.Doc.ID
|
||||
@@ -509,29 +506,6 @@ func (s *PipelineExecutor) defaultRunPipeline(ctx context.Context, dsl string) (
|
||||
return payload, dsl, nil
|
||||
}
|
||||
|
||||
func (s *PipelineExecutor) resolveEmbeddingModel(tenantID, kbID, modelID string) (*models.EmbeddingModel, error) {
|
||||
_ = modelID
|
||||
if strings.TrimSpace(kbID) != "" {
|
||||
if s.taskCtx != nil && s.taskCtx.KB.ID == kbID && strings.TrimSpace(s.taskCtx.KB.EmbdID) != "" {
|
||||
return s.getEmbeddingModelFunc(tenantID, s.taskCtx.KB.EmbdID)
|
||||
}
|
||||
if s.getKnowledgebaseByIDFunc == nil {
|
||||
return nil, fmt.Errorf("knowledgebase resolver unavailable")
|
||||
}
|
||||
kb, err := s.getKnowledgebaseByIDFunc(kbID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if kb != nil && strings.TrimSpace(kb.EmbdID) != "" {
|
||||
return s.getEmbeddingModelFunc(tenantID, kb.EmbdID)
|
||||
}
|
||||
}
|
||||
if s.taskCtx != nil && strings.TrimSpace(s.taskCtx.KB.EmbdID) != "" {
|
||||
return s.getEmbeddingModelFunc(tenantID, s.taskCtx.KB.EmbdID)
|
||||
}
|
||||
return nil, fmt.Errorf("embedding requested but dataset has no embd_id configured")
|
||||
}
|
||||
|
||||
func extractDataflowPipelinePayload(dsl string, out map[string]any) (map[string]any, error) {
|
||||
if out == nil {
|
||||
return nil, nil
|
||||
|
||||
@@ -114,7 +114,7 @@ func TestNewDataflowService_Basic(t *testing.T) {
|
||||
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.getEmbeddingModelFunc == nil || svc.loadDSLFunc == nil || svc.runPipelineFunc == 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")
|
||||
}
|
||||
}
|
||||
@@ -414,13 +414,9 @@ func TestRunDataflow_NormalizedEmpty(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunDataflow_FullFlow(t *testing.T) {
|
||||
stub := &stubDriver{}
|
||||
var progressCalls []float64
|
||||
var progressMsgs []string
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).
|
||||
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
return makeTestEmbeddingModel(stub, 100), nil
|
||||
}).
|
||||
WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
||||
return nil, nil
|
||||
}).
|
||||
@@ -529,7 +525,6 @@ func TestExtractDataflowPipelinePayload_ErrorsOnMultipleTerminals(t *testing.T)
|
||||
}
|
||||
|
||||
func TestDataflowService_Run_MainFlowWithStubs(t *testing.T) {
|
||||
stub := &stubDriver{}
|
||||
logged := false
|
||||
inserted := false
|
||||
var progressCalls []float64
|
||||
@@ -546,9 +541,6 @@ func TestDataflowService_Run_MainFlowWithStubs(t *testing.T) {
|
||||
},
|
||||
}, dsl, nil
|
||||
}).
|
||||
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
return makeTestEmbeddingModel(stub, 100), nil
|
||||
}).
|
||||
WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) {
|
||||
inserted = true
|
||||
return nil, nil
|
||||
@@ -676,75 +668,63 @@ func makeEmbeddingModelForResolver() *models.EmbeddingModel {
|
||||
return models.NewEmbeddingModel(&stubDriver{}, strPtr("embed"), &models.APIConfig{}, 128)
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_ResolveEmbeddingModel_IgnoresModelIDAndUsesDatasetEmbedding(t *testing.T) {
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0)
|
||||
var gotTenantID, gotModelID string
|
||||
svc.getEmbeddingModelFunc = func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
gotTenantID, gotModelID = tenantID, embdID
|
||||
return makeEmbeddingModelForResolver(), nil
|
||||
}
|
||||
model, err := svc.resolveEmbeddingModel("tenant-1", "kb-1", "override-model")
|
||||
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("resolveEmbeddingModel: %v", err)
|
||||
t.Fatalf("resolver: %v", err)
|
||||
}
|
||||
if model == nil {
|
||||
t.Fatal("expected model")
|
||||
if emb == nil {
|
||||
t.Fatal("expected embedder")
|
||||
}
|
||||
if gotTenantID != "tenant-1" || gotModelID != "embd-1" {
|
||||
t.Fatalf("resolver args = (%q, %q), want (tenant-1, embd-1)", gotTenantID, gotModelID)
|
||||
if gotTenantID != "tenant-1" || gotEmbdID != "explicit-embd" {
|
||||
t.Fatalf("resolver args = (%q, %q), want (tenant-1, explicit-embd)", gotTenantID, gotEmbdID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_ResolveEmbeddingModel_KnowledgebaseModel(t *testing.T) {
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0)
|
||||
var gotModelID string
|
||||
svc.getEmbeddingModelFunc = func(_ string, embdID string) (*models.EmbeddingModel, error) {
|
||||
gotModelID = embdID
|
||||
return makeEmbeddingModelForResolver(), nil
|
||||
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)
|
||||
}
|
||||
_, err := svc.resolveEmbeddingModel("tenant-1", "kb-1", "")
|
||||
if err != nil {
|
||||
t.Fatalf("resolveEmbeddingModel: %v", err)
|
||||
}
|
||||
if gotModelID != "embd-1" {
|
||||
t.Fatalf("got model id %q, want embd-1", gotModelID)
|
||||
if gotEmbdID != "lookup-embd" {
|
||||
t.Fatalf("got embd id %q, want lookup-embd", gotEmbdID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_ResolveEmbeddingModel_KnowledgebaseLookupFallback(t *testing.T) {
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0)
|
||||
svc.taskCtx.KB.ID = "other-kb"
|
||||
var gotModelID string
|
||||
svc.getEmbeddingModelFunc = func(_ string, embdID string) (*models.EmbeddingModel, error) {
|
||||
gotModelID = embdID
|
||||
return makeEmbeddingModelForResolver(), nil
|
||||
}
|
||||
svc.getKnowledgebaseByIDFunc = func(kbID string) (*entity.Knowledgebase, error) {
|
||||
if kbID != "kb-2" {
|
||||
t.Fatalf("kb lookup id = %q, want kb-2", kbID)
|
||||
}
|
||||
return &entity.Knowledgebase{ID: "kb-2", EmbdID: "lookup-embd"}, nil
|
||||
}
|
||||
_, err := svc.resolveEmbeddingModel("tenant-1", "kb-2", "")
|
||||
if err != nil {
|
||||
t.Fatalf("resolveEmbeddingModel: %v", err)
|
||||
}
|
||||
if gotModelID != "lookup-embd" {
|
||||
t.Fatalf("got model id %q, want lookup-embd", gotModelID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_ResolveEmbeddingModel_MissingDatasetEmbeddingReturnsError(t *testing.T) {
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0)
|
||||
svc.taskCtx.KB.EmbdID = ""
|
||||
svc.getKnowledgebaseByIDFunc = func(string) (*entity.Knowledgebase, error) {
|
||||
return &entity.Knowledgebase{ID: "kb-1", EmbdID: ""}, nil
|
||||
}
|
||||
svc.getEmbeddingModelFunc = func(_ string, _ string) (*models.EmbeddingModel, error) {
|
||||
t.Fatal("knowledgebase resolver should not be called")
|
||||
return nil, nil
|
||||
}
|
||||
_, err := svc.resolveEmbeddingModel("tenant-1", "kb-1", "")
|
||||
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")
|
||||
}
|
||||
@@ -753,22 +733,22 @@ func TestPipelineExecutor_ResolveEmbeddingModel_MissingDatasetEmbeddingReturnsEr
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineExecutor_TokenizerEmbedderResolver_FallsBackToTaskContext(t *testing.T) {
|
||||
svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0)
|
||||
var gotTenantID, gotModelID string
|
||||
svc.getEmbeddingModelFunc = func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
gotTenantID, gotModelID = tenantID, embdID
|
||||
return makeEmbeddingModelForResolver(), nil
|
||||
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")
|
||||
}
|
||||
resolver := svc.tokenizerEmbedderResolver()
|
||||
emb, err := resolver("", "", "")
|
||||
if err != nil {
|
||||
t.Fatalf("resolver: %v", err)
|
||||
}
|
||||
if emb == nil {
|
||||
t.Fatal("expected embedder")
|
||||
}
|
||||
if gotTenantID != "tenant-1" || gotModelID != "embd-1" {
|
||||
t.Fatalf("resolver args = (%q, %q), want (tenant-1, embd-1)", gotTenantID, gotModelID)
|
||||
if !strings.Contains(err.Error(), "neither embedding_model nor kb_id") {
|
||||
t.Fatalf("err = %v, want neither embedding_model nor kb_id", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func testStrPtr(s string) *string { return &s }
|
||||
@@ -65,10 +64,7 @@ func newNoopDataflowService(ctx *TaskContext, dataflowID string) (*PipelineExecu
|
||||
return nil
|
||||
}).
|
||||
WithDocService(&stubDocService{}).
|
||||
WithChunkCounter(&stubChunkCounter{}).
|
||||
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
return nil, nil
|
||||
})
|
||||
WithChunkCounter(&stubChunkCounter{})
|
||||
return svc, nil
|
||||
}
|
||||
|
||||
@@ -203,10 +199,7 @@ func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) {
|
||||
return nil
|
||||
}).
|
||||
WithDocService(&stubDocService{}).
|
||||
WithChunkCounter(&stubChunkCounter{}).
|
||||
WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) {
|
||||
return nil, nil
|
||||
})
|
||||
WithChunkCounter(&stubChunkCounter{})
|
||||
return svc, nil
|
||||
})
|
||||
|
||||
|
||||
@@ -1,54 +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 router — admin_routes.go registers the Phase 6 per-tenant
|
||||
// canvas-runtime override endpoint on the existing v1 admin group. It is
|
||||
// kept separate from router.go so future admin endpoints can land here
|
||||
// without churn in the main route table.
|
||||
package router
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"ragflow/internal/handler"
|
||||
)
|
||||
|
||||
// RegisterAdminRuntimeRoutes wires the canvas-runtime override endpoint
|
||||
// onto an existing /admin RouterGroup. The caller is expected to be the
|
||||
// authorised v1 group; this function is intentionally agnostic of the
|
||||
// full path prefix so the same registration helper works for the main
|
||||
// server and any future admin sub-app.
|
||||
//
|
||||
// The single route is:
|
||||
//
|
||||
// POST /api/v1/admin/canvas-runtime/:tenant_id
|
||||
// body: {"runtime": "go" | "python" | "auto"}
|
||||
// response: 200 {"code":0,"tenant_id":...,"runtime":...,"message":"ok"}
|
||||
//
|
||||
// The handler h must be non-nil. A handler with a nil selector (e.g.
|
||||
// the server started before Redis was reachable) still serves this
|
||||
// route — SetTenantRuntime responds with HTTP 500 and
|
||||
// ErrSelectorNotConfigured. The previous version of this function
|
||||
// silently no-op'd on a nil handler, which made the route disappear
|
||||
// after a Redis outage at boot and only re-appear on the next process
|
||||
// restart. Review follow-up: keep the route hot, surface a clear error
|
||||
// to the operator.
|
||||
func RegisterAdminRuntimeRoutes(g *gin.RouterGroup, h *handler.AdminRuntimeHandler) {
|
||||
if g == nil || h == nil {
|
||||
return
|
||||
}
|
||||
g.POST("/canvas-runtime/:tenant_id", h.SetTenantRuntime)
|
||||
}
|
||||
@@ -1,103 +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
|
||||
//
|
||||
|
||||
package router
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"ragflow/internal/agent/runtime"
|
||||
"ragflow/internal/handler"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
}
|
||||
|
||||
func TestAdminRuntimeRoutes_Registered(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
|
||||
selector := runtime.NewSelector(rdb, nil)
|
||||
h := handler.NewAdminRuntimeHandler(selector)
|
||||
|
||||
eng := gin.New()
|
||||
v1 := eng.Group("/api/v1")
|
||||
admin := v1.Group("/admin")
|
||||
RegisterAdminRuntimeRoutes(admin, h)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"runtime": "go"})
|
||||
req := httptest.NewRequest(http.MethodPost,
|
||||
"/api/v1/admin/canvas-runtime/tenant_123", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
if !bytes.Contains(w.Body.Bytes(), []byte(`"runtime":"go"`)) {
|
||||
t.Errorf("response body missing runtime:go: %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminRuntimeRoutes_NilSafety(t *testing.T) {
|
||||
// A nil router group or handler must not panic; the helper is
|
||||
// documented as a no-op in that case so wiring bugs surface as
|
||||
// missing routes rather than nil-deref panics.
|
||||
RegisterAdminRuntimeRoutes(nil, nil)
|
||||
// Just reaching here without panicking is the test.
|
||||
}
|
||||
|
||||
// TestAdminRuntimeRoutes_StaysRegisteredWithNilSelector locks in the
|
||||
// review follow-up: when the server starts before Redis is reachable
|
||||
// the handler is constructed with a nil selector. The route MUST
|
||||
// still be registered and MUST return ErrSelectorNotConfigured (HTTP
|
||||
// 500), not a 404. The previous version of the wiring made the route
|
||||
// vanish in this scenario, which stranded canary operators with an
|
||||
// opaque 404 until the next process restart.
|
||||
func TestAdminRuntimeRoutes_StaysRegisteredWithNilSelector(t *testing.T) {
|
||||
h := handler.NewAdminRuntimeHandler(nil) // nil selector — Redis unavailable
|
||||
|
||||
eng := gin.New()
|
||||
v1 := eng.Group("/api/v1")
|
||||
admin := v1.Group("/admin")
|
||||
RegisterAdminRuntimeRoutes(admin, h)
|
||||
|
||||
body, _ := json.Marshal(map[string]string{"runtime": "go"})
|
||||
req := httptest.NewRequest(http.MethodPost,
|
||||
"/api/v1/admin/canvas-runtime/tenant_123", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
eng.ServeHTTP(w, req)
|
||||
|
||||
if w.Code == http.StatusNotFound {
|
||||
t.Fatalf("route returned 404 — the route must stay registered even when the selector is nil; body=%s", w.Body.String())
|
||||
}
|
||||
if w.Code != http.StatusOK {
|
||||
// 200/500 both acceptable; the contract is "not 404" so the
|
||||
// operator sees a uniform surface and can read the error in the
|
||||
// body. The handler currently returns 500 with
|
||||
// ErrSelectorNotConfigured; we assert the body contains that
|
||||
// string for a useful diagnostic.
|
||||
if !bytes.Contains(w.Body.Bytes(), []byte("selector not configured")) {
|
||||
t.Errorf("body missing 'selector not configured' diagnostic; got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -51,7 +51,6 @@ type Router struct {
|
||||
pluginHandler *handler.PluginHandler
|
||||
modelHandler *handler.ModelHandler
|
||||
fileCommitHandler *handler.FileCommitHandler
|
||||
adminRuntimeHandler *handler.AdminRuntimeHandler
|
||||
botHandler *handler.BotHandler
|
||||
componentsHandler *handler.ComponentsHandler
|
||||
}
|
||||
@@ -84,7 +83,6 @@ func NewRouter(
|
||||
pluginHandler *handler.PluginHandler,
|
||||
modelHandler *handler.ModelHandler,
|
||||
fileCommitHandler *handler.FileCommitHandler,
|
||||
adminRuntimeHandler *handler.AdminRuntimeHandler,
|
||||
openaiChatHandler *handler.OpenAIChatHandler,
|
||||
botHandler *handler.BotHandler,
|
||||
componentsHandler *handler.ComponentsHandler,
|
||||
@@ -117,7 +115,6 @@ func NewRouter(
|
||||
pluginHandler: pluginHandler,
|
||||
modelHandler: modelHandler,
|
||||
fileCommitHandler: fileCommitHandler,
|
||||
adminRuntimeHandler: adminRuntimeHandler,
|
||||
botHandler: botHandler,
|
||||
componentsHandler: componentsHandler,
|
||||
}
|
||||
@@ -570,12 +567,6 @@ func (r *Router) Setup(engine *gin.Engine) {
|
||||
v1.GET("/components", r.componentsHandler.Get)
|
||||
}
|
||||
|
||||
// Admin routes — Phase 6 per-tenant canvas runtime override.
|
||||
// RegisterAdminRuntimeRoutes lives in admin_routes.go; a nil
|
||||
// handler is tolerated and yields a no-op registration.
|
||||
admin := v1.Group("/admin")
|
||||
RegisterAdminRuntimeRoutes(admin, r.adminRuntimeHandler)
|
||||
|
||||
connectors := v1.Group("/connectors")
|
||||
{
|
||||
connectors.GET("/", r.connectorHandler.ListConnectors)
|
||||
|
||||
Reference in New Issue
Block a user