Files
ragflow/internal/ingestion/component/extractor_test.go
Jack 8669c469d5 fix(ingestion): align laws DSL with Python — heading fallback, colon-title, short-line filter, remove_toc, and image extension mapping (#17200)
## Summary

This PR aligns the Go ingestion pipeline's **Laws** DSL template with
the Python implementation by fixing heading-detection gaps, adds
image-extension support, refactors the **Extractor** component's LLM
resolution, hardens heading detection for CJK text, and makes the
Extractor accept the Python DSL prompt key names
(`sys_prompt`/`prompts`) alongside the Go names.

## Changes

### 1. Picture file-type detection (`internal/utility/file.go`)

Adds explicit mapping for common image extensions (png, jpg, jpeg, gif,
bmp, tiff, tif, webp, svg, ico, avif, heic, apng) → `FileTypeVISUAL`,
with regression tests.

### 2. Laws DSL heading-detection alignment
(`internal/ingestion/component/chunker/`)

Four fixes to `resolveTitleLevels`:

| Fix | What changed | Why |
|-----|-------------|-----|
| **DOCX `ck_type` fallback** | `ckType` field on `lineRecord`,
propagated from `ChunkDoc.CKType` in `recordsFromStructured`. When
`ck_type=="heading"`, assign `fallbackLevel`. | office_oxide extracts
DOCX heading metadata, but the info was lost before reaching the heading
detector. Word headings whose text doesn't match any regex (e.g.
"Introduction") were treated as body. |
| **`make_colon_as_title` promotion** | `isColonTitle()`: promotes lines
ending with `:`/`:` that have sentence-ending punctuation before the
colon and ≥32 runes between them. | Mirrors Python's
`make_colon_as_title` in `rag/nlp/__init__.py`. Triple guard prevents
false positives. |
| **Short/numeric line filter** | Lines with ≤1 rune or purely numeric
are pinned to body level. | Mirrors Python `tree_merge`'s filter of
`sections` where `len(...) <= 1` or `re.match(r"[0-9]+$", ...)`. |
| **PDF `remove_toc`** | `"remove_toc": true` added to the PDF parser
setup in `ingestion_pipeline_laws.json`. | The Go PDF parser already
supports TOC removal; the Book template already enables it. |

### 3. Extractor llm_id resolution
(`internal/ingestion/component/extractor.go`)

Refactored to handle both **bare tenant_model UUIDs** and **composite
model@provider** strings via the shared `resolveModelConfig`
(`dispatch_model.go`):

- **`resolveExtractorChatConfig`** — UUID path calls
`resolveModelConfigByID` directly (one DB hit); composite path goes
through `resolveModelConfig`. Added `isBareTenantModelID` pre-check for
clear errors when a UUID doesn't exist.
- **`resolveExtractorChatTarget`** — propagates resolution errors
instead of silently returning empty driver.
- **`Chat()`** — removed `driver = "dummy"` fallback. Missing driver is
now an explicit error.
- **Removed dead code**: `splitExtractorLLID`,
`findExtractorSoleActiveInstance`.

### 4. `InjectExtractorLLMID` — fallback when no user config
(`internal/common/parser_config.go`)

Injects the tenant's global default LLM into extractor components **only
when their `llm_id` is empty**. Preserves user-selected UUID or
model@provider values.

Priority: user-configured llm_id > tenant global default > error (no
silent dummy fallback).

### 5. `ResponseHeaderTimeout` increase
(`internal/entity/models/base_model.go`)

`ResponseHeaderTimeout` 60s → 120s in `NewDriverHTTPClient`. Reasoning
models with large extraction prompts can take longer than 60s to produce
the first response token.

### 6. CJK rune-aware heading detection
(`internal/ingestion/component/chunker/title.go`)

Two byte-vs-rune bugs that only manifest on CJK text:

| Fix | What changed | Why |
|-----|-------------|-----|
| **`isColonTitle` byte offset** | `body[lastPunct+1:]` →
`body[lastPunct+runeLen:]` via `utf8.DecodeRuneInString` |
`strings.LastIndexAny` returns a byte index; `+1` skips only 1 byte,
corrupting multi-byte CJK punctuation (e.g. `。` = 3 bytes) and inflating
the rune count past the 32-rune threshold → false-positive heading
promotion. |
| **Short-line filter byte count** | `len(text) <= 1` →
`utf8.RuneCountInString(text) <= 1` | Go `len` is UTF-8 bytes; a single
CJK char (3 bytes) passed the filter, but Python's `len` returns 1 →
mismatch. |

### 7. `extractor_tag.go` — log error when llm fails

When `resolveExtractorChatTarget` returned an error, `runAutoTags` will
log error.

### 8. Python DSL prompt-key compatibility
(`internal/ingestion/component/extractor.go`)

The Resume DSL template uses Python-side key names (`sys_prompt`,
`prompts`). `NewExtractorComponent` now accepts them as fallbacks
alongside the Go names:

- `system_prompt` (Go) ← `sys_prompt` (Python) as fallback
- `prompt` (Go string) ← `prompts` (Python array `[{"role","content"}]`,
takes `[0].content`) as fallback

Mirrors the alias pattern already in `internal/agent/component/llm.go`.
`resolveInputs` accepts per-call `sys_prompt` override too.

## Remaining gaps vs Python

| Gap | Scope | Impact |
|-----|-------|--------|
| **TOC removal for TXT/MD/HTML** | Python's `remove_contents_table`
works on all text formats; Go's `remove_toc` is PDF-only. | Low —
plain-text documents rarely contain structured TOCs. |
| **Regex pattern details** | Minor differences in quantifiers, missing
H5/H6 markdown patterns, missing 4-level numbering pattern. | Low — Go's
variants are stricter; DOCX headings are covered by `ck_type` fallback.
|

## Testing

- `TestHierarchyTitleChunker_CKTypeHeadingFallback` — DOCX `ck_type`
heading promotion
- `TestHierarchyTitleChunker_ColonTitlePromotion` /
`_ColonTitleShortLine_Negative` — colon-title promotion + guard
- `TestIsColonTitle_CJKEdgeCase` / `TestIsColonTitle_ASCII_NoRegression`
— CJK byte-offset fix + ASCII regression
- `TestHierarchyTitleChunker_ColonTitlePromotion_CJK_EdgeCase` — CJK
colon edge case through full pipeline
- `TestHierarchyTitleChunker_ShortSingleCJKLineFilter` — single CJK char
filtered to body
- `TestHierarchyTitleChunker_ShortNumericLineFilter` — purely numeric
lines filtered
- `TestGetFileType_ImageExtensions` / `_ExistingFormats_NoRegression` —
image extension mapping
- `TestInjectExtractorLLMID_SkipWhenUUID` / `_SkipWhenComposite` /
`_InjectWhenEmpty` — llm_id injection guard
- `TestIsBareTenantModelID` — UUID detection
- `TestResolveExtractorChatTarget_AtSplitFallback` / `_NoDriver` — @
split fallback without DB
- `TestNewExtractorComponent_SysPromptAlias` / `_PromptsArray` /
`_PromptsArray_PromptWins` / `_SystemPromptWinsOverSysPrompt` — Python
key compatibility
- `TestBuildDOCXJSONSections_List` / `_TextBox` / `_MixedWithList` —
DOCX list/text_box parsing
- Full ingestion test suite passes (chunker, pipeline, task, service,
component packages)
2026-07-22 19:14:32 +08:00

779 lines
25 KiB
Go

//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package component
import (
"context"
"errors"
"strings"
"sync"
"sync/atomic"
"testing"
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(context.Background(), 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.
func TestExtractorComponent_Invoke_LLMError(t *testing.T) {
withStubChatInvoker(t,
stubResponse{Err: errors.New("upstream llm unavailable")},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "summary",
LLMID: "gpt-4o-mini",
}}
_, err := c.Invoke(context.Background(), 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_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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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(context.Background(), 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) {
withStubChatInvoker(t,
stubResponse{Content: "ok for chunk 0"},
stubResponse{Err: errors.New("chunk-1-boom")},
)
c := &ExtractorComponent{Param: schema.ExtractorParam{
FieldName: "out",
}}
_, err := c.Invoke(context.Background(), 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(context.Background(), 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_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 {
ck := ck
wg.Add(1)
go func() {
defer wg.Done()
_, err := c.Invoke(context.Background(), 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) {
driver, modelName, apiKey, baseURL, err := resolveExtractorChatTarget(
context.Background(), "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) {
driver, modelName, _, _, err := resolveExtractorChatTarget(
context.Background(), "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)
}
}