mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-10 13:24:05 +08:00
1217 lines
47 KiB
Go
1217 lines
47 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 the model solver injected by the ingestion task package
|
|
// (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"
|
|
"regexp"
|
|
"slices"
|
|
"strings"
|
|
"unicode"
|
|
"unicode/utf8"
|
|
|
|
"go.uber.org/zap"
|
|
"gorm.io/gorm"
|
|
|
|
"ragflow/internal/agent/runtime"
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/ingestion/chunkcache"
|
|
"ragflow/internal/ingestion/component/globals"
|
|
"ragflow/internal/ingestion/component/schema"
|
|
"ragflow/internal/tokenizer"
|
|
)
|
|
|
|
const ComponentNameTokenizer = schema.ComponentNameTokenizer
|
|
|
|
// 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]+$`)
|
|
|
|
// Table chunks carry their row rendered as an HTML <table> on Text (produced
|
|
// by the chunker), so the embedding input intentionally includes the table
|
|
// markup — there is no markup to strip. The legacy htmlTableRE stays removed.
|
|
|
|
// EmbeddingResult carries a vector plus the model-reported token usage
|
|
// for that input batch entry.
|
|
// esKeywordMaxTermBytes is the upper bound for a single term stored in an
|
|
// Elasticsearch keyword field.
|
|
const esKeywordMaxTermBytes = 32766
|
|
|
|
// sanitizeKeywordTerm returns a keyword that fits into an Elasticsearch keyword
|
|
// field. Small terms are returned unchanged; oversized terms are truncated at a
|
|
// UTF-8 character boundary and trailing whitespace is stripped.
|
|
func sanitizeKeywordTerm(term string) string {
|
|
term = strings.TrimSpace(term)
|
|
if term == "" {
|
|
return ""
|
|
}
|
|
if len(term) <= esKeywordMaxTermBytes {
|
|
return term
|
|
}
|
|
length := 0
|
|
end := 0
|
|
for _, r := range term {
|
|
size := utf8.RuneLen(r)
|
|
if length+size > esKeywordMaxTermBytes {
|
|
break
|
|
}
|
|
length += size
|
|
end += size
|
|
}
|
|
return strings.TrimRightFunc(term[:end], unicode.IsSpace)
|
|
}
|
|
|
|
type EmbeddingResult struct {
|
|
Vector []float64
|
|
TokenCount int
|
|
}
|
|
|
|
// Embedder is the testability seam for the embedding branch.
|
|
type Embedder interface {
|
|
MaxTokens() int
|
|
BatchSize() int
|
|
Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error)
|
|
}
|
|
|
|
// EmbedderTrimmer is an OPTIONAL extension of Embedder, implemented by embedders
|
|
// that can count with the embedding model's own tokenizer (or a calibrated upper
|
|
// bound of it) instead of cl100k_base. When present the component uses it for
|
|
// every truncation, because cl100k and a model's own tokenizer disagree by up to
|
|
// ~2% and the sign of the disagreement depends on the content: a chunk that
|
|
// cl100k scores just under the limit can be over it in the model's tokenizer,
|
|
// which the provider answers with a 400 and, before this, the loss of the whole
|
|
// document. See internal/tokenizer/embedding_token_limits.md.
|
|
type EmbedderTrimmer interface {
|
|
Trim(text string) (trimmed string, tokens int)
|
|
}
|
|
|
|
// EmbedderMaxResolver is the optional counterpart that exposes the input window
|
|
// to honour: the model's declared value, then the provider catalog's
|
|
// context_length. Without it the component has to guess, and guessing 8192
|
|
// overshoots the window of every model with a smaller one.
|
|
type EmbedderMaxResolver interface {
|
|
ResolveMaxTokens() int
|
|
}
|
|
|
|
// trimForEmbedding trims `text` the best way the embedder allows.
|
|
func trimForEmbedding(embedder Embedder, text string) string {
|
|
if trimmer, ok := embedder.(EmbedderTrimmer); ok {
|
|
trimmed, _ := trimmer.Trim(text)
|
|
return trimmed
|
|
}
|
|
return truncateForEmbedding(text, resolveEmbedderMaxTokens(embedder))
|
|
}
|
|
|
|
// resolveEmbedderMaxTokens is the window the embedder will actually honour.
|
|
func resolveEmbedderMaxTokens(embedder Embedder) int {
|
|
if resolver, ok := embedder.(EmbedderMaxResolver); ok {
|
|
return resolver.ResolveMaxTokens()
|
|
}
|
|
return embedder.MaxTokens()
|
|
}
|
|
|
|
// tableStructureTagRe mirrors the Python tokenizer's cleaned_txt step
|
|
// (rag/flow/tokenizer/tokenizer.py:86): table-structure tags are replaced
|
|
// with a space before the emptiness check and before the text is embedded,
|
|
// so a chunk that is nothing but an empty table carries no embeddable
|
|
// content. The attribute group is deliberately bounded to 12 non-tag
|
|
// characters, exactly like the Python pattern: a tag with longer attributes
|
|
// survives the strip. Do not "fix" the bound — it is parity, not a bug.
|
|
var tableStructureTagRe = regexp.MustCompile(`</?(table|td|caption|tr|th)( [^<>]{0,12})?>`)
|
|
|
|
// cleanEmbedText applies the table-structure strip and trims, mirroring
|
|
// Python's cleaned_txt. The stripped text is both the emptiness check and
|
|
// what gets embedded — Python appends cleaned_txt to the embedding batch
|
|
// (tokenizer.py:89), not the raw chunk.
|
|
func cleanEmbedText(raw string) string {
|
|
return strings.TrimSpace(tableStructureTagRe.ReplaceAllString(raw, " "))
|
|
}
|
|
|
|
// EmbedderResolver resolves the embedder and its dataset-bound embedding-model
|
|
// id for one tokenizer invocation. The resolver derives the model exclusively
|
|
// from the knowledgebase's configured embd_id (see internal/ingestion/task/
|
|
// embedder.go); it never reads any embedding identifier from the DSL. The
|
|
// returned embdID is the stable string the tokenizer keys its per-chunk
|
|
// embedding cache on, so a KB whose embedding model changes yields a different
|
|
// key and never serves a stale vector.
|
|
type EmbedderResolver func(ctx context.Context, tenantID, kbID string) (Embedder, string, 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 selects the embedding
|
|
// model. The model always comes from the dataset;
|
|
// the DSL never configures it.
|
|
// 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
|
|
}
|
|
|
|
// 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()
|
|
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 {
|
|
// Replace (not append) so fields = ["summary"] does not
|
|
// silently embed the default "text" as well. Blank names
|
|
// are dropped: an empty selection keeps the default
|
|
// instead of selecting no text at all.
|
|
var fields []string
|
|
add := func(s string) {
|
|
if s = strings.TrimSpace(s); s != "" {
|
|
fields = append(fields, s)
|
|
}
|
|
}
|
|
switch t := v.(type) {
|
|
case string:
|
|
add(t)
|
|
case []any:
|
|
for _, x := range t {
|
|
if s, ok := x.(string); ok {
|
|
add(s)
|
|
}
|
|
}
|
|
case []string:
|
|
for _, s := range t {
|
|
add(s)
|
|
}
|
|
}
|
|
if len(fields) > 0 {
|
|
p.Fields = fields
|
|
}
|
|
}
|
|
}
|
|
if err := p.Validate(); err != nil {
|
|
return nil, fmt.Errorf("tokenizer: param check: %w", err)
|
|
}
|
|
return &TokenizerComponent{param: p, resolver: resolver}, nil
|
|
}
|
|
|
|
// 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": "Knowledgebase identifier used to resolve the bound embedding model (kb.embd_id). The embedding model is taken exclusively from the dataset; the DSL must not configure it.",
|
|
"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", "")
|
|
|
|
// 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, "")
|
|
if toks := declaredTitleTokens(chunks); len(toks) > 0 {
|
|
titleStem = strings.TrimSpace(titleStem + " " + strings.Join(toks, " "))
|
|
}
|
|
|
|
// chunk_order_int is the position of the chunk in the (post-filter) reading
|
|
// sequence. It is set unconditionally on every surviving chunk so that all
|
|
// retrievable chunks carry a stable reading-order index on every path, not
|
|
// just chunks+full_text. Because it enumerates the slice after filtering,
|
|
// the values are contiguous and 0-based within Go.
|
|
for i := range chunks {
|
|
chunks[i].ChunkOrderInt = intPtr(i)
|
|
}
|
|
|
|
language := globals.GlobalOrInput(ctx, inputs, "lang", "English")
|
|
|
|
if slices.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),
|
|
}
|
|
copyPipelineControlValues(out, inputs)
|
|
|
|
// 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, name, chunks, chunkcache.Client())
|
|
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 copyPipelineControlValues(output, input map[string]any) {
|
|
for _, key := range []string{"wiki_active_map_states"} {
|
|
if value, exists := input[key]; exists {
|
|
output[key] = value
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, name string, chunks []schema.ChunkDoc, store chunkcache.Store) ([]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, embdID, err := resolver(ctx, tenantID, kbID)
|
|
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")
|
|
}
|
|
|
|
// store may be nil (no redis configured / unit test); cache lookups then
|
|
// degrade to a no-op and every chunk is embedded normally.
|
|
|
|
// resolved[i] holds the content embedding vector for chunks[i] once known:
|
|
// either served from the per-chunk cache (isHit) or produced by the batch
|
|
// Encode below. nil means the chunk produced no embeddable text. trunc is the
|
|
// exact text fed to the model, carried so the Set-side key matches the
|
|
// Get-side key (the key includes the embedded text, not just the chunk id).
|
|
type resolvedVec struct {
|
|
content []float64
|
|
isHit bool
|
|
trunc string
|
|
}
|
|
resolved := make([]*resolvedVec, len(chunks))
|
|
texts := make([]string, 0, len(chunks))
|
|
pairs := make([]int, 0, len(chunks))
|
|
// truncs[i] is the embedded text for texts[i] / pairs[i]; carried so the
|
|
// Set-side key matches the Get-side key for freshly embedded content.
|
|
truncs := make([]string, 0, len(chunks))
|
|
// cacheHits counts the chunks served from the per-chunk embedding cache;
|
|
// they are already-resolved work, so the progress fraction starts from them.
|
|
cacheHits := 0
|
|
// fieldsResolvedNoText counts chunks where the configured fields produced
|
|
// nothing at all, which is the signature of a `fields` value naming a chunk
|
|
// field this pipeline does not produce.
|
|
fieldsResolvedNoText := 0
|
|
// emptyAfterTrim counts chunks whose text survived that check but produced
|
|
// nothing after trimming, which means the embedder's trimmer returned no
|
|
// text for non-empty input.
|
|
emptyAfterTrim := 0
|
|
for i, ck := range chunks {
|
|
raw := concatFields(ck, c.param.Fields)
|
|
// Python parity (tokenizer.py:86-89): strip table-structure tags
|
|
// before the emptiness check and embed the stripped text, so an
|
|
// empty-table chunk contributes nothing to the batch.
|
|
txt := cleanEmbedText(raw)
|
|
if txt == "" {
|
|
// Two different causes reach this branch and only one is a defect:
|
|
// no configured field produced text at all (misconfiguration),
|
|
// versus a chunk that really is only table tags or whitespace (the
|
|
// supported empty-table case, which must stay silent). Only a
|
|
// literally empty concatFields result is the first — whitespace-only
|
|
// content is still content a field produced.
|
|
if raw == "" {
|
|
fieldsResolvedNoText++
|
|
}
|
|
continue
|
|
}
|
|
// Per-chunk embedding cache: identical (dataset embd_id, chunk, embedded
|
|
// text) triples reuse the previous content vector, skipping the embed
|
|
// round-trip on resume. The key includes the text that is actually
|
|
// embedded — after field selection (c.param.Fields) and truncation — not
|
|
// just the chunk id. A tokenizer-config or model change alters the
|
|
// embedded input; reusing the prior vector for that case would serve a
|
|
// stale embedding for up to the cache TTL. embdID must be non-empty or
|
|
// every model would collapse onto one key and served vectors could come
|
|
// from a different model.
|
|
trunc := trimForEmbedding(embedder, txt)
|
|
// A provider rejects the WHOLE batch when any input carries no content
|
|
// token — vLLM answers 400 "Input validation error: `inputs` cannot be
|
|
// empty" — so an empty input must never reach the batch. The
|
|
// EmbedderTrimmer seam promises nothing about the result, so a trimmer
|
|
// that returns no text for non-empty input is handled by embedding the
|
|
// pre-trim text: that keeps this chunk's vector requirement intact
|
|
// (validateTokenizerOutputs decides embeddability on the pre-trim text)
|
|
// and lets the embedder's own over-limit ladder re-cut it if the
|
|
// untrimmed text is too long.
|
|
if strings.TrimSpace(trunc) == "" {
|
|
emptyAfterTrim++
|
|
trunc = txt
|
|
}
|
|
if chunkID, ok := ck.GetExtraString("id"); ok && embdID != "" && store != nil {
|
|
if cached, hit := chunkcache.Get(ctx, store, chunkcache.Key("emb", embdID, chunkID, trunc)); hit {
|
|
var vec []float64
|
|
if err := json.Unmarshal([]byte(cached), &vec); err == nil && len(vec) > 0 {
|
|
resolved[i] = &resolvedVec{content: vec, isHit: true, trunc: trunc}
|
|
cacheHits++
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
texts = append(texts, trunc)
|
|
truncs = append(truncs, trunc)
|
|
pairs = append(pairs, i)
|
|
}
|
|
// Both warnings sit outside the early return on purpose: a mixed batch can
|
|
// embed some chunks and drop others, and the dropped ones must still be
|
|
// reported. Each reports how many chunks it covers, not the batch size.
|
|
if fieldsResolvedNoText > 0 {
|
|
common.Warn("Tokenizer: configured fields produced no text for some chunks; they were skipped",
|
|
zap.Strings("fields", c.param.Fields),
|
|
zap.Int("chunks", fieldsResolvedNoText),
|
|
zap.Int("total_chunks", len(chunks)))
|
|
}
|
|
if emptyAfterTrim > 0 {
|
|
common.Warn("Tokenizer: embedder returned no text after trimming; embedded the untrimmed text instead",
|
|
zap.Int("chunks", emptyAfterTrim),
|
|
zap.Int("total_chunks", len(chunks)))
|
|
}
|
|
if len(texts) == 0 && cacheHits == 0 {
|
|
// Nothing to embed and nothing to serve from cache: return before the
|
|
// title encode, exactly like Python's `if not texts: return` — an
|
|
// all-skipped batch must not burn a title embedding. (When the cache
|
|
// serves every chunk the flow continues so those vectors attach.)
|
|
return chunks, 0, nil
|
|
}
|
|
|
|
trimmedName := strings.TrimSpace(name)
|
|
var (
|
|
titleVec []float64
|
|
tokenCount int
|
|
hasTitleVec bool
|
|
)
|
|
// go_intentional (A3): when the upstream name is empty we skip title
|
|
// weighting entirely (hasTitleVec stays false, so the merged vector is the
|
|
// content vector alone). From the end-user perspective an empty title must
|
|
// not contribute the filename embedding weight; the Python DSL instead
|
|
// computes 0.1*emb(""), injecting an undefined bias into every chunk. Go's
|
|
// skip is the correct behavior (go_intentional). Do NOT "align" this to
|
|
// the DSL.
|
|
if trimmedName == "" {
|
|
common.Warn("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))
|
|
batchSize := embedder.BatchSize()
|
|
if batchSize <= 0 {
|
|
return nil, 0, fmt.Errorf("tokenizer: embedder reported non-positive batch size %d", batchSize)
|
|
}
|
|
embedTotal := cacheHits + len(texts)
|
|
for start := 0; start < len(texts); start += batchSize {
|
|
end := start + batchSize
|
|
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...)
|
|
runtime.ReportComponentFraction(ctx, float64(cacheHits+end)/float64(embedTotal))
|
|
}
|
|
if len(texts) == 0 && embedTotal > 0 {
|
|
// With no cache misses the batch loop never runs, so the phase would
|
|
// otherwise report nothing at all; every content embedding is already
|
|
// resolved, which is the loop's completion value.
|
|
runtime.ReportComponentFraction(ctx, 1)
|
|
}
|
|
|
|
titleWeight := c.param.FilenameEmbdWeight
|
|
// Wire freshly embedded content into the resolved slice; cache hits already
|
|
// carry their vector from the loop above.
|
|
for i, idx := range pairs {
|
|
resolved[idx] = &resolvedVec{content: contentResults[i].Vector, trunc: truncs[i]}
|
|
}
|
|
for i, re := range resolved {
|
|
if re == nil {
|
|
continue
|
|
}
|
|
merged := append([]float64(nil), re.content...)
|
|
if hasTitleVec {
|
|
merged, err = mergeEmbeddingVectors(titleVec, re.content, titleWeight)
|
|
if err != nil {
|
|
return nil, 0, fmt.Errorf("tokenizer: merge vectors: %w", err)
|
|
}
|
|
}
|
|
if err := chunks[i].SetExtraValue(fmt.Sprintf("q_%d_vec", len(merged)), merged); err != nil {
|
|
return nil, 0, fmt.Errorf("tokenizer: vector marshal: %w", err)
|
|
}
|
|
// Backfill the per-chunk cache only for freshly embedded content; cache
|
|
// hits are left untouched. The key intentionally omits title weighting:
|
|
// the content vector is title-weight independent, so identical chunk
|
|
// content reuses across runs even when only the filename weight changes.
|
|
// The embedded text (trunc) is part of the key so a config/model change
|
|
// that alters the embedded input forces a fresh embed rather than serving
|
|
// a stale vector.
|
|
if !re.isHit {
|
|
if chunkID, ok := chunks[i].GetExtraString("id"); ok && embdID != "" && store != nil {
|
|
if b, merr := json.Marshal(re.content); merr == nil {
|
|
chunkcache.Set(ctx, store, chunkcache.Key("emb", embdID, chunkID, re.trunc), string(b))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return chunks, tokenCount, nil
|
|
}
|
|
|
|
// defaultEmbeddingTokenLimit is the safe fallback used when an embedder reports
|
|
// no token limit at all. It is deliberately conservative: overshooting a model's
|
|
// window is a rejected request, while undershooting only truncates.
|
|
const defaultEmbeddingTokenLimit = tokenizer.EmbeddingTokenLimitDefault
|
|
|
|
// truncateForEmbedding keeps the first `limit` cl100k tokens of text so it fits
|
|
// the embedding model's limit.
|
|
//
|
|
// This is the LEGACY path: it is used only when the embedder cannot trim for
|
|
// itself (i.e. it does not implement EmbedderTrimmer), such as a test stub. The
|
|
// production embedder counts with the model's own tokenizer instead, because
|
|
// cl100k alone is not the model's tokenizer and the difference is
|
|
// content-dependent.
|
|
//
|
|
// Two changes from the original implementation, both fixing the same incident:
|
|
//
|
|
// - the margin is proportional (tokenizer.EmbeddingTokenLimit) instead of a
|
|
// flat 10 tokens, which was 0.12% of an 8192-token window — an order of
|
|
// magnitude smaller than the disagreement between two tokenizers;
|
|
// - the fallback limit is the resolved window (declared value, then the
|
|
// catalog's context_length, then a conservative default) rather than a
|
|
// hard-coded 8192, which overshoots the window of every smaller model.
|
|
func truncateForEmbedding(text string, maxTokens int) string {
|
|
declared := tokenizer.ResolveEmbeddingMaxTokens(maxTokens, 0)
|
|
limit := tokenizer.EmbeddingTokenLimit(declared)
|
|
return tokenizer.TrimContentToTokenLimit(text, limit)
|
|
}
|
|
|
|
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)
|
|
}
|
|
// Keep only chunks that have retrievable content: a chunk is dropped only
|
|
// when canonical text and both media-context fields are empty.
|
|
// Context-bearing media chunks may intentionally have no display text, but
|
|
// their surrounding prose is still searchable and must survive to the
|
|
// tokenizer.
|
|
filtered := raw[:0]
|
|
for _, ck := range raw {
|
|
if strings.TrimSpace(ck.Text) == "" &&
|
|
strings.TrimSpace(ck.ContextAbove) == "" && strings.TrimSpace(ck.ContextBelow) == "" {
|
|
continue
|
|
}
|
|
filtered = append(filtered, ck)
|
|
}
|
|
return filtered
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// declaredTitleTokens returns the tokens of the title a document declares about itself, when its
|
|
// header block carries one ("title: Pragyan Ojha", "name: Lesley Manyathela", "fullname: ...").
|
|
//
|
|
// Why it exists: a corpus whose files are named by id ("93372.md") otherwise leaves title_tks
|
|
// holding nothing but the id, while the retriever weights title_tks on a par with content_ltks -
|
|
// so the declared title is a ranking signal that costs nothing to index. Label words, the word
|
|
// "wikipedia", English stopwords and pure numbers are dropped: they are noise that would match
|
|
// unrelated queries once the field is weighted.
|
|
//
|
|
// Only the first non-empty chunk is inspected, because the header block (when present at all)
|
|
// belongs to the head of the document.
|
|
//
|
|
// The extraction is script-agnostic. A word that is not pure ASCII — Chinese, Japanese, Korean,
|
|
// Cyrillic, Arabic, Greek, Hebrew, Thai, or Latin carrying diacritics — is kept whole and handed
|
|
// to the analyzer downstream, which is what segments it. Splitting such a word here, on ASCII
|
|
// boundaries, is the older behaviour and it dissolved the whole title into nothing.
|
|
// Tokens that ARE pure ASCII keep the older treatment: englishStopwords are dropped (the
|
|
// tokenizer keeps function words such as "the", "for" and "and" as-is — verified via the ES
|
|
// _analyze API on title_tks and via tokenizer.Tokenize — so they have to be filtered here), and
|
|
// so are the label words and pure numbers. The stopword list stays English-only on purpose: a
|
|
// multilingual one would have to be invented, and getting a language wrong silently deletes real
|
|
// content words from that language's titles.
|
|
// This replaced an earlier length heuristic (keep only tokens >= 5 chars) which also removed real
|
|
// content words such as "icc", "cup", "john", "star", "bros".
|
|
var englishStopwords = map[string]bool{
|
|
"a": true, "an": true, "and": true, "are": true, "as": true, "at": true, "be": true, "been": true,
|
|
"but": true, "by": true, "d": true, "for": true, "from": true, "had": true, "has": true, "have": true,
|
|
"he": true, "her": true, "him": true, "his": true, "i": true, "if": true, "in": true, "into": true,
|
|
"is": true, "it": true, "its": true, "ll": true, "m": true, "me": true, "my": true, "no": true,
|
|
"not": true, "of": true, "on": true, "or": true, "our": true, "re": true, "s": true, "she": true,
|
|
"so": true, "t": true, "that": true, "the": true, "their": true, "them": true, "then": true,
|
|
"there": true, "these": true, "they": true, "this": true, "to": true, "too": true, "up": true,
|
|
"us": true, "ve": true, "was": true, "we": true, "were": true, "what": true, "when": true,
|
|
"where": true, "which": true, "who": true, "will": true, "with": true, "y": true,
|
|
"you": true, "your": true,
|
|
}
|
|
|
|
// asciiOnly reports whether s is made of ASCII bytes alone. A pure-ASCII string is the only thing
|
|
// ASCII decoding produces, so a byte scan decides it.
|
|
func asciiOnly(s string) bool {
|
|
for i := 0; i < len(s); i++ {
|
|
if s[i] >= utf8.RuneSelf {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
// allDigits reports whether s is non-empty and made only of decimal digits, in any script.
|
|
func allDigits(s string) bool {
|
|
if s == "" {
|
|
return false
|
|
}
|
|
for _, r := range s {
|
|
if !unicode.IsDigit(r) {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
// trimWordEdges drops leading and trailing runes that are neither letters nor digits, so the
|
|
// punctuation a title tends to carry — 《三国演义》, «Война и мир», (1917), a trailing comma — does
|
|
// not reach the token stream.
|
|
func trimWordEdges(s string) string {
|
|
return strings.TrimFunc(s, func(r rune) bool {
|
|
return !unicode.IsLetter(r) && !unicode.IsDigit(r)
|
|
})
|
|
}
|
|
|
|
func declaredTitleTokens(chunks []schema.ChunkDoc) []string {
|
|
const (
|
|
headerLines = 20
|
|
maxTokens = 20
|
|
)
|
|
var values []string
|
|
for i := range chunks {
|
|
head := chunks[i].Text
|
|
if head == "" {
|
|
continue
|
|
}
|
|
lines := strings.Split(head, "\n")
|
|
if len(lines) > headerLines {
|
|
lines = lines[:headerLines]
|
|
}
|
|
for _, line := range lines {
|
|
l := strings.ToLower(strings.TrimSpace(line))
|
|
for _, key := range []string{"title:", "name:", "fullname:"} {
|
|
if strings.HasPrefix(l, key) {
|
|
values = append(values, l[len(key):])
|
|
}
|
|
}
|
|
}
|
|
break
|
|
}
|
|
if len(values) == 0 {
|
|
return nil
|
|
}
|
|
skip := map[string]bool{"title": true, "name": true, "fullname": true, "wikipedia": true}
|
|
seen := map[string]bool{}
|
|
var out []string
|
|
// keep records one accepted token and reports whether the cap is now reached.
|
|
keep := func(token string) bool {
|
|
if seen[token] {
|
|
return false
|
|
}
|
|
seen[token] = true
|
|
out = append(out, token)
|
|
return len(out) >= maxTokens
|
|
}
|
|
for _, v := range values {
|
|
// Split on whitespace, NOT on ASCII boundaries. A word in another script is a word, and
|
|
// the analyzer downstream segments it (RAGFlow's tokenizer handles CJK); the
|
|
// ASCII-boundary splitter dissolved such a word into fragments of length zero.
|
|
for _, word := range strings.Fields(v) {
|
|
word = trimWordEdges(word)
|
|
if word == "" {
|
|
continue
|
|
}
|
|
if !asciiOnly(word) {
|
|
// Pure numbers stay out of the title channel whichever script writes them: the
|
|
// corpus file stem is numeric, so a numeric title token matches every document.
|
|
if !allDigits(word) && keep(word) {
|
|
return out
|
|
}
|
|
continue
|
|
}
|
|
for _, w := range strings.FieldsFunc(word, func(r rune) bool {
|
|
return !(r >= 'a' && r <= 'z') && !(r >= '0' && r <= '9')
|
|
}) {
|
|
// Pure numbers stay out (years, ids, and the corpus file stem are all numeric
|
|
// noise for the title channel); everything else is kept unless it is an English
|
|
// stopword, a label or a duplicate.
|
|
if englishStopwords[w] || skip[w] || allDigits(w) {
|
|
continue
|
|
}
|
|
if keep(w) {
|
|
return out
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// 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 selects the analyzer language (Snowball stemmer, or diacritics
|
|
// folding without stemming for Slovak/Czech), 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]
|
|
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 != "" {
|
|
// Split keywords on the ENGLISH COMMA ONLY. The keyword_prompt
|
|
// contract specifies "delimited by ENGLISH COMMA", so CJK commas,
|
|
// semicolons and newlines stay part of the keyword rather than
|
|
// acting as separators. strings.Split also preserves empty
|
|
// elements, matching Python's "a,,b".split(",") == ["a","","b"].
|
|
// Each piece is then bounded to the ES keyword field byte limit and
|
|
// truncated at a UTF-8 character boundary.
|
|
rawParts := strings.Split(kw, ",")
|
|
kwdParts := make([]string, len(rawParts))
|
|
for i, part := range rawParts {
|
|
// Empty parts are intentionally preserved to match the contract
|
|
// tested by TestTokenizerComponent_ImportantKwd_PreservesEmptyElements.
|
|
kwdParts[i] = sanitizeKeywordTerm(part)
|
|
}
|
|
if err = ck.SetExtraValue("important_kwd", kwdParts); 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 := schema.ContextualText(*ck); 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(schema.ContextualText(ck))
|
|
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.
|
|
//
|
|
// go_intentional (A4): the kbID != "" guard is deliberate. Each dataset
|
|
// configures its own embedding model, so an empty kb_id (e.g. a canvas-debug
|
|
// dry run) must NOT fall back to the tenant's default embedding model — doing
|
|
// so would produce vectors a dataset cannot actually use at retrieval time.
|
|
// This is a deliberate, go_intentional divergence. Do NOT "align" this to a
|
|
// path that injects a default embedding.
|
|
func shouldHaveEmbedding(searchMethods []string, kbID string) bool {
|
|
return slices.Contains(searchMethods, "embedding") && kbID != ""
|
|
}
|
|
|
|
func validateTokenizerOutputs(chunks []schema.ChunkDoc, searchMethods, fields []string, kbID string) error {
|
|
needFullText := slices.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(schema.ContextualText(ck)) != ""
|
|
}
|
|
|
|
// requiresEmbeddingVector uses the same cleaning as the embed path
|
|
// (cleanEmbedText): a chunk whose cleaned text is empty was skipped from the
|
|
// embedding batch and therefore legitimately carries no vector.
|
|
func requiresEmbeddingVector(ck schema.ChunkDoc, fields []string) bool {
|
|
return cleanEmbedText(concatFields(ck, fields)) != ""
|
|
}
|
|
|
|
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 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(),
|
|
})
|
|
}
|