fix(ingestion): split Extractor call() into callRaw/callText/callStructured (#18038)

2026-08-13 04:13:35 +08:00 · 2026-08-10 17:18:24 +08:00
parent 3a2a5f87dc
commit ca9e62ee59
2 changed files with 200 additions and 64 deletions
--- a/internal/ingestion/component/extractor.go
+++ b/internal/ingestion/component/extractor.go
@@ -665,7 +665,7 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map
 		}

 		if len(in.chunks) == 0 {
-			ans, callErr := c.call(timeoutCtx, db, in, "")
+			ans, callErr := c.callText(timeoutCtx, db, in, "")
 			if callErr != nil {
 				return callErr
 			}
@@ -709,7 +709,7 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map
 				strings.Contains(in.systemPrompt, "{text}") || strings.Contains(in.systemPrompt, "{chunks}") {
 				callChunkText = ""
 			}
-			ans, callErr := c.call(timeoutCtx, db, callIn, callChunkText)
+			ans, callErr := c.callText(timeoutCtx, db, callIn, callChunkText)
 			if callErr != nil {
 				return fmt.Errorf("chunk %d: %w", i, callErr)
 			}
@@ -747,11 +747,10 @@ func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, i
 		prompt:       "Output: ",
 		temperature:  &kwTemp,
 	}
-	result, err := c.call(ctx, db, kwIn, "")
+	resultStr, err := c.callText(ctx, db, kwIn, "")
 	if err != nil {
 		return err
 	}
-	resultStr, _ := result.(string)
 	resultStr = cleanExtractionResult(resultStr)
 	if resultStr == "" {
 		return nil
@@ -783,11 +782,10 @@ func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB,
 		prompt:       "Output: ",
 		temperature:  &qTemp,
 	}
-	result, err := c.call(ctx, db, qIn, "")
+	resultStr, err := c.callText(ctx, db, qIn, "")
 	if err != nil {
 		return err
 	}
-	resultStr, _ := result.(string)
 	resultStr = cleanExtractionResult(resultStr)
 	if resultStr == "" {
 		return nil
@@ -931,31 +929,16 @@ func (c *ExtractorComponent) runEnableMetadata(ctx context.Context, db *gorm.DB,
 			prompt:       "Output: ",
 			temperature:  &metaTemp,
 		}
-		result, err := c.call(ctx, db, metaIn, "")
+		parsed, err = c.callStructured(ctx, db, metaIn, "")
 		if err != nil {
 			return err
 		}
-		// call JSON-parses a JSON object response into a map; plain-text
-		// (or </think>-prefixed) responses come back as a string. Handle
-		// both so a valid ```json reply is never silently dropped.
-		var parsedObj map[string]any
-		switch r := result.(type) {
-		case map[string]any:
-			parsedObj = r
-		case string:
-			s := cleanExtractionResult(r)
-			if s == "" {
-				return nil
-			}
-			var ok bool
-			parsedObj, ok = tryParseJSONObject(s)
-			if !ok || len(parsedObj) == 0 {
-				return nil
-			}
-		default:
+		if parsed == nil {
+			// Non-JSON or empty response — nothing to extract, not an error.
+			// Matches Python gen_metadata treating an unparseable reply as an
+			// empty result rather than failing the chunk.
 			return nil
 		}
-		parsed = parsedObj
 		setMetadataLLMCache(ctx, in.llmID, schemaStr, chunkText, parsed)
 	}
 	// Merge into the chunk metadata map, preserving existing keys.
@@ -1027,6 +1010,38 @@ func setMetadataLLMCache(ctx context.Context, llmID, schemaJSON, chunkText strin
 	client.Set(ctx, metadataLLMCacheKey(llmID, schemaJSON, chunkText), string(data), metadataLLMCacheTTL)
 }

+// toolCallRE matches a <tool_call>...</tool_call> block, mirroring Python's
+// `re.sub(r"<tool_call>.*?</tool_call>", "", txt, flags=re.DOTALL)` at
+// llm_service.py:461. Non-greedy so consecutive tool_call blocks are each
+// stripped.
+var toolCallRE = regexp.MustCompile(`(?s)<tool_call>.*?</tool_call>`)
+
+// cleanLLMText applies the two-step LLM-layer cleanup that Python performs
+// in LLMBundle.async_chat (llm_service.py:459-461) for every response,
+// before any parsing:
+//
+//  1. Reasoning content: when the response BEGINS with `<think>`, take
+//     everything after the LAST `</think>`. A `<think>` that appears after
+//     a non-empty prefix is treated as ordinary content and preserved, so a
+//     legitimate prefix is never stripped. Python's _remove_reasoning_content
+//     (llm_service.py:332-346) finds the first `<think>` anywhere; restricting
+//     to a leading `<think>` avoids deleting real content that merely contains
+//     a think-looking tag. A leading `<think>` with no closing `</think>` is
+//     kept unchanged.
+//  2. Tool-call blocks: remove `<tool_call>...</tool_call>` spans.
+//
+// The result is a plain string, ready for the caller to consume as text or
+// to parse explicitly.
+func cleanLLMText(s string) string {
+	if strings.HasPrefix(s, "<think>") {
+		if j := strings.LastIndex(s, "</think>"); j >= 0 {
+			s = s[j+len("</think>"):]
+		}
+	}
+	s = toolCallRE.ReplaceAllString(s, "")
+	return strings.TrimSpace(s)
+}
+
 // cleanExtractionResult strips `</think>` tags and rejects `**ERROR**` responses,
 // matching Python's keyword_extraction and question_proposal post-processing.
 func cleanExtractionResult(s string) string {
@@ -1097,11 +1112,24 @@ func isRetryableLLMError(err error) bool {
 	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
-// callers can rely on a structured value downstream.
-func (c *ExtractorComponent) call(ctx context.Context, db *gorm.DB, in extractorInputs, chunkText string) (any, error) {
+// callRaw dispatches one LLM chat call for the supplied chunk text (empty
+// string in the no-chunk fast path) and returns the raw response. It holds
+// the shared plumbing — driver/target resolution, message assembly,
+// temperature override, retry — but deliberately does NOT clean or parse the
+// response. Callers pick the view they need:
+//
+//   - callText wraps callRaw with the LLM-layer two-step cleanup (think +
+//     tool_call) and returns a plain string — the field-extraction path,
+//     matching Python _generate_async.
+//   - callStructured wraps callRaw with cleanup + explicit JSON parsing and
+//     returns a map — the metadata path, matching Python gen_metadata.
+//
+// Splitting this way restores Python's separation of concerns: the LLM layer
+// (async_chat) cleans, the caller decides whether to parse. A single
+// all-purpose call() that best-effort parses would hand field extraction a
+// map it does not want (silently dropped downstream) while leaving metadata
+// to re-parse — and could not add cleaning without breaking JSON structure.
+func (c *ExtractorComponent) callRaw(ctx context.Context, db *gorm.DB, in extractorInputs, chunkText string) (*extractorChatResponse, error) {
 	driver, modelName, apiKey, baseURL, err := resolveExtractorChatTarget(ctx, db, in.llmID)
 	if err != nil {
 		return nil, err
@@ -1131,20 +1159,45 @@ func (c *ExtractorComponent) call(ctx context.Context, db *gorm.DB, in extractor
 	}, isRetryableLLMError); err != nil {
 		return nil, err
 	}
-	raw := strings.TrimSpace(resp.Content)
-	if raw == "" {
-		// No response — emit empty string so downstream code
-		// can distinguish from "LLM errored" via the error
-		// path above.
-		return "", nil
+	if resp == nil {
+		// Defensive: a provider adapter returning (nil, nil) would otherwise
+		// panic on resp.Content below. Surface a diagnosable error instead.
+		return nil, fmt.Errorf("extractor: chat: nil response from invoker")
 	}
-	// Best-effort JSON parse: a JSON object response is the
-	// canonical structured-extraction shape. Other shapes are
-	// returned verbatim so the caller can decide.
-	if parsed, ok := tryParseJSONObject(raw); ok {
-		return parsed, nil
+	return resp, nil
+}
+
+// callText runs one chat call and returns the cleaned text response. Used by
+// the field-extraction path (empty-chunks fast path and per-chunk writes),
+// which must always produce a string — matching Python's _generate_async
+// (llm.py:374-377) returning a raw string with no JSON parsing.
+func (c *ExtractorComponent) callText(ctx context.Context, db *gorm.DB, in extractorInputs, chunkText string) (string, error) {
+	resp, err := c.callRaw(ctx, db, in, chunkText)
+	if err != nil {
+		return "", err
 	}
-	return raw, nil
+	return cleanLLMText(resp.Content), nil
+}
+
+// callStructured runs one chat call, cleans the response, and explicitly
+// parses it as a JSON object. Used by the metadata path, which needs a
+// structured value to merge into chunk metadata — matching Python
+// gen_metadata's json_repair.loads. A non-JSON or empty response returns
+// (nil, nil) so the caller treats it as "nothing extracted", not an error.
+func (c *ExtractorComponent) callStructured(ctx context.Context, db *gorm.DB, in extractorInputs, chunkText string) (map[string]any, error) {
+	resp, err := c.callRaw(ctx, db, in, chunkText)
+	if err != nil {
+		return nil, err
+	}
+	s := cleanLLMText(resp.Content)
+	if s == "" {
+		return nil, nil
+	}
+	parsed, ok := tryParseJSONObject(s)
+	if !ok {
+		return nil, nil
+	}
+	return parsed, nil
 }

 // resolveExtractorChatTarget resolves the llm_id into driver / model /
--- a/internal/ingestion/component/extractor_test.go
+++ b/internal/ingestion/component/extractor_test.go
@@ -286,10 +286,12 @@ func TestExtractorComponent_Invoke_UnknownProvider(t *testing.T) {
 	}
 }

-// 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) {
+// TestExtractorComponent_Invoke_KeepsJSONAsString verifies a JSON object
+// response from the LLM is written to the chunk's field_name value as a
+// plain string — matching Python's _generate_async, which returns the raw
+// string with no JSON parsing. (The Extractor does NOT parse field-extraction
+// results; only the metadata path, via callStructured, parses explicitly.)
+func TestExtractorComponent_Invoke_KeepsJSONAsString(t *testing.T) {
 	withStubChatInvoker(t,
 		stubResponse{Content: `{"answer": 42, "tags": ["a", "b"]}`},
 	)
@@ -305,24 +307,20 @@ func TestExtractorComponent_Invoke_ParsesJSON(t *testing.T) {
 		t.Fatalf("Invoke: %v", err)
 	}
 	chunks := out["chunks"].([]map[string]any)
-	got, ok := chunks[0]["extraction"].(map[string]any)
+	got, ok := chunks[0]["extraction"].(string)
 	if !ok {
-		t.Fatalf("extraction should be parsed JSON object, got %T", chunks[0]["extraction"])
+		t.Fatalf("extraction should be a string (no JSON parse on field extraction), 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))
+	if got != `{"answer": 42, "tags": ["a", "b"]}` {
+		t.Errorf("extraction = %q, want the raw JSON string", got)
 	}
 }

-// 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) {
+// TestExtractorComponent_Invoke_KeepsJSONStringInFence verifies a JSON
+// response wrapped in a markdown code fence is stored as the raw string —
+// the code fence is not stripped on the field-extraction path (Python's
+// _generate_async returns the raw text untouched).
+func TestExtractorComponent_Invoke_KeepsJSONStringInFence(t *testing.T) {
 	withStubChatInvoker(t,
 		stubResponse{Content: "```json\n{\"summary\": \"hello\"}\n```"},
 	)
@@ -336,12 +334,12 @@ func TestExtractorComponent_Invoke_ParsesJSONInFence(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
-	got, ok := out["chunks"].([]map[string]any)[0]["out"].(map[string]any)
+	got, ok := out["chunks"].([]map[string]any)[0]["out"].(string)
 	if !ok {
-		t.Fatalf("out should be parsed JSON object, got %T", out["chunks"].([]map[string]any)[0]["out"])
+		t.Fatalf("out should be a string, got %T", out["chunks"].([]map[string]any)[0]["out"])
 	}
-	if got["summary"] != "hello" {
-		t.Errorf("summary = %v, want hello", got["summary"])
+	if got != "```json\n{\"summary\": \"hello\"}\n```" {
+		t.Errorf("out = %q, want the raw fenced JSON string", got)
 	}
 }

@@ -1266,6 +1264,91 @@ func TestCleanExtractionResult_LastThinkTag(t *testing.T) {
 	}
 }

+// TestCleanLLMText verifies the two-step LLM-layer cleanup that Python
+// applies in LLMBundle.async_chat (llm_service.py:459-461): reasoning
+// content is stripped only when a leading <think> has a matching closing
+// </think> after it, and <tool_call>...</tool_call> blocks are removed.
+func TestCleanLLMText(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		want string
+	}{
+		{
+			name: "think block stripped",
+			in:   "<think>reasoning</think>the answer",
+			want: "the answer",
+		},
+		{
+			name: "close without open kept",
+			in:   "abc</think>def",
+			want: "abc</think>def", // no leading <think> → unchanged
+		},
+		{
+			name: "prefix before think kept",
+			in:   "prefix<think>reason</think>answer",
+			want: "prefix<think>reason</think>answer", // <think> not at start → content preserved
+		},
+		{
+			name: "open without close kept",
+			in:   "<think>unclosed",
+			want: "<think>unclosed", // no </think> → unchanged
+		},
+		{
+			name: "tool_call block removed",
+			in:   "before<tool_call>{\"name\":\"x\"}</tool_call>after",
+			want: "beforeafter",
+		},
+		{
+			name: "consecutive tool_call blocks",
+			in:   "a<tool_call>1</tool_call>b<tool_call>2</tool_call>c",
+			want: "abc",
+		},
+		{
+			name: "think then tool_call",
+			in:   "<think>r</think>out<tool_call>t</tool_call>end",
+			want: "outend",
+		},
+		{
+			name: "plain text",
+			in:   "  plain answer  ",
+			want: "plain answer",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := cleanLLMText(tt.in); got != tt.want {
+				t.Errorf("cleanLLMText(%q) = %q, want %q", tt.in, got, tt.want)
+			}
+		})
+	}
+}
+
+// TestExtractorComponent_callStructured verifies the metadata path parses a
+// JSON object response into a map, and returns (nil, nil) for a non-JSON or
+// empty response (nothing extracted, not an error).
+func TestExtractorComponent_callStructured(t *testing.T) {
+	withStubChatInvoker(t, stubResponse{Content: `{"a": 1}`})
+	c := &ExtractorComponent{}
+	got, err := c.callStructured(t.Context(), nil, extractorInputs{llmID: "m"}, "")
+	if err != nil {
+		t.Fatalf("callStructured: %v", err)
+	}
+	if got["a"].(float64) != 1 {
+		t.Errorf("parsed = %v, want map with a=1", got)
+	}
+
+	// Non-JSON response → (nil, nil), not an error.
+	withStubChatInvoker(t, stubResponse{Content: "this is not JSON"})
+	got, err = c.callStructured(t.Context(), nil, extractorInputs{llmID: "m"}, "")
+	if err != nil {
+		t.Fatalf("callStructured on non-JSON: %v", err)
+	}
+	if got != nil {
+		t.Errorf("non-JSON response should yield nil map, got %v", got)
+	}
+}
+
 // 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