Files
ragflow/internal/ingestion/component/extractor_test.go
Jack 76aaecc284 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>
2026-07-28 19:22:18 +08:00

1213 lines
40 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package component
import (
"context"
"errors"
"ragflow/internal/dao"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
eschema "github.com/cloudwego/eino/schema"
"ragflow/internal/agent/runtime"
"ragflow/internal/ingestion/component/schema"
)
// stubExtractorChatInvoker is the test seam for the package-level
// extractorChatInvoker. It records every call (for assertions) and
// returns canned responses configured per-test. Concurrent-safe so
// it can backstop concurrent test cases without rewriting.
type stubExtractorChatInvoker struct {
mu sync.Mutex
// responses is consumed in order; remaining entries are returned
// as the wrap-error. tests set entries == call count they expect.
responses []stubResponse
// lastReq records the most recent call's request for inspection
// (e.g. driver / model name resolved from the llm_id).
lastReq extractorChatRequest
calls atomic.Int32
}
// stubResponse couples a Content value and an Err. tests populate
// either field — Err takes precedence over Content when non-nil.
type stubResponse struct {
Content string
Err error
}
func (s *stubExtractorChatInvoker) Chat(_ context.Context, req extractorChatRequest) (*extractorChatResponse, error) {
s.calls.Add(1)
s.mu.Lock()
s.lastReq = req
var resp stubResponse
if len(s.responses) > 0 {
resp = s.responses[0]
s.responses = s.responses[1:]
}
s.mu.Unlock()
if resp.Err != nil {
return nil, resp.Err
}
return &extractorChatResponse{Content: resp.Content}, nil
}
func (s *stubExtractorChatInvoker) Calls() int32 { return s.calls.Load() }
// withStubChatInvoker installs a stub invoker for the duration of
// the test and restores the production invoker on cleanup.
func withStubChatInvoker(t *testing.T, responses ...stubResponse) *stubExtractorChatInvoker {
t.Helper()
prev := defaultExtractorChatInvoker
stub := &stubExtractorChatInvoker{responses: responses}
SetExtractorChatInvoker(stub)
t.Cleanup(func() {
SetExtractorChatInvoker(prev)
})
return stub
}
// TestExtractorComponent_Registered verifies the init() registration
// is visible to the runtime registry (Phase 4 / API layer
// depends on this).
func TestExtractorComponent_Registered(t *testing.T) {
factory, cat, md, ok := runtime.DefaultRegistry.Lookup("Extractor")
if !ok {
t.Fatal("Extractor not registered in runtime.DefaultRegistry")
}
if cat != runtime.CategoryIngestion {
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
}
if factory == nil {
t.Error("factory is nil")
}
if md.Inputs == nil || len(md.Inputs) == 0 {
t.Errorf("metadata.Inputs empty: %v", md.Inputs)
}
if md.Outputs == nil || len(md.Outputs) == 0 {
t.Errorf("metadata.Outputs empty: %v", md.Outputs)
}
if _, has := md.Outputs["chunks"]; !has {
t.Errorf("metadata.Outputs missing %q", "chunks")
}
if _, has := md.Outputs["output_format"]; !has {
t.Errorf("metadata.Outputs missing %q", "output_format")
}
}
// TestExtractorComponent_Invoke_HappyPath covers the per-chunk
// fan-out: two chunks in → two LLM calls → each chunk enriched
// with the field_name key.
func TestExtractorComponent_Invoke_HappyPath(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "answer for chunk 1"},
stubResponse{Content: "answer for chunk 2"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "summary",
LLMID: "gpt-4o-mini",
Prompt: "Summarize:",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{
{"text": "first text"},
{"text": "second text"},
},
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
chunks, ok := out["chunks"].([]map[string]any)
if !ok {
t.Fatalf("chunks key missing or wrong shape: %T", out["chunks"])
}
if len(chunks) != 2 {
t.Fatalf("chunks len = %d, want 2", len(chunks))
}
if chunks[0]["summary"] != "answer for chunk 1" {
t.Errorf("chunk[0].summary = %v, want %q", chunks[0]["summary"], "answer for chunk 1")
}
if chunks[1]["summary"] != "answer for chunk 2" {
t.Errorf("chunk[1].summary = %v, want %q", chunks[1]["summary"], "answer for chunk 2")
}
if out["output_format"] != "chunks" {
t.Errorf("output_format = %v, want chunks", out["output_format"])
}
}
// TestExtractorComponent_Invoke_LLMError verifies a mock LLM
// error is surfaced through Invoke with the component-name prefix
// 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: errSentinel},
stubResponse{Err: errSentinel},
stubResponse{Err: errSentinel},
stubResponse{Err: errSentinel},
)
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": "x"}},
})
if err == nil {
t.Fatal("Invoke returned nil error")
}
if !strings.HasPrefix(err.Error(), "extractor:") {
t.Errorf("error should be wrapped with 'extractor:' prefix, got %v", err)
}
if !strings.Contains(err.Error(), "upstream llm unavailable") {
t.Errorf("error should chain underlying error, got %v", err)
}
}
// 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;
// the Extractor is provider-agnostic via llm_id; the 8 missing
// are edge cases that do not block Phase 2.5").
//
// Design note: every other test in this file drives the
// invoker through the production Component.Invoke path with a
// canned-response invoker installed via SetExtractorChatInvoker
// (the test seam). That seam accepts a pre-resolved driver
// path; it cannot model the eino factory's default-branch
// behaviour for an unknown driver. This test exercises the
// production chat-invoker directly to pin that branch — the
// production code path the real Extractor will hit when the
// DSL references a provider that is not in the 48/56 covered
// set.
//
// The contract under test:
// - The call MUST NOT panic.
// - On unknown driver, the factory's default branch routes to
// a DummyModel that returns a deterministic error string
// (we assert the error contains that sentinel so future
// maintainers see the wiring goes through the factory,
// not bypassed by a hand-rolled default).
func TestExtractorComponent_Invoke_UnknownProvider(t *testing.T) {
inv := &einoExtractorChatInvoker{}
resp, err := inv.Chat(context.Background(), extractorChatRequest{
Driver: "definitely-not-a-real-provider-xyz",
ModelName: "anything",
})
// Either an error is returned OR a non-nil response is produced
// by the DummyModel fallback. The contract is "no panic"; both
// of these outcomes are acceptable. We only fail the test if
// BOTH error and response are empty (which would indicate a
// silent no-op).
if err == nil && resp == nil {
t.Fatal("production invoker returned nil error AND nil response for unknown driver — silent no-op")
}
// When an error IS returned, it must mention the driver name so
// operators can correlate the failure back to the DSL config.
if err != nil {
// Acceptable error patterns for an unknown driver:
// - mentions the driver name (correlatable for operators)
// - "no driver"/"unknown" sentinels (typed error)
// - "not implemented" (the eino dummy model fallback path)
if !strings.Contains(err.Error(), "definitely-not-a-real-provider-xyz") &&
!strings.Contains(err.Error(), "no driver") &&
!strings.Contains(err.Error(), "unknown") &&
!strings.Contains(err.Error(), "not implemented") {
t.Errorf("unknown-driver error should mention the driver name or a typed/typed-sentinel substring; got: %v", err)
}
}
}
// TestExtractorComponent_Invoke_ParsesJSON verifies a JSON object
// response from the LLM is parsed into the chunk's field_name
// value (matching the python set_output contract).
func TestExtractorComponent_Invoke_ParsesJSON(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: `{"answer": 42, "tags": ["a", "b"]}`},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "extraction",
Prompt: "extract:",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{{"text": "doc"}}},
)
if err != nil {
t.Fatalf("Invoke: %v", err)
}
chunks := out["chunks"].([]map[string]any)
got, ok := chunks[0]["extraction"].(map[string]any)
if !ok {
t.Fatalf("extraction should be parsed JSON object, got %T", chunks[0]["extraction"])
}
if got["answer"].(float64) != 42 {
t.Errorf("answer = %v, want 42", got["answer"])
}
tags, _ := got["tags"].([]any)
if len(tags) != 2 {
t.Errorf("tags len = %d, want 2", len(tags))
}
}
// TestExtractorComponent_Invoke_ParsesJSONInFence verifies the
// common LLM response shape — JSON wrapped in a markdown code
// fence — parses cleanly. Mirrors the behaviour the agent
// canvas applies (e.g. llm_retry_test.go matchOutputStructure).
func TestExtractorComponent_Invoke_ParsesJSONInFence(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "```json\n{\"summary\": \"hello\"}\n```"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "out",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{{"text": "x"}}},
)
if err != nil {
t.Fatalf("Invoke: %v", err)
}
got, ok := out["chunks"].([]map[string]any)[0]["out"].(map[string]any)
if !ok {
t.Fatalf("out should be parsed JSON object, got %T", out["chunks"].([]map[string]any)[0]["out"])
}
if got["summary"] != "hello" {
t.Errorf("summary = %v, want hello", got["summary"])
}
}
// TestExtractorComponent_Invoke_HandlesMalformedJSON verifies a
// non-JSON response surfaces as the raw string under the
// destination field — not an error. The python Extractor
// accepts whatever the LLM emits; downstream callers decide
// what to do with it.
func TestExtractorComponent_Invoke_HandlesMalformedJSON(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "this is not JSON at all"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "raw",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{{"text": "x"}}},
)
if err != nil {
t.Fatalf("Invoke returned error on non-JSON: %v", err)
}
got := out["chunks"].([]map[string]any)[0]["raw"]
if got != "this is not JSON at all" {
t.Errorf("raw = %v, want %q", got, "this is not JSON at all")
}
}
// TestExtractorComponent_Invoke_TOCNotPorted asserts the
// field_name=="toc" branch is gated by a clear error so a future
// migration to the Go TOC generator doesn't accidentally fall
// through to chunk iteration.
func TestExtractorComponent_Invoke_TOCNotPorted(t *testing.T) {
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "toc",
}}
_, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{{"text": "x"}}},
)
if err == nil {
t.Fatal("expected error for field_name=toc, got nil")
}
if !strings.Contains(err.Error(), "toc") {
t.Errorf("error should mention toc: %v", err)
}
if !strings.Contains(err.Error(), "not yet ported") {
t.Errorf("error should call out parity gap: %v", err)
}
}
// TestExtractorComponent_Invoke_NoChunksFastPath verifies the
// no-chunks input still produces a one-element chunks slice
// (mirrors python _invoke line 110 fallback).
func TestExtractorComponent_Invoke_NoChunksFastPath(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "single-shot answer"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "answer",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
chunks, ok := out["chunks"].([]map[string]any)
if !ok {
t.Fatalf("chunks missing or wrong shape")
}
if len(chunks) != 1 {
t.Fatalf("chunks len = %d, want 1", len(chunks))
}
if chunks[0]["answer"] != "single-shot answer" {
t.Errorf("answer = %v, want %q", chunks[0]["answer"], "single-shot answer")
}
}
func TestExtractorComponent_Invoke_JSONListInput(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "json chunk answer"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "answer",
}}
out, err := c.Invoke(t.Context(), nil, map[string]any{
"json": []map[string]any{{"text": "json payload chunk"}},
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
chunks, ok := out["chunks"].([]map[string]any)
if !ok || len(chunks) != 1 {
t.Fatalf("chunks malformed: %v", out["chunks"])
}
if chunks[0]["answer"] != "json chunk answer" {
t.Errorf("answer = %v, want %q", chunks[0]["answer"], "json chunk answer")
}
}
// TestExtractorComponent_Invoke_PerCallLLMIDOverride verifies an
// inputs["llm_id"] override wins over Param.LLMID and reaches
// the chat invoker verbatim (the per-call override is the
// explicit test seam for runtime reconfiguration).
func TestExtractorComponent_Invoke_PerCallLLMIDOverride(t *testing.T) {
stub := withStubChatInvoker(t,
stubResponse{Content: "ok"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "out",
LLMID: "static-llm",
}}
_, err := c.Invoke(t.Context(), nil, map[string]any{
"llm_id": "override-llm",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
stub.mu.Lock()
defer stub.mu.Unlock()
if stub.lastReq.ModelName != "override-llm" {
t.Errorf("ModelName = %q, want override-llm", stub.lastReq.ModelName)
}
}
// TestExtractorComponent_Invoke_CompositeLLMID verifies the
// composite "gpt-4o-mini@openai" form is split into driver and
// model before reaching the chat invoker. Matches the canonical
// composite llm_id convention used throughout the codebase
// (see internal/agent/component/llm_credentials.go:parseLLMIDParts).
func TestExtractorComponent_Invoke_CompositeLLMID(t *testing.T) {
stub := withStubChatInvoker(t,
stubResponse{Content: "ok"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "out",
LLMID: "gpt-4o-mini@openai",
}}
if _, err := c.Invoke(t.Context(), nil, map[string]any{}); err != nil {
t.Fatalf("Invoke: %v", err)
}
stub.mu.Lock()
defer stub.mu.Unlock()
if stub.lastReq.Driver != "openai" {
t.Errorf("Driver = %q, want openai", stub.lastReq.Driver)
}
if stub.lastReq.ModelName != "gpt-4o-mini" {
t.Errorf("ModelName = %q, want gpt-4o-mini", stub.lastReq.ModelName)
}
}
// TestExtractorComponent_Invoke_ChunkIndexInError verifies the
// error message includes the failing chunk index so a long
// 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: 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",
}}
_, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{
{"text": "first"},
{"text": "second"},
},
})
if err == nil {
t.Fatal("Invoke returned nil error")
}
if !strings.Contains(err.Error(), "chunk 1") {
t.Errorf("error should mention chunk 1 (zero-indexed): %v", err)
}
if !strings.Contains(err.Error(), "chunk-1-boom") {
t.Errorf("error should chain underlying error: %v", err)
}
}
// TestExtractorComponent_NewExtractorComponent_ParamCheck covers
// the construction-time Validate() rejection of an empty
// field_name (matches python check_empty "Result Destination").
func TestExtractorComponent_NewExtractorComponent_ParamCheck(t *testing.T) {
c, err := NewExtractorComponent(map[string]any{})
if err != nil {
t.Fatalf("expected nil error, got %v", err)
}
if c == nil {
t.Fatal("expected non-nil component")
}
}
// TestExtractorComponent_NewExtractorComponent_Happy covers the
// parse path of every supported key; the param block coming out
// should round-trip cleanly through Invoke.
func TestExtractorComponent_NewExtractorComponent_Happy(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
c, err := NewExtractorComponent(map[string]any{
"field_name": "summary",
"llm_id": "openai/gpt-4o-mini",
"system_prompt": "You are a precise summarizer.",
"prompt": "Summarize:",
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
if _, err = c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{{"text": "x"}}},
); err != nil {
t.Fatalf("Invoke: %v", err)
}
}
// TestNewExtractorComponent_SysPromptAlias verifies that "sys_prompt"
// (the Python DSL name) is accepted as a fallback for SystemPrompt.
func TestNewExtractorComponent_SysPromptAlias(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"sys_prompt": "You are a Python DSL prompt.",
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
if ec.Param.SystemPrompt != "You are a Python DSL prompt." {
t.Errorf("SystemPrompt = %q, want %q", ec.Param.SystemPrompt, "You are a Python DSL prompt.")
}
}
// TestNewExtractorComponent_PromptsArray verifies that the Python DSL
// "prompts" array format is parsed into Param.Prompt.
func TestNewExtractorComponent_PromptsArray(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"prompts": []any{
map[string]any{
"content": "Analyze: {TitleChunker:FlatMiceFix@chunks}",
"role": "user",
},
},
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
want := "Analyze: {TitleChunker:FlatMiceFix@chunks}"
if ec.Param.Prompt != want {
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, want)
}
}
// TestNewExtractorComponent_PromptsArray_PromptWins verifies that
// "prompt" (string) takes priority over "prompts" (array) when both
// are present in the DSL params.
func TestNewExtractorComponent_PromptsArray_PromptWins(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"prompt": "Direct prompt wins.",
"prompts": []any{
map[string]any{
"content": "Should be ignored.",
"role": "user",
},
},
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
if ec.Param.Prompt != "Direct prompt wins." {
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, "Direct prompt wins.")
}
}
// TestNewExtractorComponent_PromptsString verifies that a bare-string
// "prompts" (the shape emitted by the front-end graph.nodes form and
// the dsl/testdata templates) is normalized into Param.Prompt, mirroring
// Python agent/component/llm.py:119-120 which coerces a string prompts
// into [{"role":"user","content":prompts}]. Without this normalization
// the string form is silently dropped (the .([]any) assertion fails).
func TestNewExtractorComponent_PromptsString(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"prompts": "Content: {TitleChunker:FlatMiceFix@chunks}",
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
want := "Content: {TitleChunker:FlatMiceFix@chunks}"
if ec.Param.Prompt != want {
t.Errorf("Prompt = %q, want %q (string prompts should be normalized)", ec.Param.Prompt, want)
}
}
// TestNewExtractorComponent_PromptsString_PromptWins verifies that
// "prompt" (string) still takes priority over a string-form "prompts"
// when both are present, matching the prompt>prompts precedence of
// the list-form path (TestNewExtractorComponent_PromptsArray_PromptWins).
func TestNewExtractorComponent_PromptsString_PromptWins(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"prompt": "Direct prompt wins.",
"prompts": "Should be ignored.",
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
if ec.Param.Prompt != "Direct prompt wins." {
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, "Direct prompt wins.")
}
}
// TestNewExtractorComponent_SystemPromptWinsOverSysPrompt verifies
// that "system_prompt" takes priority over "sys_prompt".
func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) {
withStubChatInvoker(t, stubResponse{Content: "ok"})
comp, err := NewExtractorComponent(map[string]any{
"field_name": "out",
"system_prompt": "system_prompt wins.",
"sys_prompt": "sys_prompt ignored.",
})
if err != nil {
t.Fatalf("NewExtractorComponent: %v", err)
}
ec := comp.(*ExtractorComponent)
if ec.Param.SystemPrompt != "system_prompt wins." {
t.Errorf("SystemPrompt = %q, want %q", ec.Param.SystemPrompt, "system_prompt wins.")
}
}
// TestExtractorComponent_InputsOutputs_NonEmpty is the shape
// assertion Phase 4's API endpoint relies on.
func TestExtractorComponent_InputsOutputs_NonEmpty(t *testing.T) {
c := &ExtractorComponent{}
ins := c.Inputs()
outs := c.Outputs()
if len(ins) == 0 {
t.Error("Inputs() returned empty map")
}
if len(outs) == 0 {
t.Error("Outputs() returned empty map")
}
if _, ok := outs["chunks"]; !ok {
t.Errorf("Outputs() missing %q", "chunks")
}
if _, ok := outs["output_format"]; !ok {
t.Errorf("Outputs() missing %q", "output_format")
}
}
// TestSplitExtractorLLID covers the composite-id parser in
// isolation — keeps the matrix of edge cases at one call site
// so a regression is easy to attribute. The "@" separator is
// the canonical composite llm_id form used throughout the
// codebase (see internal/agent/component/llm_credentials.go).
func TestSplitExtractorLLID(t *testing.T) {
cases := []struct {
in string
wantModel string
wantProvider string
wantOK bool
}{
{"gpt-4o-mini@openai", "gpt-4o-mini", "openai", true},
{"bare-model", "bare-model", "", false},
{"trailing@", "trailing", "", true},
{"@leading", "", "leading", true},
{"", "", "", false},
}
for _, tc := range cases {
t.Run(tc.in, func(t *testing.T) {
model, provider, ok := splitExtractorLLIDPair(tc.in)
if ok != tc.wantOK {
t.Errorf("ok = %v, want %v", ok, tc.wantOK)
}
if model != tc.wantModel {
t.Errorf("model = %q, want %q", model, tc.wantModel)
}
if provider != tc.wantProvider {
t.Errorf("provider = %q, want %q", provider, tc.wantProvider)
}
})
}
}
// TestTryParseJSONObject covers the best-effort JSON parser
// independently of the LLM seam so its matrix of edge cases is
// easy to attribute.
func TestTryParseJSONObject(t *testing.T) {
cases := []struct {
name string
in string
wantOK bool
wantKey string // when wantOK=true, expected key in the parsed map
}{
{name: "object", in: `{"a":1}`, wantOK: true, wantKey: "a"},
{name: "object with fence", in: "```json\n{\"a\":1}\n```", wantOK: true, wantKey: "a"},
{name: "fence without json tag", in: "```\n{\"a\":1}\n```", wantOK: true, wantKey: "a"},
{name: "plain string", in: "hello", wantOK: false},
{name: "array", in: `[1,2]`, wantOK: false},
{name: "empty object", in: `{}`, wantOK: false},
{name: "empty", in: ``, wantOK: false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
parsed, ok := tryParseJSONObject(tc.in)
if ok != tc.wantOK {
t.Fatalf("ok = %v, want %v (got %v)", ok, tc.wantOK, parsed)
}
if ok && tc.wantKey != "" {
if _, has := parsed[tc.wantKey]; !has {
t.Errorf("parsed map missing %q: %v", tc.wantKey, parsed)
}
}
})
}
}
// TestExtractorComponent_ConcurrentInvoke verifies the chat
// invoker swap is safe under concurrent Invoke calls. This is
// the canary for SetExtractorChatInvoker and the package-level
// RWMutex contract — a data race here breaks race detector.
func TestExtractorComponent_ConcurrentInvoke(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Content: "1"},
stubResponse{Content: "2"},
stubResponse{Content: "3"},
stubResponse{Content: "4"},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "out",
}}
chunks := []map[string]any{
{"text": "a"}, {"text": "b"}, {"text": "c"}, {"text": "d"},
}
var wg sync.WaitGroup
errs := make(chan error, len(chunks))
for _, ck := range chunks {
wg.Add(1)
go func() {
defer wg.Done()
_, err := c.Invoke(t.Context(), nil, map[string]any{
"chunks": []map[string]any{ck},
})
if err != nil {
errs <- err
}
}()
}
wg.Wait()
close(errs)
for err := range errs {
t.Errorf("Invoke error under concurrency: %v", err)
}
}
// silence unused-import vet warnings for eschema in case the
// test file is built without the import ever being referenced
// (it currently isn't, but pinning the import keeps test-side
// imports honest if helpers move around in future revisions).
var _ = eschema.Message{}
// TestIsBareTenantModelID verifies UUID detection.
func TestIsBareTenantModelID(t *testing.T) {
tests := []struct {
input string
want bool
}{
{"9e819c2442b14f9dab46062916e29195", true},
{"ABCDEFabcdef01234567890123456789", true},
{"9e819c2442b14f9dab46062916e2919", false}, // 31 chars
{"9e819c2442b14f9dab46062916e29195X", false}, // 33 chars
{"gpt-4o-mini@openai", false},
{"", false},
{"not-a-uuid", false},
}
for _, tc := range tests {
got := isBareTenantModelID(tc.input)
if got != tc.want {
t.Errorf("isBareTenantModelID(%q) = %v, want %v", tc.input, got, tc.want)
}
}
}
// TestResolveExtractorChatTarget_AtSplitFallback verifies the @ split
// fallback path works without canvas state (unit test compatibility).
func TestResolveExtractorChatTarget_AtSplitFallback(t *testing.T) {
ctx := t.Context()
driver, modelName, apiKey, baseURL, err := resolveExtractorChatTarget(
ctx, dao.DB, "gpt-4o-mini@openai")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if driver != "openai" {
t.Errorf("driver = %q, want openai", driver)
}
if modelName != "gpt-4o-mini" {
t.Errorf("modelName = %q, want gpt-4o-mini", modelName)
}
if apiKey != "" || baseURL != "" {
t.Errorf("apiKey/baseURL should be empty in fallback path")
}
}
// TestResolveExtractorChatTarget_NoDriver verifies a non-@ plain string
// without canvas state returns no driver (passes through to Chat()).
func TestResolveExtractorChatTarget_NoDriver(t *testing.T) {
ctx := t.Context()
driver, modelName, _, _, err := resolveExtractorChatTarget(
ctx, dao.DB, "plain-name")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if driver != "" {
t.Errorf("driver should be empty for plain name, got %q", driver)
}
if modelName != "plain-name" {
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)
}
}