From 76aaecc28448e4f7c72a719997a6bca78c1ab1ca Mon Sep 17 00:00:00 2001 From: Jack Date: Tue, 28 Jul 2026 19:22:18 +0800 Subject: [PATCH] fix(ingestion): improve Extractor/LLM robustness and Python->Go parity (#17470) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR hardens the Go ingestion **Extractor** and **LLM retry** paths and closes several Python->Go parity gaps in the keyword/question/tag extraction flow. - **Generic retry utility** (`internal/common/retry.go`): `RetryWithBackoff` with exponential backoff (default 3 retries, 2s initial delay, capped at 1m), context-aware sleep, and a `maxRetries<=0` fast path. Covered by `internal/common/retry_test.go`. - **LLM retry reuse**: `agent/component/llm_retry.go` now delegates to `common.RetryWithBackoff` instead of an inline loop (behavior preserved: ctx cancellation short-circuits the backoff). - **Extractor LLM calls** (`extractor.go`): - `call()` now retries transient LLM failures via `RetryWithBackoff` (retry exhaustion fails the chunk instead of silently skipping). - Sets `temperature = 0.2`, matching Python `generator.py:230,245`. - Runs keyword and question extraction **concurrently** per chunk when both are enabled (`task_executor.py:444-448`), with mutex-guarded map writes to avoid data races. - Substitutes `{field_name}` placeholders (including `{chunks}` -> chunk text) in `prompt`/`system_prompt` before the call, mirroring Python `string_format` (`extractor.py:102-103`); unmatched placeholders are left as-is. - Falls back to the **tenant default chat model** when `llm_id` is empty (`task_executor.py:573-574`). - Strips `` **greedily** (`strings.LastIndex`) in `cleanExtractionResult`. - **Auto-tagging** (`extractor_tag.go`): drops the `in.llmID != ""` guards so an empty `llm_id` no longer skips tagging (uses the tenant default model), and strips `` greedily in `parseTaggerResponse`. - **Docs**: fixes a misleading `PresentationChunker` docstring that claimed per-slide `image`/`position` output (the PPTX path emits none — unlike PDF), and removes a stale `docs/migration_python_go_diff.md` reference in `media_dispatch.go`. ## Test plan - `bash build.sh --test ./internal/common/...` — passes (new retry utility + tests). - `bash build.sh --test ./internal/ingestion/component/...` — passes (extractor/chunker/schema). - `gofmt` and lefthook pre-commit checks pass. Note: the personal `docs/migration_python_go_diff.md` working notebook in the tree is intentionally **not** part of this PR. --------- Co-authored-by: CodeBuddy --- internal/agent/component/llm_retry.go | 42 +- internal/common/retry.go | 76 ++++ internal/common/retry_test.go | 157 +++++++ .../component/chunker/presentation.go | 6 +- internal/ingestion/component/extractor.go | 294 +++++++++++-- internal/ingestion/component/extractor_tag.go | 6 +- .../ingestion/component/extractor_tag_test.go | 48 +++ .../ingestion/component/extractor_test.go | 396 +++++++++++++++++- .../ingestion/component/media_dispatch.go | 3 +- 9 files changed, 963 insertions(+), 65 deletions(-) create mode 100644 internal/common/retry.go create mode 100644 internal/common/retry_test.go diff --git a/internal/agent/component/llm_retry.go b/internal/agent/component/llm_retry.go index b654f111f9..1b3b3fe651 100644 --- a/internal/agent/component/llm_retry.go +++ b/internal/agent/component/llm_retry.go @@ -26,6 +26,8 @@ import ( "fmt" "time" + "ragflow/internal/common" + "gorm.io/gorm" ) @@ -119,34 +121,16 @@ func (r *retryInvoker) Invoke(ctx context.Context, db *gorm.DB, req ChatInvokeRe if r.inner == nil { return nil, fmt.Errorf("component: retryInvoker: nil inner") } - delay := r.initialDelay - var lastErr error - for attempt := 0; attempt <= r.maxRetries; attempt++ { - resp, err := r.inner.Invoke(ctx, db, req) - if err == nil { - return resp, nil - } - lastErr = err - if attempt == r.maxRetries { - break - } - // Honour ctx cancellation during backoff. A short-circuited - // sleep avoids hanging on shutdown when a long initialDelay - // would otherwise block the goroutine. - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(delay): - } - // Cap the doubling at a sane upper bound (1 minute). Without - // this a misconfigured initialDelay (e.g. 10s) plus 5 retries - // would sleep 10+20+40+80+160 = 310s before giving up. - if delay > 0 { - delay *= 2 - if delay > time.Minute { - delay = time.Minute - } - } + var resp *ChatInvokeResponse + err := common.RetryWithBackoff(ctx, r.maxRetries, r.initialDelay, func() error { + r, e := r.inner.Invoke(ctx, db, req) + resp = r + return e + }) + // On failure, return nil (not the last partial response) so callers + // that check err first never dereference a half-formed resp. + if err != nil { + return nil, err } - return nil, fmt.Errorf("component: LLM: chat failed after %d retries: %w", r.maxRetries, lastErr) + return resp, nil } diff --git a/internal/common/retry.go b/internal/common/retry.go new file mode 100644 index 0000000000..757cfe37fe --- /dev/null +++ b/internal/common/retry.go @@ -0,0 +1,76 @@ +// Package common — generic retry utility with exponential backoff. +// +// RetryWithBackoff calls fn up to maxRetries+1 times (one initial +// attempt + maxRetries retries), sleeping delay*2^attempt between +// failures (capped at 1 minute). The sleep honours ctx cancellation. +// +// maxRetries <= 0 yields a single attempt (no retries). +// initialDelay <= 0 results in no delay between retries. +// +// Usage: +// +// err := RetryWithBackoff(ctx, 3, 2*time.Second, func() error { +// return someFailableOperation() +// }) + +package common + +import ( + "context" + "fmt" + "time" +) + +const ( + // DefaultRetryMax is the default number of retries (3). + DefaultRetryMax = 3 + // DefaultRetryDelay is the initial backoff delay (2s). + DefaultRetryDelay = 2 * time.Second +) + +// RetryWithBackoff retries fn on error with exponential backoff. +// Returns nil on the first successful attempt. Returns the last +// error wrapped with the retry count when all attempts fail. +// +// An optional shouldRetry predicate lets callers abort on +// non-transient errors: when supplied and it returns false for a +// given error, RetryWithBackoff stops immediately and returns that +// error without further backoff. A nil predicate (or a nil function +// in the slice) retries on every error, preserving the original +// behavior. +func RetryWithBackoff(ctx context.Context, maxRetries int, initialDelay time.Duration, fn func() error, shouldRetry ...func(error) bool) error { + if maxRetries <= 0 { + return fn() + } + canRetry := func(err error) bool { + if len(shouldRetry) == 0 || shouldRetry[0] == nil { + return true + } + return shouldRetry[0](err) + } + delay := initialDelay + var lastErr error + for attempt := 0; attempt <= maxRetries; attempt++ { + err := fn() + if err == nil { + return nil + } + lastErr = err + if !canRetry(err) { + return err + } + if attempt == maxRetries { + break + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(delay): + } + delay *= 2 + if delay > time.Minute { + delay = time.Minute + } + } + return fmt.Errorf("failed after %d retries: %w", maxRetries, lastErr) +} diff --git a/internal/common/retry_test.go b/internal/common/retry_test.go new file mode 100644 index 0000000000..8bd166ea26 --- /dev/null +++ b/internal/common/retry_test.go @@ -0,0 +1,157 @@ +package common + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestRetryWithBackoff_SuccessOnFirstAttempt(t *testing.T) { + var calls atomic.Int32 + err := RetryWithBackoff(t.Context(), 3, time.Millisecond, func() error { + calls.Add(1) + return nil + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if calls.Load() != 1 { + t.Errorf("calls = %d, want 1", calls.Load()) + } +} + +func TestRetryWithBackoff_SuccessOnThirdAttempt(t *testing.T) { + var calls atomic.Int32 + err := RetryWithBackoff(t.Context(), 3, time.Millisecond, func() error { + if calls.Add(1) < 3 { + return errors.New("transient") + } + return nil + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if calls.Load() != 3 { + t.Errorf("calls = %d, want 3", calls.Load()) + } +} + +func TestRetryWithBackoff_AllRetriesExhausted(t *testing.T) { + sentinel := errors.New("persistent failure") + err := RetryWithBackoff(t.Context(), 3, time.Millisecond, func() error { + return sentinel + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if !errors.Is(err, sentinel) { + t.Errorf("error should wrap sentinel, got %v", err) + } + if !strings.Contains(err.Error(), "after 3 retries") { + t.Errorf("error should mention retry count: %v", err) + } +} + +func TestRetryWithBackoff_ContextCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + var calls atomic.Int32 + errChan := make(chan error, 1) + go func() { + errChan <- RetryWithBackoff(ctx, 3, 10*time.Second, func() error { + calls.Add(1) + cancel() // cancel before retry sleep + return errors.New("fail") + }) + }() + err := <-errChan + if !errors.Is(err, context.Canceled) { + t.Errorf("expected context.Canceled, got %v", err) + } + if calls.Load() != 1 { + t.Errorf("calls = %d, want 1 (no retry after cancel)", calls.Load()) + } +} + +func TestRetryWithBackoff_ContextDeadlineExceeded(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Millisecond) + defer cancel() + var calls atomic.Int32 + err := RetryWithBackoff(ctx, 3, time.Second, func() error { + calls.Add(1) + return errors.New("fail") + }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected DeadlineExceeded, got %v", err) + } + // Should have made at least 1 call; ctx deadline stops further retries + if calls.Load() < 1 { + t.Errorf("expected at least 1 call, got %d", calls.Load()) + } +} + +func TestRetryWithBackoff_ZeroMaxRetries(t *testing.T) { + var calls atomic.Int32 + err := RetryWithBackoff(t.Context(), 0, 0, func() error { + calls.Add(1) + return errors.New("fail") + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if calls.Load() != 1 { + t.Errorf("calls = %d, want 1 (single attempt)", calls.Load()) + } +} + +func TestRetryWithBackoff_NegativeMaxRetries(t *testing.T) { + var calls atomic.Int32 + err := RetryWithBackoff(t.Context(), -1, 0, func() error { + calls.Add(1) + return errors.New("fail") + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if calls.Load() != 1 { + t.Errorf("calls = %d, want 1 (negative → single attempt)", calls.Load()) + } +} + +// TestRetryWithBackoff_NonRetryableStops verifies that when a +// shouldRetry predicate returns false, RetryWithBackoff aborts +// immediately instead of burning the remaining retries/backoff. +func TestRetryWithBackoff_NonRetryableStops(t *testing.T) { + var calls atomic.Int32 + sentinel := errors.New("auth failure") + err := RetryWithBackoff(t.Context(), 3, 10*time.Second, func() error { + calls.Add(1) + return sentinel + }, func(err error) bool { + return !errors.Is(err, sentinel) + }) + if !errors.Is(err, sentinel) { + t.Errorf("expected sentinel, got %v", err) + } + if calls.Load() != 1 { + t.Errorf("calls = %d, want 1 (non-retryable → single attempt)", calls.Load()) + } +} + +// TestRetryWithBackoff_NilPredicateRetriesAll confirms that omitting +// the predicate preserves the original blind-retry behavior. +func TestRetryWithBackoff_NilPredicateRetriesAll(t *testing.T) { + var calls atomic.Int32 + err := RetryWithBackoff(t.Context(), 2, time.Millisecond, func() error { + calls.Add(1) + return errors.New("transient") + }) + if err == nil { + t.Fatal("expected error, got nil") + } + if calls.Load() != 3 { + t.Errorf("calls = %d, want 3 (all retries attempted)", calls.Load()) + } +} diff --git a/internal/ingestion/component/chunker/presentation.go b/internal/ingestion/component/chunker/presentation.go index 63b51113fb..78abeff1d4 100644 --- a/internal/ingestion/component/chunker/presentation.go +++ b/internal/ingestion/component/chunker/presentation.go @@ -22,8 +22,10 @@ // Unlike TokenChunker (which merges slides into a single chunk) or // OneChunker (which collapses many slides into one), PresentationChunker // keeps each slide as the unit of chunking. The upstream parser produces -// one structured record per slide with text, image, page_number, and -// position information; this chunker passes each through unchanged. +// one record per slide with text and slide_number; this chunker passes +// each through unchanged. Note that the PPTX/PPT path does not emit image +// or position information (unlike the PDF path), so slide chunks carry no +// image and no bbox-based preview positioning. package chunker import ( diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index f0d9ea28f3..6d3b6a2ee9 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -63,6 +63,7 @@ package component import ( "context" "encoding/json" + "errors" "fmt" "regexp" "sort" @@ -90,6 +91,19 @@ const componentNameExtractor = "Extractor" // is configured. const extractorTimeout = 600 * time.Second +// extractorRetryMax and extractorRetryDelay are package-level vars +// so tests can override them (extractorRetryDelay → time.Millisecond) +// to avoid multi-second retry sleeps. Production defaults match +// common.RetryWithBackoff defaults (3 retries, 2s initial delay). +var ( + extractorRetryMax = common.DefaultRetryMax + extractorRetryDelay = common.DefaultRetryDelay +) + +// extractorTemperature mirrors Python's 0.2 default for keyword +// and question extraction calls (generator.py:230,245). +const extractorTemperature = 0.2 + const ( autoKeywordPrompt = `## Role You are a text analyzer. @@ -425,6 +439,12 @@ type extractorInputs struct { prompt string lang string chunks []map[string]any + // temperature overrides the LLM temperature for this call. A + // nil value leaves the request's Temperature unset so the model + // (or the chat-model default) decides, matching Python's generic + // Extractor path. The keyword/question helpers set it to + // extractorTemperature (0.2) to mirror generator.py. + temperature *float64 } // resolveInputs overlays per-call inputs on top of the @@ -561,19 +581,67 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map text, _ = ck["text"].(string) } - if c.Param.AutoKeywords > 0 { - if err := c.runAutoKeywords(timeoutCtx, db, in, ck, text); err != nil { - return fmt.Errorf("chunk %d keywords: %w", i, err) + if c.Param.AutoKeywords > 0 && c.Param.AutoQuestions > 0 { + // Run keyword and question extraction concurrently + // per chunk (mirrors Python's ThreadPoolExecutor: + // task_executor.py:444-448). The shared chunk map is + // guarded by mu inside runAutoKeywords/runAutoQuestions, + // which also short-circuit when the key already exists. + var kwErr, qErr error + var mu sync.Mutex + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + if err := c.runAutoKeywords(timeoutCtx, db, in, ck, text, &mu); err != nil { + kwErr = err + } + }() + go func() { + defer wg.Done() + if err := c.runAutoQuestions(timeoutCtx, db, in, ck, text, &mu); err != nil { + qErr = err + } + }() + wg.Wait() + if kwErr != nil { + return fmt.Errorf("chunk %d keywords: %w", i, kwErr) } - } - if c.Param.AutoQuestions > 0 { - if err := c.runAutoQuestions(timeoutCtx, db, in, ck, text); err != nil { - return fmt.Errorf("chunk %d questions: %w", i, err) + if qErr != nil { + return fmt.Errorf("chunk %d questions: %w", i, qErr) + } + } else { + if c.Param.AutoKeywords > 0 { + if err := c.runAutoKeywords(timeoutCtx, db, in, ck, text, nil); err != nil { + return fmt.Errorf("chunk %d keywords: %w", i, err) + } + } + if c.Param.AutoQuestions > 0 { + if err := c.runAutoQuestions(timeoutCtx, db, in, ck, text, nil); err != nil { + return fmt.Errorf("chunk %d questions: %w", i, err) + } } } if in.fieldName != "" { - ans, callErr := c.call(timeoutCtx, db, in, text) + // Substitute {field_name} placeholders with current + // chunk field values, mirroring Python's + // string_format at extractor.py:103. + callIn := in + callIn.prompt = substituteChunkPlaceholders(in.prompt, ck, text) + callIn.systemPrompt = substituteChunkPlaceholders(in.systemPrompt, ck, text) + // buildExtractorMessages appends chunkText to the user + // message unconditionally. When the chunk text was already + // embedded via a {text}/{chunks} placeholder above, passing + // it again would duplicate the content in the final prompt. + // Suppress the automatic append in that case (the keyword/ + // question helper paths do the same by passing ""). + callChunkText := text + if strings.Contains(in.prompt, "{text}") || strings.Contains(in.prompt, "{chunks}") || + strings.Contains(in.systemPrompt, "{text}") || strings.Contains(in.systemPrompt, "{chunks}") { + callChunkText = "" + } + ans, callErr := c.call(timeoutCtx, db, callIn, callChunkText) if callErr != nil { return fmt.Errorf("chunk %d: %w", i, callErr) } @@ -594,14 +662,30 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map }, nil } -func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string) error { - if _, exists := ck["important_kwd"]; exists { +// runAutoKeywords extracts keywords for the current chunk and stores +// them on ck["important_kwd"]. mu may be nil (sequential path); when +// non-nil it serializes the shared chunk-map accesses so concurrent +// keyword/question goroutines stay race-free. The existence check and +// the map writes are both guarded by mu. Keyword extraction pins +// temperature to extractorTemperature (0.2) to mirror generator.py. +func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string, mu *sync.Mutex) error { + if mu != nil { + mu.Lock() + } + _, exists := ck["important_kwd"] + if mu != nil { + mu.Unlock() + } + if exists { return nil } - kwIn := in - kwIn.prompt = "Output: " - kwIn.systemPrompt = fmt.Sprintf(autoKeywordPrompt, c.Param.AutoKeywords, chunkText) - kwIn.fieldName = "" + kwTemp := extractorTemperature + kwIn := extractorInputs{ + llmID: in.llmID, + systemPrompt: fmt.Sprintf(autoKeywordPrompt, c.Param.AutoKeywords, chunkText), + prompt: "Output: ", + temperature: &kwTemp, + } result, err := c.call(ctx, db, kwIn, "") if err != nil { return err @@ -615,23 +699,42 @@ func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, i if len(kwds) == 0 { return nil } - ck["important_kwd"] = kwds tok := tokenizer.New(in.lang) tks, tkErr := tok.Tokenize(strings.Join(kwds, " ")) + if mu != nil { + mu.Lock() + } + ck["important_kwd"] = kwds if tkErr == nil { ck["important_tks"] = tks } + if mu != nil { + mu.Unlock() + } return nil } -func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string) error { - if _, exists := ck["question_kwd"]; exists { +// runAutoQuestions extracts questions for the current chunk and stores +// them on ck["question_kwd"]. See runAutoKeywords for the mu contract +// and the temperature pin. +func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string, mu *sync.Mutex) error { + if mu != nil { + mu.Lock() + } + _, exists := ck["question_kwd"] + if mu != nil { + mu.Unlock() + } + if exists { return nil } - qIn := in - qIn.prompt = "Output: " - qIn.systemPrompt = fmt.Sprintf(autoQuestionPrompt, c.Param.AutoQuestions, chunkText) - qIn.fieldName = "" + qTemp := extractorTemperature + qIn := extractorInputs{ + llmID: in.llmID, + systemPrompt: fmt.Sprintf(autoQuestionPrompt, c.Param.AutoQuestions, chunkText), + prompt: "Output: ", + temperature: &qTemp, + } result, err := c.call(ctx, db, qIn, "") if err != nil { return err @@ -653,19 +756,25 @@ func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, if len(filtered) == 0 { return nil } - ck["question_kwd"] = filtered tok := tokenizer.New(in.lang) tks, tkErr := tok.Tokenize(strings.Join(filtered, "\n")) + if mu != nil { + mu.Lock() + } + ck["question_kwd"] = filtered if tkErr == nil { ck["question_tks"] = tks } + if mu != nil { + mu.Unlock() + } return nil } // cleanExtractionResult strips `` tags and rejects `**ERROR**` responses, // matching Python's keyword_extraction and question_proposal post-processing. func cleanExtractionResult(s string) string { - if i := strings.Index(s, ""); i >= 0 { + if i := strings.LastIndex(s, ""); i >= 0 { s = s[i+len(""):] } s = strings.TrimSpace(s) @@ -690,6 +799,48 @@ func splitKeywords(s string) []string { return result } +// nonRetryableStatusRE matches HTTP client-error status codes that +// signal a permanent (non-transient) condition and therefore must +// NOT be retried. The word boundaries are essential: a bare +// substring check on "400" would wrongly flag phrasing such as +// "context deadline exceeded after 400ms" (a transient timeout, +// retryable) as non-retryable. \b ensures we only match a standalone +// 3-digit status token, so "400ms" / "4000" do not match. 429 and +// 5xx are deliberately absent — they stay retryable. +var nonRetryableStatusRE = regexp.MustCompile(`\b(?:400|401|403|404|405|422)\b`) + +// isRetryableLLMError classifies an LLM chat error as worth +// retrying. The production chat invoker returns opaque errors: +// configuration failures (missing model/driver) before any API +// call, and the provider SDK's raw error after the call. We treat +// context cancellation/deadline as terminal, plus a lightweight +// heuristic for non-transient auth/client errors. Anything +// unrecognized defaults to retryable so genuinely transient 5xx / +// 429 / network blips keep retrying (matching the prior blind-retry +// behavior). +func isRetryableLLMError(err error) bool { + if err == nil { + return true + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + msg := strings.ToLower(err.Error()) + for _, s := range []string{ + "unauthorized", "authentication", "api key", + "bad request", "not found", "model not found", "content filter", + "no driver resolved", "model_name is required", "resolve driver", + } { + if strings.Contains(msg, s) { + return false + } + } + if nonRetryableStatusRE.MatchString(msg) { + return false + } + return true +} + // call dispatches one LLM chat call for the supplied chunk text // (empty string in the no-chunk fast path). The result is the // raw string from the model — JSON parsing happens here so @@ -701,14 +852,27 @@ func (c *ExtractorComponent) call(ctx context.Context, db *gorm.DB, in extractor } msgs := buildExtractorMessages(in.systemPrompt, in.prompt, chunkText, in.chunks) inv := getExtractorChatInvoker() - resp, err := inv.Chat(ctx, extractorChatRequest{ + req := extractorChatRequest{ Driver: driver, ModelName: modelName, APIKey: apiKey, BaseURL: baseURL, Messages: msgs, - }) - if err != nil { + } + // Only override the temperature when the caller set one. A nil + // Temperature lets the model / chat-model default decide, matching + // Python's generic Extractor path; keyword/question helpers set + // extractorTemperature (0.2) to mirror generator.py. + if in.temperature != nil { + temp := *in.temperature + req.Temperature = &temp + } + var resp *extractorChatResponse + if err := common.RetryWithBackoff(ctx, extractorRetryMax, extractorRetryDelay, func() error { + r, e := inv.Chat(ctx, req) + resp = r + return e + }, isRetryableLLMError); err != nil { return nil, err } raw := strings.TrimSpace(resp.Content) @@ -738,6 +902,14 @@ func resolveExtractorChatTarget(ctx context.Context, db *gorm.DB, llmID string) } } + // When llmID is empty, try tenant default chat model + // (mirrors Python task_executor.py:573-574 fallback). + if llmID == "" { + if cfg := resolveExtractorChatDefaultConfig(ctx, db); cfg.driver != "" { + return cfg.driver, cfg.modelName, cfg.apiKey, cfg.baseURL, nil + } + } + cfg, cfgErr := resolveExtractorChatConfig(ctx, db, llmID) if cfgErr != nil { return "", "", "", "", cfgErr @@ -822,6 +994,43 @@ func resolveExtractorChatConfig(ctx context.Context, db *gorm.DB, compositeLLMID }, nil } +// resolveExtractorChatDefaultConfig resolves the tenant's default chat +// model when no explicit llm_id is provided. Returns empty config when +// no canvas state or tenant_id is available (unit-test context). +func resolveExtractorChatDefaultConfig(ctx context.Context, db *gorm.DB) extractorChatConfig { + state, _, err := runtime.GetStateFromContext[*runtime.CanvasState](ctx) + if err != nil || state == nil { + return extractorChatConfig{} + } + tidVal, _ := state.GetGlobal("tenant_id") + tid, _ := tidVal.(string) + if tid == "" { + return extractorChatConfig{} + } + + driver, modelName, apiConfig, _, err := resolveTenantModelByType(ctx, db, tid, entity.ModelTypeChat) + if err != nil || driver == nil { + return extractorChatConfig{} + } + + apiKey := "" + baseURL := "" + if apiConfig != nil { + if apiConfig.ApiKey != nil { + apiKey = *apiConfig.ApiKey + } + if apiConfig.BaseURL != nil { + baseURL = *apiConfig.BaseURL + } + } + return extractorChatConfig{ + driver: strings.ToLower(driver.Name()), + modelName: modelName, + apiKey: apiKey, + baseURL: baseURL, + } +} + // isBareTenantModelID reports whether s is a 32-character hex string // (a tenant_model primary key), as opposed to a composite "model@provider". func isBareTenantModelID(s string) bool { @@ -930,6 +1139,39 @@ func substitutePromptPlaceholders(prompt string, chunks []map[string]any) string // placeholders (a future per-component substitution extends here). var placeholderRE = regexp.MustCompile(`\{[A-Za-z0-9_]+:[A-Za-z0-9_]+@chunks\}`) +// simplePlaceholderRE matches simple {field_name} placeholders +// (no colons, no @chunks). Used by substituteChunkPlaceholders to +// replace {text}, {content_with_weight}, etc. with the current +// chunk's field values, mirroring Python's string_format +// (agent/component/base.py:602-609). +var simplePlaceholderRE = regexp.MustCompile(`\{[A-Za-z_][A-Za-z0-9_]*\}`) + +// substituteChunkPlaceholders replaces {field_name} placeholders in +// the prompt with values from the current chunk map. The special +// alias "chunks" maps to chunkText (the current chunk's primary +// text), matching Python's `args[chunks_key] = ck["text"]` at +// extractor.py:102. Unmatched placeholders are left as-is. +func substituteChunkPlaceholders(prompt string, ck map[string]any, chunkText string) string { + if prompt == "" || ck == nil { + return prompt + } + // Build lookup: chunk fields + "chunks" alias + lookup := make(map[string]string, len(ck)+1) + for k, v := range ck { + lookup[k] = fmt.Sprintf("%v", v) + } + if _, has := lookup["chunks"]; !has { + lookup["chunks"] = chunkText + } + return simplePlaceholderRE.ReplaceAllStringFunc(prompt, func(match string) string { + key := match[1 : len(match)-1] // strip { } + if val, ok := lookup[key]; ok { + return val + } + return match // leave unknown placeholders as-is + }) +} + // tryParseJSONObject tries to parse s as a JSON object. Returns // (parsed, true) on success; (nil, false) on parse error or when // s is not a JSON object. Trims common markdown code fences diff --git a/internal/ingestion/component/extractor_tag.go b/internal/ingestion/component/extractor_tag.go index f4c8a14988..1b1a221000 100644 --- a/internal/ingestion/component/extractor_tag.go +++ b/internal/ingestion/component/extractor_tag.go @@ -166,12 +166,12 @@ func (c *ExtractorComponent) runAutoTags(ctx context.Context, db *gorm.DB, in ex matched := matchAndTagChunk(d, indexed.examples, indexed.tagTokens, indexed.allTags, topN) if matched != nil { examples = append(examples, *matched) - } else if in.llmID != "" { + } else { docsToTag = append(docsToTag, d) } } - if len(docsToTag) > 0 && in.llmID != "" { + if len(docsToTag) > 0 { driver, model, apiKey, baseURL, err := resolveExtractorChatTarget(ctx, db, in.llmID) if err != nil { common.Warn("extractor tag: resolve model failed, skipping LLM tagging", zap.Error(err)) @@ -728,7 +728,7 @@ func buildTaggerPrompt(topN int, tagSetStr string, examples []schema.TaggedChunk func parseTaggerResponse(raw string, topN int) map[string]int { raw = strings.TrimSpace(raw) - if idx := strings.Index(raw, ""); idx >= 0 { + if idx := strings.LastIndex(raw, ""); idx >= 0 { raw = strings.TrimSpace(raw[idx+len(""):]) } if strings.Contains(raw, "**ERROR**") { diff --git a/internal/ingestion/component/extractor_tag_test.go b/internal/ingestion/component/extractor_tag_test.go index 42c56e632a..6161281e1a 100644 --- a/internal/ingestion/component/extractor_tag_test.go +++ b/internal/ingestion/component/extractor_tag_test.go @@ -401,3 +401,51 @@ func TestJsonRepairExtract(t *testing.T) { }) } } + +// TestParseTaggerResponse_LastThinkTag verifies that when the LLM +// response contains multiple tags, parseTaggerResponse strips +// up to the LAST one (greedy, matching Python's re.sub), not just the +// first (which would leave a residual think block). +func TestParseTaggerResponse_LastThinkTag(t *testing.T) { + tests := []struct { + name string + raw string + wantKeys []string + }{ + { + name: "single think block", + raw: `reasoning{"tag": 5}`, + wantKeys: []string{"tag"}, + }, + { + name: "nested think blocks", + raw: `outermidinner{"real": 3}`, + wantKeys: []string{"real"}, + }, + { + name: "no think tag", + raw: `{"simple": 7}`, + wantKeys: []string{"simple"}, + }, + { + name: "error sentinel", + raw: "**ERROR**", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := parseTaggerResponse(tt.raw, 5) + if len(tt.wantKeys) == 0 { + if len(got) != 0 { + t.Errorf("expected empty result, got %v", got) + } + return + } + for _, k := range tt.wantKeys { + if _, ok := got[k]; !ok { + t.Errorf("expected key %q in result, got %v", k, got) + } + } + }) + } +} diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index 85b4cc347a..cff5bb206b 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -24,6 +24,7 @@ import ( "sync" "sync/atomic" "testing" + "time" eschema "github.com/cloudwego/eino/schema" @@ -158,10 +159,22 @@ func TestExtractorComponent_Invoke_HappyPath(t *testing.T) { // TestExtractorComponent_Invoke_LLMError verifies a mock LLM // error is surfaced through Invoke with the component-name prefix -// so the upstream pipeline can attribute failures. +// so the upstream pipeline can attribute failures. After retry +// (RetryWithBackoff: 3 retries), the error chains the cause. func TestExtractorComponent_Invoke_LLMError(t *testing.T) { + // Fast retry for tests — avoid multi-second sleeps. + prevMax, prevDelay := extractorRetryMax, extractorRetryDelay + extractorRetryMax, extractorRetryDelay = 3, time.Millisecond + t.Cleanup(func() { + extractorRetryMax, extractorRetryDelay = prevMax, prevDelay + }) + + errSentinel := errors.New("upstream llm unavailable") withStubChatInvoker(t, - stubResponse{Err: errors.New("upstream llm unavailable")}, + stubResponse{Err: errSentinel}, + stubResponse{Err: errSentinel}, + stubResponse{Err: errSentinel}, + stubResponse{Err: errSentinel}, ) c := &ExtractorComponent{Param: schema.ExtractorParam{ @@ -182,6 +195,41 @@ func TestExtractorComponent_Invoke_LLMError(t *testing.T) { } } +// TestExtractorComponent_Invoke_RetrySucceeds verifies that a transient +// LLM error is retried (RetryWithBackoff), and the invocation succeeds +// once the LLM recovers. +func TestExtractorComponent_Invoke_RetrySucceeds(t *testing.T) { + prevMax, prevDelay := extractorRetryMax, extractorRetryDelay + extractorRetryMax, extractorRetryDelay = 3, time.Millisecond + t.Cleanup(func() { + extractorRetryMax, extractorRetryDelay = prevMax, prevDelay + }) + + stub := withStubChatInvoker(t, + stubResponse{Err: errors.New("transient")}, + stubResponse{Err: errors.New("transient")}, + stubResponse{Content: "recovered"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + FieldName: "summary", + LLMID: "gpt-4o-mini", + }} + out, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"text": "x"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if s, _ := chunks[0]["summary"].(string); s != "recovered" { + t.Errorf("summary = %q, want recovered", s) + } + if calls := stub.Calls(); calls != 3 { + t.Errorf("calls = %d, want 3 (2 transient + 1 success)", calls) + } +} + // TestExtractorComponent_Invoke_UnknownProvider asserts the // production (eino) chat invoker handles an unregistered driver // without panicking, per plan §8 Q1 ("48/56 providers covered; @@ -449,9 +497,19 @@ func TestExtractorComponent_Invoke_CompositeLLMID(t *testing.T) { // pipeline run surfaces which input document triggered the LLM // failure (mirrors python's per-chunk progress call at line 105). func TestExtractorComponent_Invoke_ChunkIndexInError(t *testing.T) { + prevMax, prevDelay := extractorRetryMax, extractorRetryDelay + extractorRetryMax, extractorRetryDelay = 3, time.Millisecond + t.Cleanup(func() { + extractorRetryMax, extractorRetryDelay = prevMax, prevDelay + }) + + errBoom := errors.New("chunk-1-boom") withStubChatInvoker(t, stubResponse{Content: "ok for chunk 0"}, - stubResponse{Err: errors.New("chunk-1-boom")}, + stubResponse{Err: errBoom}, // chunk 1: attempt 0 + stubResponse{Err: errBoom}, // attempt 1 + stubResponse{Err: errBoom}, // attempt 2 + stubResponse{Err: errBoom}, // attempt 3 (last retry) ) c := &ExtractorComponent{Param: schema.ExtractorParam{ FieldName: "out", @@ -820,3 +878,335 @@ func TestResolveExtractorChatTarget_NoDriver(t *testing.T) { t.Errorf("modelName = %q, want plain-name", modelName) } } + +// TestExtractorComponent_Invoke_TemperatureSet verifies the keyword +// extraction LLM chat call receives Temperature=0.2, matching Python's +// keyword_extraction and question_proposal defaults (generator.py:230,245). +// Field extraction intentionally runs on a separate call and uses the +// model default (see TestExtractorComponent_Invoke_FieldNameTemperatureDefault), +// so this test enables only AutoKeywords to assert the 0.2 pin directly. +func TestExtractorComponent_Invoke_TemperatureSet(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "keyword, extraction"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + LLMID: "gpt-4o-mini", + AutoKeywords: 3, + }} + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"text": "document content"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + stub.mu.Lock() + defer stub.mu.Unlock() + if stub.lastReq.Temperature == nil { + t.Fatal("Temperature is nil, want 0.2") + } + if *stub.lastReq.Temperature != 0.2 { + t.Errorf("Temperature = %v, want 0.2", *stub.lastReq.Temperature) + } + if stub.calls.Load() != 1 { + t.Errorf("expected exactly 1 LLM call (keyword), got %d", stub.calls.Load()) + } +} + +// TestExtractorComponent_Invoke_FieldNameTemperatureDefault verifies +// that the generic field-extraction path leaves Temperature unset +// (model/default), unlike keyword/question which pin 0.2 — matching +// Python's generic Extractor behavior. +func TestExtractorComponent_Invoke_FieldNameTemperatureDefault(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "extracted"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + FieldName: "summary", + LLMID: "gpt-4o-mini", + }} + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"text": "document content"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + stub.mu.Lock() + defer stub.mu.Unlock() + if stub.lastReq.Temperature != nil { + t.Errorf("Temperature = %v, want nil (field extraction uses model default)", *stub.lastReq.Temperature) + } +} + +// TestIsRetryableLLMError locks in the retry-classification heuristic, +// especially the word-boundary guard that prevents a transient timeout +// message ("...after 400ms") from being misclassified as a permanent +// HTTP 400 and dropped. +func TestIsRetryableLLMError(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + {name: "nil is retryable", err: nil, want: true}, + {name: "context canceled is terminal", err: context.Canceled, want: false}, + {name: "deadline exceeded is terminal", err: context.DeadlineExceeded, want: false}, + { + name: "wrapped deadline with 400ms must stay retryable", + err: errors.New("context deadline exceeded after 400ms"), + want: true, + }, + {name: "429 stays retryable", err: errors.New("429 Too Many Requests"), want: true}, + {name: "503 stays retryable", err: errors.New("503 Service Unavailable"), want: true}, + {name: "401 unauthorized is terminal", err: errors.New("HTTP 401 Unauthorized"), want: false}, + {name: "403 forbidden is terminal", err: errors.New("403 forbidden"), want: false}, + {name: "404 not found is terminal", err: errors.New("HTTP 404 Not Found"), want: false}, + {name: "405 method not allowed is terminal", err: errors.New("405 Method Not Allowed"), want: false}, + {name: "422 unprocessable is terminal", err: errors.New("422 Unprocessable Entity"), want: false}, + {name: "bad request is terminal", err: errors.New("400 Bad Request: malformed"), want: false}, + {name: "api key phrase is terminal", err: errors.New("invalid api key"), want: false}, + {name: "no driver phrase is terminal", err: errors.New("no driver resolved for llm_id"), want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isRetryableLLMError(tt.err); got != tt.want { + t.Errorf("isRetryableLLMError(%v) = %v, want %v", tt.err, got, tt.want) + } + }) + } +} + +// TestCleanExtractionResult_LastThinkTag verifies that when the LLM +// response contains multiple tags, cleanExtractionResult strips +// up to the LAST one (greedy, matching Python's re.sub), not just the +// first (which would leave a residual think block in the output). +func TestCleanExtractionResult_LastThinkTag(t *testing.T) { + tests := []struct { + name string + in string + want string + }{ + { + name: "single think block", + in: "reasoningthe answer", + want: "the answer", + }, + { + name: "nested think blocks", + in: "outermidinnerfinal output", + want: "final output", + }, + { + name: "no think tag", + in: "plain answer", + want: "plain answer", + }, + { + name: "think tag without close", + in: "unclosed", + want: "unclosed", + }, + { + name: "error sentinel", + in: "valid output**ERROR**extra", + want: "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := cleanExtractionResult(tt.in) + if got != tt.want { + t.Errorf("cleanExtractionResult(%q) = %q, want %q", tt.in, got, tt.want) + } + }) + } +} + +// TestExtractorComponent_Invoke_ConcurrentKeywordsAndQuestions verifies +// that when both auto_keywords and auto_questions are enabled, both +// LLM calls are dispatched per chunk and results land on the chunk +// (matching Python's ThreadPoolExecutor concurrency: task_executor.py:444-448). +func TestExtractorComponent_Invoke_ConcurrentKeywordsAndQuestions(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "alpha, beta"}, // chunk 0 keywords + stubResponse{Content: "what is it?\nwhy?"}, // chunk 0 questions + stubResponse{Content: "gamma, delta"}, // chunk 1 keywords + stubResponse{Content: "how?\nwhen?"}, // chunk 1 questions + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + LLMID: "gpt-4o-mini", + AutoKeywords: 2, + AutoQuestions: 2, + }} + out, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "first doc"}, + {"text": "second doc"}, + }, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + chunks, ok := out["chunks"].([]map[string]any) + if !ok || len(chunks) != 2 { + t.Fatalf("expected 2 chunks, got %v", out["chunks"]) + } + + // Both chunks should have keywords and questions populated. + for i, ck := range chunks { + kwds, hasKW := ck["important_kwd"].([]string) + if !hasKW || len(kwds) == 0 { + t.Errorf("chunk %d: missing important_kwd", i) + } + qs, hasQ := ck["question_kwd"].([]string) + if !hasQ || len(qs) == 0 { + t.Errorf("chunk %d: missing question_kwd", i) + } + } + + if calls := stub.Calls(); calls != 4 { + t.Errorf("expected 4 LLM calls (2 chunks × 2 types), got %d", calls) + } +} + +// TestResolveExtractorChatTarget_EmptyLLMID verifies that when llmID is +// empty, resolveExtractorChatTarget falls back to the tenant default chat +// model (via resolveTenantModelByType), matching Python's behavior +// (task_executor.py:573-574 never skips tagging on empty llm_id). +// When no canvas state is available (unit-test context), returns empty +// driver — callers like runAutoTags check driver!="" before using it. +func TestResolveExtractorChatTarget_EmptyLLMID(t *testing.T) { + // Without canvas state: empty llmID returns empty driver (no crash). + ctx := t.Context() + driver, modelName, _, _, err := resolveExtractorChatTarget(ctx, dao.DB, "") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // In test context without canvas state, neither tenant default nor @ split + // can resolve — driver ends up empty. Callers must handle this gracefully. + if driver != "" { + t.Logf("resolved empty llmID: driver=%q model=%q (tenant default might be available)", driver, modelName) + } + // Contract: no panic, no error for empty llmID. +} + +// TestExtractorComponent_Invoke_SubstitutesPlaceholders verifies that +// {field_name} placeholders in the user prompt are substituted with +// the current chunk's field values before the LLM call, matching +// Python's string_format (agent/component/base.py:602). +func TestExtractorComponent_Invoke_SubstitutesPlaceholders(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "substituted answer"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + FieldName: "summary", + Prompt: "Analyze: {text}", + LLMID: "gpt-4o-mini", + }} + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"text": "the document content"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + stub.mu.Lock() + defer stub.mu.Unlock() + var userContent string + for _, msg := range stub.lastReq.Messages { + if msg.Role == eschema.User { + userContent = msg.Content + } + } + if strings.Contains(userContent, "{text}") { + t.Errorf("prompt still contains literal {text}: %q", userContent) + } + if !strings.Contains(userContent, "the document content") { + t.Errorf("prompt missing chunk text: %q", userContent) + } + // Regression guard: when the prompt embeds {text}, the chunk text + // must appear exactly once — buildExtractorMessages must not append + // it a second time (placeholder duplication bug). + if n := strings.Count(userContent, "the document content"); n != 1 { + t.Errorf("chunk text appears %d times, want 1: %q", n, userContent) + } +} + +// TestExtractorComponent_Invoke_PlaceholderChunksAlias verifies that +// {chunks} (the Python DSL upstream key) is also substituted with +// the current chunk text. +func TestExtractorComponent_Invoke_PlaceholderChunksAlias(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "answer"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + FieldName: "out", + Prompt: "Content: {chunks}", + LLMID: "gpt-4o-mini", + }} + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"content_with_weight": "weighted doc"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + stub.mu.Lock() + defer stub.mu.Unlock() + var userContent string + for _, msg := range stub.lastReq.Messages { + if msg.Role == eschema.User { + userContent = msg.Content + } + } + if strings.Contains(userContent, "{chunks}") { + t.Errorf("prompt still contains literal {chunks}: %q", userContent) + } + if !strings.Contains(userContent, "weighted doc") { + t.Errorf("prompt missing chunk text: %q", userContent) + } + // Regression guard: {chunks} must not duplicate the chunk text. + if n := strings.Count(userContent, "weighted doc"); n != 1 { + t.Errorf("chunk text appears %d times, want 1: %q", n, userContent) + } +} + +// TestExtractorComponent_Invoke_AppendsChunkTextWhenNoPlaceholder verifies +// that when the prompt has no {text}/{chunks} placeholder, the chunk text is +// still automatically appended by buildExtractorMessages exactly once. +func TestExtractorComponent_Invoke_AppendsChunkTextWhenNoPlaceholder(t *testing.T) { + stub := withStubChatInvoker(t, + stubResponse{Content: "answer"}, + ) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + FieldName: "summary", + Prompt: "Summarize the above:", + LLMID: "gpt-4o-mini", + }} + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{{"text": "the document content"}}, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + stub.mu.Lock() + defer stub.mu.Unlock() + var userContent string + for _, msg := range stub.lastReq.Messages { + if msg.Role == eschema.User { + userContent = msg.Content + } + } + if n := strings.Count(userContent, "the document content"); n != 1 { + t.Errorf("chunk text appears %d times, want 1: %q", n, userContent) + } +} diff --git a/internal/ingestion/component/media_dispatch.go b/internal/ingestion/component/media_dispatch.go index b359776feb..4c032c2d3e 100644 --- a/internal/ingestion/component/media_dispatch.go +++ b/internal/ingestion/component/media_dispatch.go @@ -77,8 +77,7 @@ func maybeDispatchVideo( // video_url. Returning an explicit error is safer than silently producing // a description from the prompt text alone. The real implementation must // be provider-specific (OpenAI-compatible: frame extraction -> image_url; - // Gemini: raw-bytes inline_data; Qwen: file://) — see - // docs/migration_python_go_diff.md 2.7. + // Gemini: raw-bytes inline_data; Qwen: file://) return parserDispatchResult{}, true, fmt.Errorf("Parser: video parsing is not yet supported; underlying video analysis capability is pending") }