fix(ingestion): improve Extractor/LLM robustness and Python->Go parity (#17470)

## 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 <noreply@codebuddy.ai>
This commit is contained in:
Jack
2026-07-28 19:22:18 +08:00
committed by GitHub
parent 7e1ab9741b
commit 76aaecc284
9 changed files with 963 additions and 65 deletions

View File

@@ -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
}

76
internal/common/retry.go Normal file
View File

@@ -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)
}

View File

@@ -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())
}
}

View File

@@ -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 (

View File

@@ -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 `</think>` 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, "</think>"); i >= 0 {
if i := strings.LastIndex(s, "</think>"); i >= 0 {
s = s[i+len("</think>"):]
}
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

View File

@@ -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, "</think>"); idx >= 0 {
if idx := strings.LastIndex(raw, "</think>"); idx >= 0 {
raw = strings.TrimSpace(raw[idx+len("</think>"):])
}
if strings.Contains(raw, "**ERROR**") {

View File

@@ -401,3 +401,51 @@ func TestJsonRepairExtract(t *testing.T) {
})
}
}
// TestParseTaggerResponse_LastThinkTag verifies that when the LLM
// response contains multiple </think> 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: `<think>reasoning</think>{"tag": 5}`,
wantKeys: []string{"tag"},
},
{
name: "nested think blocks",
raw: `<think>outer</think>mid<think>inner</think>{"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)
}
}
})
}
}

View File

@@ -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 </think> 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: "<think>reasoning</think>the answer",
want: "the answer",
},
{
name: "nested think blocks",
in: "<think>outer</think>mid<think>inner</think>final output",
want: "final output",
},
{
name: "no think tag",
in: "plain answer",
want: "plain answer",
},
{
name: "think tag without close",
in: "<think>unclosed",
want: "<think>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)
}
}

View File

@@ -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")
}