mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-09-19 02:18:25 +08:00
fix(graph): recover JSON from malformed/truncated fences in extract_entity
Closes #1113 The graph extraction pipeline relies on a strict markdown fence regex to pull JSON out of LLM responses. In production the regex misses ~84% of real-world responses, in three ways: 1. The LLM hits max_tokens mid-output and produces no closing fence. 2. The opening fence is malformed or surrounded by prose, so the non-greedy regex fails to anchor a match. 3. The model returns raw JSON with stray backticks but no real fence. In every case extractContent fell through to returning the raw text, which then failed json.Unmarshal with errors like "invalid character '` + "`" + "' looking for beginning of value". This change keeps the existing happy path untouched and adds a conservative recovery step in the default branch of extractContent: - If an opening ``` is present, take everything after it, drop a likely language tag on the first line, cut at any trailing closing fence, and trim stray backticks/whitespace. - Otherwise, look for an outermost JSON object/array in the text using a small bracket-balanced scanner that respects string literals, so embedded {} or [] inside JSON strings don't confuse it. - Only fall back to the original raw-text behavior when neither strategy yields anything plausible. The recovery helpers (stripFencesAndExtract, extractJSONLike, isLikelyLanguageTag) are package-private and have table-driven tests in extract_entity_test.go covering the three failure patterns described in the issue, plus the previously-working fenced and bare-JSON shapes to guard against regressions.
This commit is contained in:
draix
@@ -545,11 +545,153 @@ func (f *Formater) extractContent(ctx context.Context, text string) string {
|
||||
return strings.TrimSpace(matches[0][2])
|
||||
|
||||
default:
|
||||
// Fallback strategies for cases where the fence regex fails to match.
|
||||
// This commonly happens when:
|
||||
// 1. The LLM output is truncated (no closing fence) — issue #1113 Pattern 3.
|
||||
// 2. The opening fence is malformed or surrounded by unexpected content,
|
||||
// so the non-greedy regex falls back to the raw text — issue #1113 Pattern 1.
|
||||
// Without these fallbacks, the raw text (including backticks) is passed to
|
||||
// json.Unmarshal and fails with `invalid character '`'`.
|
||||
if extracted := stripFencesAndExtract(text, f.formatType); extracted != "" {
|
||||
logger.Debugf(ctx, "no fence match, recovered content via fallback (%d bytes)", len(extracted))
|
||||
return extracted
|
||||
}
|
||||
logger.Warnf(ctx, "no match found")
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
}
|
||||
|
||||
// stripFencesAndExtract attempts to recover a parseable payload from an LLM
|
||||
// response when the strict fence regex fails. It handles three common cases:
|
||||
//
|
||||
// 1. Truncated responses with an opening ```lang fence but no closing fence
|
||||
// (LLM hit max_tokens mid-output).
|
||||
// 2. Responses where the JSON/YAML body is preceded or followed by prose
|
||||
// and the fences are present but malformed.
|
||||
// 3. Responses with no fences at all but a recognizable JSON object/array
|
||||
// embedded in surrounding text.
|
||||
//
|
||||
// It returns an empty string when no plausible payload can be recovered, so
|
||||
// callers can fall back to their own behavior.
|
||||
func stripFencesAndExtract(text string, format FormatType) string {
|
||||
trimmed := strings.TrimSpace(text)
|
||||
if trimmed == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Case 1: opening fence present (with or without language tag) but no
|
||||
// matching closing fence. Take everything after the first fence and
|
||||
// strip any trailing backticks.
|
||||
if idx := strings.Index(trimmed, "```"); idx >= 0 {
|
||||
rest := trimmed[idx+3:]
|
||||
// Drop optional language tag on the same line.
|
||||
if nl := strings.IndexByte(rest, '\n'); nl >= 0 {
|
||||
firstLine := strings.TrimSpace(rest[:nl])
|
||||
// A pure language tag is short and alphanumeric-ish.
|
||||
if firstLine == "" || isLikelyLanguageTag(firstLine) {
|
||||
rest = rest[nl+1:]
|
||||
}
|
||||
}
|
||||
// If there is a closing fence somewhere, cut at it.
|
||||
if end := strings.Index(rest, "```"); end >= 0 {
|
||||
rest = rest[:end]
|
||||
}
|
||||
rest = strings.TrimSpace(rest)
|
||||
rest = strings.Trim(rest, "`")
|
||||
rest = strings.TrimSpace(rest)
|
||||
if rest != "" {
|
||||
return rest
|
||||
}
|
||||
}
|
||||
|
||||
// Case 2: no usable fence found, but the payload may still contain a
|
||||
// JSON object/array. Extract the outermost {...} or [...] substring.
|
||||
if format == FormatTypeJSON {
|
||||
if extracted := extractJSONLike(trimmed); extracted != "" {
|
||||
return extracted
|
||||
}
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// isLikelyLanguageTag reports whether s looks like a markdown fence language
|
||||
// tag (e.g. "json", "yaml", "yml", "go"). It must be short and contain only
|
||||
// characters typical for a language identifier.
|
||||
func isLikelyLanguageTag(s string) bool {
|
||||
if s == "" || len(s) > 16 {
|
||||
return false
|
||||
}
|
||||
for _, r := range s {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z':
|
||||
case r >= 'A' && r <= 'Z':
|
||||
case r >= '0' && r <= '9':
|
||||
case r == '_' || r == '-' || r == '+':
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// extractJSONLike returns the outermost JSON object or array substring from s,
|
||||
// or an empty string if none is found. It picks whichever bracket type appears
|
||||
// first in the input, which mirrors what the LLM is most likely to have
|
||||
// produced. The returned slice is not validated as JSON; callers must still
|
||||
// json.Unmarshal it.
|
||||
func extractJSONLike(s string) string {
|
||||
objStart := strings.IndexByte(s, '{')
|
||||
arrStart := strings.IndexByte(s, '[')
|
||||
var open, closeCh byte
|
||||
var start int
|
||||
switch {
|
||||
case objStart < 0 && arrStart < 0:
|
||||
return ""
|
||||
case objStart < 0:
|
||||
open, closeCh, start = '[', ']', arrStart
|
||||
case arrStart < 0:
|
||||
open, closeCh, start = '{', '}', objStart
|
||||
case objStart < arrStart:
|
||||
open, closeCh, start = '{', '}', objStart
|
||||
default:
|
||||
open, closeCh, start = '[', ']', arrStart
|
||||
}
|
||||
// Find matching close, respecting string literals so braces/brackets
|
||||
// inside JSON strings don't unbalance the count.
|
||||
depth := 0
|
||||
inString := false
|
||||
escaped := false
|
||||
for i := start; i < len(s); i++ {
|
||||
c := s[i]
|
||||
if inString {
|
||||
if escaped {
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
switch c {
|
||||
case '\\':
|
||||
escaped = true
|
||||
case '"':
|
||||
inString = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch c {
|
||||
case '"':
|
||||
inString = true
|
||||
case open:
|
||||
depth++
|
||||
case closeCh:
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return strings.TrimSpace(s[start : i+1])
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (f *Formater) addFences(content string) string {
|
||||
content = strings.TrimSpace(content)
|
||||
return fmt.Sprintf("```%s\n%s\n```", f.formatType, content)
|
||||
|
||||
@@ -0,0 +1,261 @@
|
||||
package chatpipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestFormater_ParseGraph_FenceVariants exercises the JSON parsing path used
|
||||
// by the graph extraction pipeline against the LLM response shapes that
|
||||
// caused issue #1113. Each case feeds a raw LLM response string to
|
||||
// Formater.ParseGraph and asserts the resulting graph data, or that the
|
||||
// error path is preserved for genuinely invalid input.
|
||||
func TestFormater_ParseGraph_FenceVariants(t *testing.T) {
|
||||
const validJSON = `[
|
||||
{"entity": "Alice", "entity_attributes": ["person"]},
|
||||
{"entity": "Bob", "entity_attributes": ["person"]},
|
||||
{"entity1": "Alice", "entity2": "Bob", "relation": "knows"}
|
||||
]`
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
input string
|
||||
wantNodes int
|
||||
wantRels int
|
||||
wantErr bool
|
||||
errContains string
|
||||
}{
|
||||
{
|
||||
name: "wrapped in ```json fence",
|
||||
input: "```json\n" + validJSON + "\n```",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "wrapped in plain ``` fence (no language tag)",
|
||||
input: "```\n" + validJSON + "\n```",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "no fences at all (raw JSON)",
|
||||
input: validJSON,
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "leading prose then ```json fence",
|
||||
input: "Here is the extracted graph:\n\n```json\n" + validJSON + "\n```",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "trailing prose after closing fence",
|
||||
input: "```json\n" + validJSON + "\n```\n\nHope this helps!",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "extra surrounding whitespace and newlines",
|
||||
input: "\n\n ```json\n\n" + validJSON + "\n\n``` \n",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
// Issue #1113 Pattern 3: LLM hit max_tokens, no closing fence.
|
||||
// The response is structurally a JSON array we can still parse.
|
||||
name: "truncated response, opening ```json fence with no closer",
|
||||
input: "```json\n" + validJSON,
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
// Issue #1113 Pattern 1: bare backticks/markdown around JSON
|
||||
// without a well-formed fence pair.
|
||||
name: "stray backticks around JSON",
|
||||
input: "`" + validJSON + "`",
|
||||
wantNodes: 2,
|
||||
wantRels: 1,
|
||||
},
|
||||
{
|
||||
name: "JSON object embedded in prose (single dict)",
|
||||
input: "Result: {\"entity\": \"Alice\", \"entity_attributes\": [\"person\"]} -- end.",
|
||||
wantNodes: 1,
|
||||
wantRels: 0,
|
||||
},
|
||||
{
|
||||
name: "empty input",
|
||||
input: "",
|
||||
wantErr: true,
|
||||
errContains: "empty",
|
||||
},
|
||||
{
|
||||
name: "whitespace only",
|
||||
input: " \n\t ",
|
||||
wantErr: true,
|
||||
errContains: "empty",
|
||||
},
|
||||
{
|
||||
name: "fenced but body is invalid JSON",
|
||||
input: "```json\nnot json at all\n```",
|
||||
wantErr: true,
|
||||
errContains: "parse",
|
||||
},
|
||||
{
|
||||
name: "no recoverable JSON, only prose",
|
||||
input: "Sorry, I cannot extract a graph from this text.",
|
||||
wantErr: true,
|
||||
errContains: "parse",
|
||||
},
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
f := NewFormater()
|
||||
graph, err := f.ParseGraph(ctx, tc.input)
|
||||
if tc.wantErr {
|
||||
if err == nil {
|
||||
t.Fatalf("expected error, got nil (graph=%+v)", graph)
|
||||
}
|
||||
if tc.errContains != "" && !strings.Contains(err.Error(), tc.errContains) {
|
||||
t.Fatalf("error %q does not contain %q", err.Error(), tc.errContains)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if graph == nil {
|
||||
t.Fatalf("expected non-nil graph")
|
||||
}
|
||||
if got := len(graph.Node); got != tc.wantNodes {
|
||||
t.Errorf("nodes: got %d, want %d (graph=%+v)", got, tc.wantNodes, graph)
|
||||
}
|
||||
if got := len(graph.Relation); got != tc.wantRels {
|
||||
t.Errorf("relations: got %d, want %d (graph=%+v)", got, tc.wantRels, graph)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestExtractJSONLike covers the JSON-substring extraction helper in
|
||||
// isolation. The helper is used by the fallback path in extractContent when
|
||||
// fences are missing or malformed.
|
||||
func TestExtractJSONLike(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
}{
|
||||
{"empty", "", ""},
|
||||
{"plain object", `{"a": 1}`, `{"a": 1}`},
|
||||
{"plain array", `[1, 2, 3]`, `[1, 2, 3]`},
|
||||
{"object in prose", `noise {"a": 1} tail`, `{"a": 1}`},
|
||||
{"array preferred when first", `tail [1] {"a":1}`, `[1]`},
|
||||
{"object preferred when first", `tail {"a":1} [1]`, `{"a":1}`},
|
||||
{"nested braces", `{"a": {"b": [1,2]}}`, `{"a": {"b": [1,2]}}`},
|
||||
{"brace inside string literal", `{"a": "}{not real}"}`, `{"a": "}{not real}"}`},
|
||||
{"escaped quote inside string", `{"a": "he said \"hi\""}`, `{"a": "he said \"hi\""}`},
|
||||
{"unbalanced object returns empty", `{"a": 1`, ""},
|
||||
{"no json", `just words`, ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := extractJSONLike(tc.in)
|
||||
if got != tc.want {
|
||||
t.Errorf("extractJSONLike(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestStripFencesAndExtract focuses on the fence-recovery helper used as a
|
||||
// last resort when the main fence regex fails. Behavior must remain
|
||||
// conservative: return an empty string when nothing plausible can be
|
||||
// recovered, so the caller can fall through to existing behavior.
|
||||
func TestStripFencesAndExtract(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in string
|
||||
format FormatType
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "open fence with json tag and no close",
|
||||
in: "```json\n{\"a\":1}",
|
||||
format: FormatTypeJSON,
|
||||
want: `{"a":1}`,
|
||||
},
|
||||
{
|
||||
name: "open fence with no tag and no close",
|
||||
in: "```\n[1,2]",
|
||||
format: FormatTypeJSON,
|
||||
want: `[1,2]`,
|
||||
},
|
||||
{
|
||||
name: "well-formed fence still recovers body",
|
||||
in: "```json\n{\"a\":1}\n```",
|
||||
format: FormatTypeJSON,
|
||||
want: `{"a":1}`,
|
||||
},
|
||||
{
|
||||
name: "no fence but embedded json object",
|
||||
in: "Sure! {\"a\":1} done.",
|
||||
format: FormatTypeJSON,
|
||||
want: `{"a":1}`,
|
||||
},
|
||||
{
|
||||
name: "no fence and no json",
|
||||
in: "just prose",
|
||||
format: FormatTypeJSON,
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "empty input",
|
||||
in: "",
|
||||
format: FormatTypeJSON,
|
||||
want: "",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := stripFencesAndExtract(tc.in, tc.format)
|
||||
if got != tc.want {
|
||||
t.Errorf("stripFencesAndExtract(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsLikelyLanguageTag guards the heuristic used to drop language-tag
|
||||
// lines after an opening fence in the recovery path.
|
||||
func TestIsLikelyLanguageTag(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want bool
|
||||
}{
|
||||
{"json", true},
|
||||
{"yaml", true},
|
||||
{"yml", true},
|
||||
{"go", true},
|
||||
{"c++", true},
|
||||
{"objective-c", true},
|
||||
{"", false},
|
||||
{"this is not a tag", false},
|
||||
{`{"a":1}`, false},
|
||||
{strings.Repeat("a", 17), false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.in, func(t *testing.T) {
|
||||
if got := isLikelyLanguageTag(tc.in); got != tc.want {
|
||||
t.Errorf("isLikelyLanguageTag(%q) = %v, want %v", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user