mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-29 04:08:12 +08:00
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:
@@ -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
76
internal/common/retry.go
Normal 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)
|
||||
}
|
||||
157
internal/common/retry_test.go
Normal file
157
internal/common/retry_test.go
Normal 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())
|
||||
}
|
||||
}
|
||||
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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**") {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user