mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 21:37:33 +08:00
## Summary
Adds a side-effect-free DataFlow canvas **debug (dry-run) mode** plus a
**debug run log with a "View result" panel**, so a canvas can be
executed synchronously and inspected end-to-end (per-component progress
and parsed chunks) without persisting anything.
### Dry-run execution (inline parsed chunks)
- `task/debug.go`: `NewDebugTaskContext` builds an in-memory
`TaskContext` with **`KB.ID == ""` — the single debug signal used across
the ingestion pipeline**. A canvas debug run has no knowledgebase, and
production ingestion always supplies one, so `kb_id == ""` occurs ONLY
in debug mode. Components gate their own side effects on this signal
without any dedicated debug vocabulary (the former `CANVAS_DEBUG_DOC_ID`
marker constant is removed).
- `task/pipeline_executor.go`: `validateTaskContext` no longer requires
a KB when `KB.ID == ""` (debug); debug runs return `collectDebugOutput`
(chunks) instead of a no-op; uploaded bytes are delivered as
`inputs['binary']` for doc-less runs; `injectDebugPageCap` caps the
parser to the first pages for a fast preview via the production
`override_params` channel (Parser cpnID + family).
- `component/tokenizer.go`: `shouldHaveEmbedding` skips embedding when
`kb_id == ""` — the embedder is configured on the knowledgebase, so a
debug run has nothing to resolve against and stays side-effect free.
- `chunker/register.go`: chunk images are uploaded to MinIO only when a
KB is present (persist run). **In debug mode the raw image bytes are
intentionally dropped (`delete(ck, "image")`)** — the debug preview does
not render chunk images, and dropping the bytes keeps them out of memory
and out of the Redis-stored debug log. This is a deliberate trade-off,
not an oversight.
- `component/file.go`: pass through in-memory binary bytes, skipping
`doc_id` -> storage resolution.
- `handler/agent.go` + `agent_webhook.go`: detect `dataflow_canvas` and
run a sync debug returning chunks inline on the existing
chat/completions endpoint; reject DataFlow canvases from webhooks (fixes
the previously dead `== "DataFlow"` check; mirrors Python
`agent_api.py`).
- `parser_dispatch.go`: export `ParserFileFamily` for the executor's
page-cap injection.
### Debug run log + "View result"
Mirrors Python's debug-log contract so the front-end can replay each
component's progress and parsed output:
- `task/debug_log_sink.go`: a `DebugLogSink` records every component's
lifecycle into a `[{component_id, trace}]` array (each trace entry
carries `message`, `progress`, `timestamp`, `elapsed_time`). `Flush`
appends a terminal `END` marker whose first trace message is non-empty
so the front-end detects completion. On failure the END marker is
prefixed `[ERROR]` yet still carries the run, so the failure timeline
renders instead of being stuck empty. Timestamps and `elapsed_time` are
in seconds (matching the rest of the app).
- `task/debug_result_dsl.go`: `BuildDebugResultDSL` builds the `dsl` the
END marker carries — the Go analogue of Python's `Graph.__str__` +
END-marker `dsl` in `rag/flow/pipeline.py`. It combines the static DSL
structure (component_name / downstream / params / graph.nodes) with the
run output map (`output["state"][<id>]`) to emit, per component,
`obj.params.outputs[<format>].value` (chunks / text / json / html /
markdown) — the exact keys the front-end `dataflow-result` page reads to
render each step's parsed chunks. Raw embedding vectors (including the
dimension-scoped `q_<dim>_vec` keys) are stripped so the stored log
stays Python-scale.
- `task/pipeline_executor.go`: after the run, attach the built `dsl` to
the END marker via the `ResultSink` capability.
- `handler/agent.go`: `runCanvasPipelineDebug` generates a stable
`message_id` up-front and always flushes the log (success or failure);
`respondWithDebugResult` returns `message_id` in **both** the success
and the error envelope so the front-end can poll the log. The debug-log
endpoint `GET /agents/:id/logs/:message_id` serves the array.
- `web/src/pages/agent/hooks/use-run-dataflow.ts`: on a run failure,
also surface `message_id` via `setMessageId` so the log sheet renders
the failure timeline (the `[ERROR]` END marker is already written).
Guarded by `if (msgId)`, so it is a safe no-op when the back-end does
not return an id.
## Behavioral notes
- Debug parses only the first pages (`debugPageCapPages`) for a fast
preview; an explicit `pages` cap already present in the ParserConfig is
respected.
- Debug mode does not keep chunk images (see above) and does not compute
embeddings — it exercises parse + chunk only.
## Test plan
- Go: `debug_test.go`, `debug_log_sink_test.go` (trace pairing, END
marker, `[ERROR]` prefix, fractional-second timestamp/elapsed_time, size
caps, and a real-pipeline test asserting the END-marker `dsl` carries
non-empty per-component `params.outputs` with chunks),
`debug_result_dsl_test.go` (flat and real nested `output["state"]`
shapes, vector stripping, format priority),
`debug_pages_integration_test.go`, `pipeline_executor_persist_test.go`,
`handler/agent_pipeline_debug_test.go`, `handler/agent_logs_test.go`
(incl. `TestRunCanvasPipelineDebug_ErrorStillExposesMessageID` /
`TestRespondWithDebugResult_ErrorCarriesMessageID` locking `message_id`
on failure), plus updates to `agent_test.go` / `agent_webhook_test.go` /
`chunker/image_upload_test.go` / `tokenizer*.go`.
- `go build ./...` and `./build.sh --test` for affected packages.
🤖 Generated with [CodeBuddy Code](https://cnb.cool/codebuddy)
---------
Co-authored-by: CodeBuddy Code <noreply@tencent.com>
884 lines
31 KiB
Go
884 lines
31 KiB
Go
//
|
|
// 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.
|
|
//
|
|
|
|
// Tokenizer ingestion component (Phase 2.4 of
|
|
// port-rag-flow-pipeline-to-go.md §4). Port of Python
|
|
// `rag/flow/tokenizer/tokenizer.py`. Computes (a) full-text token
|
|
// counts via the Go tokenizer package and (b) embedding vectors via
|
|
// the tenant's embedding model.
|
|
//
|
|
// SCOPE (honest):
|
|
//
|
|
// - TOKEN COUNTING: matched at the wire level. Each chunk gets
|
|
// `content_ltks` (tokenized string via `tokenizer.Tokenize`) and
|
|
// `content_sm_ltks` (fine-grained variant) when `search_method`
|
|
// includes `full_text`. `title_tks` / `title_sm_tks` mirror the
|
|
// upstream `name` field. Python uses C++ RAGAnalyzer via
|
|
// `rag_tokenizer`; the Go side goes through `internal/tokenizer`
|
|
// which itself calls into the same C++ binding (`internal/binding`).
|
|
// For non-ASCII (CJK) input, Python's `rag_tokenizer.tokenize`
|
|
// falls back gracefully; the Go path uses the CGo analyzer
|
|
// when initialized, otherwise an empty string — see
|
|
// `internal/tokenizer/tokenizer.go:Tokenize` (Infinity engine
|
|
// returns input unchanged; otherwise the C++ binding is used).
|
|
//
|
|
// - EMBEDDING MODEL RESOLUTION: mirrored. Python uses
|
|
// `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 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 `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)
|
|
// and uses an async semaphore (`embed_limiter`). The Go port
|
|
// issues ONE `Encode([]string)` call with the entire chunk
|
|
// list (AD-5a calls out "embedding calls batched, not fanned").
|
|
// Drivers that need to chunk internally can do so — the wire
|
|
// call is one round-trip.
|
|
//
|
|
// - TRACKING: TrackProgress, TrackElapsed. See
|
|
// `internal/agent/runtime/helpers.go` (plan §1 Phase 1).
|
|
// `internal/agent/runtime/helpers.go` (plan §1 Phase 1).
|
|
//
|
|
// - WHAT IS NOT PORTED:
|
|
//
|
|
// - The python `finalize_pdf_chunk` post-step — that
|
|
// normalizes PDF bbox metadata; it lives in
|
|
// `rag/flow/parser/pdf_chunk_metadata.py` and is the Parser
|
|
// component's concern (Phase 2.2).
|
|
//
|
|
// - `rag.flow.tokenizer` `thread_pool_exec` async batching +
|
|
// `embed_limiter` semaphore — replaced by the single
|
|
// batched `Encode` call.
|
|
package component
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"regexp"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"go.uber.org/zap"
|
|
"gorm.io/gorm"
|
|
|
|
"ragflow/internal/agent/runtime"
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/ingestion/component/globals"
|
|
"ragflow/internal/ingestion/component/schema"
|
|
"ragflow/internal/tokenizer"
|
|
"ragflow/internal/utility"
|
|
)
|
|
|
|
const ComponentNameTokenizer = "Tokenizer"
|
|
|
|
// embeddingBatchSize returns the embedding batch size, matching Python's
|
|
// settings.EMBEDDING_BATCH_SIZE. Reads TOKENIZER_EMBEDDING_BATCH_SIZE env
|
|
// var; defaults to 16. Invalid / non-positive values fall back to the
|
|
// default (diff Tokenizer Omission-3).
|
|
func embeddingBatchSize() int {
|
|
if v := os.Getenv("TOKENIZER_EMBEDDING_BATCH_SIZE"); v != "" {
|
|
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
|
return n
|
|
}
|
|
}
|
|
return 16
|
|
}
|
|
|
|
// titleExtRE strips a trailing file-extension (e.g. ".pdf") from the
|
|
// upstream document name before tokenizing it. Mirrors the python
|
|
// `re.sub(r"\.[a-zA-Z]+$", "", name)` in tokenizer.py:137.
|
|
var titleExtRE = regexp.MustCompile(`\.[a-zA-Z]+$`)
|
|
|
|
// htmlTableRE matches HTML table-cell tags so the embedded text fed
|
|
// to the embedding model doesn't carry raw markup. Mirrors the python
|
|
// `re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", txt)` at
|
|
// tokenizer.py:79.
|
|
var htmlTableRE = regexp.MustCompile(`</?(table|td|caption|tr|th)( [^<>]{0,12})?>`)
|
|
|
|
// EmbeddingResult carries a vector plus the model-reported token usage
|
|
// for that input batch entry.
|
|
type EmbeddingResult struct {
|
|
Vector []float64
|
|
TokenCount int
|
|
}
|
|
|
|
// Embedder is the testability seam for the embedding branch.
|
|
type Embedder interface {
|
|
MaxTokens() int
|
|
Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error)
|
|
}
|
|
|
|
// EmbedderResolver resolves the embedder for one tokenizer invocation.
|
|
// 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(ctx context.Context, 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
|
|
// rag/flow/tokenizer/tokenizer.py:Tokenizer.
|
|
//
|
|
// Inputs:
|
|
//
|
|
// tenant_id (string, optional) — used to resolve the embedding model
|
|
// 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
|
|
// markdown/text/html — scalar payload matching output_format
|
|
//
|
|
// Outputs:
|
|
//
|
|
// chunks — the chunk list with tokenized fields
|
|
// and (when embedding is requested)
|
|
// q_<n>_vec vector fields
|
|
// embedding_token_consumption — non-negative int (matches the python
|
|
// `embedding_token_consumption` output)
|
|
// output_format — always "chunks" (matches python set_output)
|
|
// _created_time / _elapsed_time — TrackElapsed bookkeeping
|
|
type TokenizerComponent struct {
|
|
param schema.TokenizerParam
|
|
resolver EmbedderResolver
|
|
embeddingModel string
|
|
}
|
|
|
|
// NewTokenizerComponent constructs a production TokenizerComponent from DSL
|
|
// params. Mirrors python `TokenizerParam` defaults (search_method =
|
|
// ["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 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
|
|
// search_method = ["full_text"] correctly disables
|
|
// embedding. Python's TokenizerParam similarly treats
|
|
// caller-supplied values as the full set.
|
|
p.SearchMethod = nil
|
|
switch t := v.(type) {
|
|
case []any:
|
|
for _, x := range t {
|
|
if s, ok := x.(string); ok {
|
|
p.SearchMethod = append(p.SearchMethod, s)
|
|
}
|
|
}
|
|
case []string:
|
|
p.SearchMethod = append(p.SearchMethod, t...)
|
|
}
|
|
}
|
|
if v, ok := params["filename_embd_weight"]; ok {
|
|
switch t := v.(type) {
|
|
case float64:
|
|
p.FilenameEmbdWeight = t
|
|
case int:
|
|
p.FilenameEmbdWeight = float64(t)
|
|
}
|
|
}
|
|
if v, ok := params["fields"]; ok {
|
|
switch t := v.(type) {
|
|
case string:
|
|
p.Fields = []string{t}
|
|
case []any:
|
|
for _, x := range t {
|
|
if s, ok := x.(string); ok {
|
|
p.Fields = append(p.Fields, s)
|
|
}
|
|
}
|
|
case []string:
|
|
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, 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 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.",
|
|
"text": "Plain-text payload when output_format == \"text\".",
|
|
"markdown": "Markdown payload when output_format == \"markdown\".",
|
|
"html": "HTML payload when output_format == \"html\".",
|
|
"name": "Upstream document name (used for title_tks and the title-blended embedding).",
|
|
}
|
|
}
|
|
|
|
// Outputs returns the parameter metadata. Mirrors python set_output
|
|
// contract for Tokenizer.
|
|
func (c *TokenizerComponent) Outputs() map[string]string {
|
|
return map[string]string{
|
|
"chunks": "Tokenized chunk list (each entry gains content_ltks / content_sm_ltks / title_tks and, when embedding is requested, q_<n>_vec).",
|
|
"embedding_token_consumption": "Non-negative token count consumed by the embedding call. Omitted when no embedding ran.",
|
|
"output_format": "Always \"chunks\" (matches python set_output).",
|
|
"_created_time": "RFC3339Nano creation timestamp (TrackElapsed).",
|
|
"_elapsed_time": "Wall-clock seconds (TrackElapsed).",
|
|
}
|
|
}
|
|
|
|
// Invoke computes tokens + embeddings for the upstream chunks.
|
|
//
|
|
// Failure modes:
|
|
//
|
|
// - "embedding" requested but resolver is nil → returns an
|
|
// error (fail-loud: same contract as python when LLMBundle is
|
|
// unconstructable).
|
|
// - Empty chunks list → returns an empty chunks output without
|
|
// panicking (python tokenizer.py:121 treats this as valid).
|
|
// - Per-chunk empty cleaned text → chunk is skipped from the
|
|
// embedding batch (python tokenizer.py:80-82 `if not cleaned_txt:
|
|
// continue`), but the chunk still carries tokenized fields if
|
|
// `full_text` is in `search_method`.
|
|
func (c *TokenizerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) {
|
|
// 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)
|
|
common.Debug("tokenizer stage",
|
|
zap.String("component", "Tokenizer"),
|
|
zap.Int("input_chunks", len(chunks)),
|
|
)
|
|
titleStem := titleExtRE.ReplaceAllString(name, "")
|
|
|
|
normalizeChunkTextFallback(chunks)
|
|
|
|
// Diff Tokenizer-8: chunk_order_int must be set on all paths
|
|
// (Python sets it unconditionally). tokenizeChunks also sets it
|
|
// for the full_text path; this covers the embedding-only path.
|
|
for i := range chunks {
|
|
if chunks[i].ChunkOrderInt == nil {
|
|
chunks[i].ChunkOrderInt = intPtr(i)
|
|
}
|
|
}
|
|
|
|
language := globals.GlobalOrInput(ctx, inputs, "lang", "English")
|
|
|
|
if contains(c.param.SearchMethod, "full_text") {
|
|
if err := tokenizeChunks(chunks, titleStem, language); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
out := map[string]any{
|
|
"output_format": "chunks",
|
|
"chunks": schema.ChunkDocsToMaps(chunks),
|
|
}
|
|
|
|
// Embedding requires a KB: the embedder (and its embd_id) is configured
|
|
// on the knowledgebase, so without kb_id there is nothing to resolve
|
|
// against. A canvas-debug (dry-run) run has kb_id == "" by construction,
|
|
// so embedding is skipped there — debug only exercises parse+chunk and
|
|
// must stay side-effect free.
|
|
if shouldHaveEmbedding(c.param.SearchMethod, kbID) {
|
|
chunks, tokenCount, err := c.embedChunks(ctx, tenantID, kbID, embeddingModel, name, chunks)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out["embedding_token_consumption"] = tokenCount
|
|
out["chunks"] = schema.ChunkDocsToMaps(chunks)
|
|
}
|
|
if err := validateTokenizerOutputs(chunks, c.param.SearchMethod, c.param.Fields, kbID); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
common.Debug("tokenizer stage",
|
|
zap.String("component", "Tokenizer"),
|
|
zap.Int("output_chunks", len(chunks)),
|
|
)
|
|
return out, nil
|
|
}
|
|
|
|
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
|
|
}
|
|
// 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
|
|
}
|
|
if resolver == nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: embedding requested but no embedder resolver configured")
|
|
}
|
|
embedder, err := resolver(ctx, tenantID, kbID, embeddingModel)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: resolve embedder: %w", err)
|
|
}
|
|
if embedder == nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: embedding requested but encoder resolution returned nil")
|
|
}
|
|
|
|
texts := make([]string, 0, len(chunks))
|
|
pairs := make([]int, 0, len(chunks))
|
|
for i, ck := range chunks {
|
|
raw := concatFields(ck, c.param.Fields)
|
|
txt := htmlTableRE.ReplaceAllString(raw, " ")
|
|
txt = strings.TrimSpace(txt)
|
|
trunc := truncateForEmbedding(txt, embedder.MaxTokens())
|
|
if txt == "" {
|
|
continue
|
|
}
|
|
texts = append(texts, trunc)
|
|
pairs = append(pairs, i)
|
|
}
|
|
if len(texts) == 0 {
|
|
return chunks, 0, nil
|
|
}
|
|
|
|
trimmedName := strings.TrimSpace(name)
|
|
var (
|
|
titleVec []float64
|
|
tokenCount int
|
|
hasTitleVec bool
|
|
)
|
|
if trimmedName == "" {
|
|
log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting")
|
|
} else {
|
|
// Encode the raw name (no TrimSpace) to mirror Python
|
|
// tokenizer.py:95 which passes name verbatim to embedding. The
|
|
// empty-name guard above still uses TrimSpace, matching Python's
|
|
// `.strip()==""` check at tokenizer.py:200.
|
|
titleResults, err := embedder.Encode(ctx, []string{name})
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: encode title: %w", err)
|
|
}
|
|
if len(titleResults) != 1 {
|
|
return nil, 0, fmt.Errorf("Tokenizer: encode title returned %d vectors for 1 chunk", len(titleResults))
|
|
}
|
|
titleVec = titleResults[0].Vector
|
|
tokenCount = titleResults[0].TokenCount
|
|
hasTitleVec = true
|
|
}
|
|
|
|
contentResults := make([]EmbeddingResult, 0, len(texts))
|
|
for start := 0; start < len(texts); start += embeddingBatchSize() {
|
|
end := start + embeddingBatchSize()
|
|
if end > len(texts) {
|
|
end = len(texts)
|
|
}
|
|
batchResults, err := embedder.Encode(ctx, texts[start:end])
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: encode: %w", err)
|
|
}
|
|
if len(batchResults) != end-start {
|
|
return nil, 0, fmt.Errorf("Tokenizer: encode returned %d vectors for %d chunks", len(batchResults), end-start)
|
|
}
|
|
for _, result := range batchResults {
|
|
tokenCount += result.TokenCount
|
|
}
|
|
contentResults = append(contentResults, batchResults...)
|
|
}
|
|
|
|
titleWeight := c.param.FilenameEmbdWeight
|
|
for i, idx := range pairs {
|
|
merged := append([]float64(nil), contentResults[i].Vector...)
|
|
if hasTitleVec {
|
|
merged, err = mergeEmbeddingVectors(titleVec, contentResults[i].Vector, titleWeight)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: merge vectors: %w", err)
|
|
}
|
|
}
|
|
if err := chunks[idx].SetExtraValue(fmt.Sprintf("q_%d_vec", len(merged)), merged); err != nil {
|
|
return nil, 0, fmt.Errorf("Tokenizer: vector marshal: %w", err)
|
|
}
|
|
}
|
|
return chunks, tokenCount, nil
|
|
}
|
|
|
|
// defaultEmbeddingTokenLimit is the safe fallback used when an embedder reports
|
|
// no token limit (maxTokens <= 0). It both prevents empty embedding inputs and
|
|
// keeps truncation active for every path instead of passing the full text through.
|
|
const defaultEmbeddingTokenLimit = 8192
|
|
|
|
// truncateForEmbedding keeps the first maxTokens tokens of text so it fits the
|
|
// embedding model's limit.
|
|
//
|
|
// For a positive maxTokens it mirrors Python common/token_utils.py:183-185
|
|
// `truncate(string, max_len)` (keep the first max_len tokens).
|
|
//
|
|
// An unconfigured embedder reports maxTokens <= 0. Rather than mirror Python's
|
|
// behaviour of returning "" (which would make the embeddings API reject the whole
|
|
// batch with "inputs cannot be empty"), Go clamps the limit to a safe default
|
|
// (defaultEmbeddingTokenLimit = 8192). This both prevents empty inputs AND keeps
|
|
// truncation active for every path (Builtin and generic) instead of silently
|
|
// passing the full, untruncated text when no limit is configured.
|
|
func truncateForEmbedding(text string, maxTokens int) string {
|
|
if maxTokens <= 0 {
|
|
maxTokens = defaultEmbeddingTokenLimit
|
|
}
|
|
// Keep a 10-token safety margin, mirroring Python's embedding path
|
|
// (rag/svr/task_executor.py uses `mdl.max_length - 10`). Only apply it
|
|
// when the limit is large enough; for small limits (<=10) keep the full
|
|
// value so the result stays non-empty instead of collapsing to "".
|
|
if maxTokens > 10 {
|
|
maxTokens -= 10
|
|
}
|
|
return tokenizer.TrimContentToTokenLimit(text, maxTokens)
|
|
}
|
|
|
|
func mergeEmbeddingVectors(titleVec, contentVec []float64, titleWeight float64) ([]float64, error) {
|
|
if len(titleVec) == 0 || len(contentVec) == 0 {
|
|
return nil, fmt.Errorf("empty embedding vector")
|
|
}
|
|
if len(titleVec) != len(contentVec) {
|
|
return nil, fmt.Errorf("unexpected embedding dimensions")
|
|
}
|
|
merged := make([]float64, len(titleVec))
|
|
for i := range titleVec {
|
|
merged[i] = titleWeight*titleVec[i] + (1-titleWeight)*contentVec[i]
|
|
}
|
|
return merged, nil
|
|
}
|
|
|
|
func decodeTokenizerFromUpstream(inputs map[string]any) (schema.TokenizerFromUpstream, error) {
|
|
var out schema.TokenizerFromUpstream
|
|
if inputs == nil {
|
|
return out, fmt.Errorf("Tokenizer: inputs map is nil")
|
|
}
|
|
data, err := json.Marshal(stripRuntimeTimestamps(inputs))
|
|
if err != nil {
|
|
return out, fmt.Errorf("Tokenizer: encode inputs: %w", err)
|
|
}
|
|
if err := json.Unmarshal(data, &out); err != nil {
|
|
return out, fmt.Errorf("Tokenizer: decode inputs: %w", err)
|
|
}
|
|
if err := out.Validate(); err != nil {
|
|
return out, fmt.Errorf("Tokenizer: input error: %w", err)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func stripRuntimeTimestamps(inputs map[string]any) map[string]any {
|
|
out := make(map[string]any, len(inputs))
|
|
for k, v := range inputs {
|
|
if k == "_created_time" || k == "_elapsed_time" {
|
|
continue
|
|
}
|
|
out[k] = v
|
|
}
|
|
return out
|
|
}
|
|
|
|
func chunksFromTokenizerUpstream(in schema.TokenizerFromUpstream) []schema.ChunkDoc {
|
|
var raw []schema.ChunkDoc
|
|
switch in.OutputFormat {
|
|
case schema.PayloadFormatChunks:
|
|
raw = cloneChunkDocs(in.Chunks)
|
|
case schema.PayloadFormatMarkdown:
|
|
raw = textPayloadToChunks(in.MarkdownResult)
|
|
case schema.PayloadFormatText:
|
|
raw = textPayloadToChunks(in.TextResult)
|
|
case schema.PayloadFormatHTML:
|
|
raw = textPayloadToChunks(in.HTMLResult)
|
|
default:
|
|
raw = cloneChunkDocs(in.JSONResult)
|
|
}
|
|
// Discard zero-value ChunkDocs (no text, no content_with_weight, no
|
|
// image, no summary) so they don't produce phantom embeddings
|
|
// downstream. Python's pipeline filters none-chunks before the
|
|
// tokenizer.
|
|
filtered := raw[:0]
|
|
for _, ck := range raw {
|
|
if isPhantomChunk(ck) {
|
|
continue
|
|
}
|
|
filtered = append(filtered, ck)
|
|
}
|
|
return filtered
|
|
}
|
|
|
|
// isPhantomChunk returns true when a ChunkDoc has no usable content for
|
|
// downstream tokenization or embedding. A chunk carrying only a Summary
|
|
// (no Text/Image/ContentWithWeight) is kept — tokenizeChunks tokenizes the
|
|
// Summary in that case. Mirrors Python's none-chunk filtering.
|
|
func isPhantomChunk(ck schema.ChunkDoc) bool {
|
|
return ck.Text == "" && ck.Image == "" && ck.ContentWithWeight == "" && ck.Summary == ""
|
|
}
|
|
|
|
func textPayloadToChunks(payload *string) []schema.ChunkDoc {
|
|
if payload == nil || strings.TrimSpace(*payload) == "" {
|
|
return []schema.ChunkDoc{}
|
|
}
|
|
return []schema.ChunkDoc{{Text: *payload}}
|
|
}
|
|
|
|
func cloneChunkDocs(in []schema.ChunkDoc) []schema.ChunkDoc {
|
|
if len(in) == 0 {
|
|
return []schema.ChunkDoc{}
|
|
}
|
|
out := make([]schema.ChunkDoc, len(in))
|
|
for i := range in {
|
|
out[i] = cloneTokenizerChunkDoc(in[i])
|
|
}
|
|
return out
|
|
}
|
|
|
|
func cloneTokenizerChunkDoc(in schema.ChunkDoc) schema.ChunkDoc {
|
|
out := in
|
|
if in.TKNums != nil {
|
|
v := *in.TKNums
|
|
out.TKNums = &v
|
|
}
|
|
if in.ChunkOrderInt != nil {
|
|
v := *in.ChunkOrderInt
|
|
out.ChunkOrderInt = &v
|
|
}
|
|
if in.PageNumber != nil {
|
|
v := *in.PageNumber
|
|
out.PageNumber = &v
|
|
}
|
|
if in.Extra != nil {
|
|
out.Extra = make(map[string]json.RawMessage, len(in.Extra))
|
|
for k, v := range in.Extra {
|
|
out.Extra[k] = append(json.RawMessage(nil), v...)
|
|
}
|
|
}
|
|
if len(in.PDFPositions) > 0 {
|
|
out.PDFPositions = append(json.RawMessage(nil), in.PDFPositions...)
|
|
}
|
|
if len(in.Positions) > 0 {
|
|
out.Positions = append(json.RawMessage(nil), in.Positions...)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// normalizeChunkTextFallback populates each chunk's "text" key
|
|
// from "content_with_weight" when "text" is absent or empty. Mirrors
|
|
// the python rag/flow/tokenizer.py:111 fallback so a chunk that
|
|
// arrives from the parser path with only the structured
|
|
// content_with_weight field still tokenizes.
|
|
//
|
|
// The function mutates the input slice in place; callers should
|
|
// not retain separate copies of the chunks map. If both fields
|
|
// are present, the existing "text" wins — preserves the python
|
|
// contract where the chunker's emitted text is authoritative.
|
|
func normalizeChunkTextFallback(chunks []schema.ChunkDoc) {
|
|
for i := range chunks {
|
|
if chunks[i].Text != "" {
|
|
continue
|
|
}
|
|
if chunks[i].ContentWithWeight != "" {
|
|
chunks[i].Text = chunks[i].ContentWithWeight
|
|
}
|
|
}
|
|
}
|
|
|
|
// tokenizeChunks annotates each chunk with title_tks, content_ltks,
|
|
// and (when applicable) question_tks / important_tks / summary fields.
|
|
// Mirrors python tokenizer.py:130-185 and rag/nlp/__init__.py tokenize() /
|
|
// tokenize_chunks().
|
|
//
|
|
// language sets the Snowball stemmer language, matching Python's
|
|
// rag_tokenizer.tokenizer.set_language(language) call inside tokenize().
|
|
func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string, language string) error {
|
|
tok := tokenizer.New(language)
|
|
for i := range chunks {
|
|
ck := &chunks[i]
|
|
ck.ChunkOrderInt = intPtr(i)
|
|
titleTk, err := tok.Tokenize(titleStem)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: title tokenize: %w", err)
|
|
}
|
|
titleSmTk, err := tok.FineGrainedTokenize(titleTk)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: title fine-grain: %w", err)
|
|
}
|
|
ck.TitleTks = titleTk
|
|
ck.TitleSmTks = titleSmTk
|
|
|
|
// Question / keyword / summary fields are optional. The python
|
|
// path branches on each independently.
|
|
if q := ck.Questions; q != "" {
|
|
if err := ck.SetExtraValue("question_kwd", strings.Split(q, "\n")); err != nil {
|
|
return fmt.Errorf("Tokenizer: question keywords marshal: %w", err)
|
|
}
|
|
qt, err := tok.Tokenize(q)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: question tokenize: %w", err)
|
|
}
|
|
if err := ck.SetExtraValue("question_tks", qt); err != nil {
|
|
return fmt.Errorf("Tokenizer: question tokens marshal: %w", err)
|
|
}
|
|
}
|
|
if kw := ck.Keywords; kw != "" {
|
|
if err := ck.SetExtraValue("important_kwd", utility.SplitKeywords(kw)); err != nil {
|
|
return fmt.Errorf("Tokenizer: keyword list marshal: %w", err)
|
|
}
|
|
it, err := tok.Tokenize(kw)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: keyword tokenize: %w", err)
|
|
}
|
|
if err := ck.SetExtraValue("important_tks", it); err != nil {
|
|
return fmt.Errorf("Tokenizer: keyword tokens marshal: %w", err)
|
|
}
|
|
}
|
|
// Keep Go: skip whitespace-only summaries so they don't shadow
|
|
// the real Text. Python's truthy check (tokenizer.py:155) treats
|
|
// " " as present and blanks out content_ltks; Go is more sensible.
|
|
if s := strings.TrimSpace(ck.Summary); s != "" {
|
|
st, err := tok.Tokenize(s)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: summary tokenize: %w", err)
|
|
}
|
|
if st == "" {
|
|
st = s
|
|
}
|
|
ck.ContentLtks = st
|
|
smt, err := tok.FineGrainedTokenize(st)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: summary fine-grain: %w", err)
|
|
}
|
|
if smt == "" {
|
|
smt = st
|
|
}
|
|
ck.ContentSmLtks = smt
|
|
} else if t := ck.Text; strings.TrimSpace(t) != "" {
|
|
tt, err := tok.Tokenize(t)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: text tokenize: %w", err)
|
|
}
|
|
if tt == "" {
|
|
tt = t
|
|
}
|
|
ck.ContentLtks = tt
|
|
smt, err := tok.FineGrainedTokenize(tt)
|
|
if err != nil {
|
|
return fmt.Errorf("Tokenizer: text fine-grain: %w", err)
|
|
}
|
|
if smt == "" {
|
|
smt = tt
|
|
}
|
|
ck.ContentSmLtks = smt
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// concatFields concatenates the configured fields of a chunk into
|
|
// a single string. Mirrors python tokenizer.py:69-79 which
|
|
// concatenates `param.fields` (string or list-of-strings per chunk).
|
|
func concatFields(ck schema.ChunkDoc, fields []string) string {
|
|
var b strings.Builder
|
|
for _, f := range fields {
|
|
switch f {
|
|
case "text":
|
|
b.WriteString(ck.Text)
|
|
case "content_with_weight":
|
|
b.WriteString(ck.ContentWithWeight)
|
|
case "questions":
|
|
b.WriteString(ck.Questions)
|
|
case "keywords":
|
|
b.WriteString(ck.Keywords)
|
|
case "summary":
|
|
b.WriteString(ck.Summary)
|
|
default:
|
|
if s, ok := ck.GetExtraString(f); ok {
|
|
b.WriteString(s)
|
|
continue
|
|
}
|
|
if values, ok := ck.GetExtraStringSlice(f); ok {
|
|
b.WriteString(strings.Join(values, "\n"))
|
|
}
|
|
}
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
// shouldHaveEmbedding reports whether the tokenizer must attach embedding
|
|
// vectors: the search method requests embedding AND a KB is present.
|
|
func shouldHaveEmbedding(searchMethods []string, kbID string) bool {
|
|
return contains(searchMethods, "embedding") && kbID != ""
|
|
}
|
|
|
|
func validateTokenizerOutputs(chunks []schema.ChunkDoc, searchMethods, fields []string, kbID string) error {
|
|
needFullText := contains(searchMethods, "full_text")
|
|
needEmbedding := shouldHaveEmbedding(searchMethods, kbID)
|
|
if !needFullText && !needEmbedding {
|
|
return nil
|
|
}
|
|
for i := range chunks {
|
|
if needFullText && requiresFullTextTokens(chunks[i]) {
|
|
if strings.TrimSpace(chunks[i].ContentLtks) == "" || strings.TrimSpace(chunks[i].ContentSmLtks) == "" {
|
|
return fmt.Errorf("Tokenizer: chunk[%d] missing full_text tokens", i)
|
|
}
|
|
}
|
|
if needEmbedding && requiresEmbeddingVector(chunks[i], fields) {
|
|
if !hasEmbeddingVector(chunks[i]) {
|
|
return fmt.Errorf("Tokenizer: chunk[%d] missing embedding vector", i)
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func requiresFullTextTokens(ck schema.ChunkDoc) bool {
|
|
return strings.TrimSpace(ck.Summary) != "" || strings.TrimSpace(ck.Text) != ""
|
|
}
|
|
|
|
func requiresEmbeddingVector(ck schema.ChunkDoc, fields []string) bool {
|
|
return strings.TrimSpace(cleanEmbeddingText(concatFields(ck, fields))) != ""
|
|
}
|
|
|
|
func cleanEmbeddingText(text string) string {
|
|
return strings.TrimSpace(htmlTableRE.ReplaceAllString(text, " "))
|
|
}
|
|
|
|
func hasEmbeddingVector(ck schema.ChunkDoc) bool {
|
|
if len(ck.Extra) == 0 {
|
|
return false
|
|
}
|
|
for key, raw := range ck.Extra {
|
|
if !strings.HasPrefix(key, "q_") || !strings.HasSuffix(key, "_vec") {
|
|
continue
|
|
}
|
|
var vec []float64
|
|
if err := json.Unmarshal(raw, &vec); err != nil {
|
|
continue
|
|
}
|
|
if len(vec) > 0 {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func getStringOr(m map[string]any, key, def string) string {
|
|
if v, ok := m[key].(string); ok && v != "" {
|
|
return v
|
|
}
|
|
return def
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
func contains(s []string, v string) bool {
|
|
return slices.Contains(s, v)
|
|
}
|
|
|
|
func intPtr(v int) *int { return &v }
|
|
|
|
// init registers Tokenizer under CategoryIngestion (plan §4
|
|
// Phase 2.4). The metadata drives Phase 4's GET /api/v1/components
|
|
// listing.
|
|
func init() {
|
|
c := &TokenizerComponent{}
|
|
runtime.MustRegister(ComponentNameTokenizer, runtime.CategoryIngestion,
|
|
func(_ string, params map[string]any) (runtime.Component, error) {
|
|
return NewTokenizerComponent(params)
|
|
},
|
|
runtime.Metadata{
|
|
Version: "1.0.0",
|
|
Inputs: c.Inputs(),
|
|
Outputs: c.Outputs(),
|
|
})
|
|
}
|