From 6dcac3d7e30acdea4e3dd71424ebee8864ed9bc6 Mon Sep 17 00:00:00 2001 From: viettranx Date: Sat, 14 Mar 2026 16:22:48 +0700 Subject: [PATCH] fix(kg): sanitize LLM JSON output before parsing (#167) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add sanitizeJSON() to fix malformed decimal numbers (e.g. '0. 85' → '0.85') and trailing commas before closing brackets. Fixes extraction failures with Gemini 2.5 Flash which occasionally produces invalid JSON. --- internal/knowledgegraph/extractor.go | 87 +++++++++++- internal/knowledgegraph/extractor_test.go | 163 ++++++++++++++++++++++ 2 files changed, 248 insertions(+), 2 deletions(-) create mode 100644 internal/knowledgegraph/extractor_test.go diff --git a/internal/knowledgegraph/extractor.go b/internal/knowledgegraph/extractor.go index 45611ef7..40b58f1a 100644 --- a/internal/knowledgegraph/extractor.go +++ b/internal/knowledgegraph/extractor.go @@ -98,12 +98,27 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR var result ExtractionResult content := strings.TrimSpace(resp.Content) content = stripCodeBlock(content) + + originalContent := content + content = sanitizeJSON(content) + if content != originalContent { + slog.Debug("kg extraction: sanitized JSON output", + "original_len", len(originalContent), + "sanitized_len", len(content), + ) + } + if err := json.Unmarshal([]byte(content), &result); err != nil { - preview := content + preview := originalContent if len(preview) > 300 { preview = preview[:300] + "..." } - slog.Warn("kg extraction: failed to parse LLM response", "error", err, "content_len", len(content), "finish_reason", resp.FinishReason, "preview", preview) + slog.Warn("kg extraction: failed to parse LLM response", + "error", err, + "content_len", len(originalContent), + "finish_reason", resp.FinishReason, + "preview", preview, + ) return nil, fmt.Errorf("parse extraction result: %w", err) } @@ -128,6 +143,73 @@ func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionR return filtered, nil } +// sanitizeJSON fixes common LLM JSON issues while preserving string values. +// It walks the JSON character-by-character, only applying fixes outside quoted strings: +// - Malformed decimals: "0. 85" → "0.85" +// - Trailing commas: [1, 2,] → [1, 2] +func sanitizeJSON(s string) string { + var b strings.Builder + b.Grow(len(s)) + + inString := false + escaped := false + + for i := 0; i < len(s); i++ { + ch := s[i] + + if escaped { + b.WriteByte(ch) + escaped = false + continue + } + + if ch == '\\' && inString { + b.WriteByte(ch) + escaped = true + continue + } + + if ch == '"' { + inString = !inString + b.WriteByte(ch) + continue + } + + if inString { + b.WriteByte(ch) + continue + } + + // Fix malformed decimals: "0. 85" → "0.85" + if ch == '.' && i > 0 && isDigit(s[i-1]) { + b.WriteByte('.') + for i+1 < len(s) && s[i+1] == ' ' { + i++ + } + continue + } + + // Fix trailing commas: skip comma if next non-whitespace is } or ] + if ch == ',' { + j := i + 1 + for j < len(s) && (s[j] == ' ' || s[j] == '\t' || s[j] == '\n' || s[j] == '\r') { + j++ + } + if j < len(s) && (s[j] == '}' || s[j] == ']') { + continue + } + } + + b.WriteByte(ch) + } + + return b.String() +} + +func isDigit(c byte) bool { + return c >= '0' && c <= '9' +} + // splitChunks splits text into chunks at paragraph boundaries (\n\n). func splitChunks(text string, maxChars int) []string { if len(text) <= maxChars { @@ -204,3 +286,4 @@ func stripCodeBlock(s string) string { } return strings.TrimSpace(s) } + diff --git a/internal/knowledgegraph/extractor_test.go b/internal/knowledgegraph/extractor_test.go new file mode 100644 index 00000000..9a98c60a --- /dev/null +++ b/internal/knowledgegraph/extractor_test.go @@ -0,0 +1,163 @@ +package knowledgegraph + +import ( + "encoding/json" + "testing" +) + +func TestSanitizeJSON(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + { + name: "valid JSON unchanged", + input: `{"confidence": 0.85}`, + want: `{"confidence": 0.85}`, + }, + { + name: "fix spaced decimal", + input: `{"confidence": 0. 85}`, + want: `{"confidence": 0.85}`, + }, + { + name: "fix multiple spaced decimals", + input: `{"a": 0. 9, "b": 1. 0}`, + want: `{"a": 0.9, "b": 1.0}`, + }, + { + name: "fix spaced decimal with multiple spaces", + input: `{"confidence": 0. 75}`, + want: `{"confidence": 0.75}`, + }, + { + name: "preserve decimal-like pattern in strings", + input: `{"description": "Founded in 2023. 15 employees"}`, + want: `{"description": "Founded in 2023. 15 employees"}`, + }, + { + name: "preserve spaced period in strings (no digit before dot)", + input: `{"description": "Mr. Smith leads the team"}`, + want: `{"description": "Mr. Smith leads the team"}`, + }, + { + name: "mixed: fix value decimal, preserve string decimal", + input: `{"description": "Version 2. 0 alpha", "confidence": 0. 9}`, + want: `{"description": "Version 2. 0 alpha", "confidence": 0.9}`, + }, + { + name: "fix trailing comma in array", + input: `{"items": [1, 2, 3,]}`, + want: `{"items": [1, 2, 3]}`, + }, + { + name: "fix trailing comma in object", + input: `{"a": 1, "b": 2,}`, + want: `{"a": 1, "b": 2}`, + }, + { + name: "trailing comma with whitespace", + input: `{"a": 1, }`, + want: `{"a": 1 }`, + }, + { + name: "trailing comma with newline", + input: "{\"a\": 1,\n}", + want: "{\"a\": 1\n}", + }, + { + name: "preserve comma in string value", + input: `{"text": "hello, world,"}`, + want: `{"text": "hello, world,"}`, + }, + { + name: "escaped quote in string", + input: `{"text": "she said \"0. 5\" loudly", "val": 0. 5}`, + want: `{"text": "she said \"0. 5\" loudly", "val": 0.5}`, + }, + { + name: "nested structure", + input: `{"entities": [{"confidence": 0. 85,}]}`, + want: `{"entities": [{"confidence": 0.85}]}`, + }, + { + name: "empty input", + input: "", + want: "", + }, + { + name: "no fixes needed", + input: `{"entities": [], "relations": []}`, + want: `{"entities": [], "relations": []}`, + }, + { + name: "truncated LLM response does not crash", + input: `{"entities": [{"name": "Facebook Ads", "confidence": 0.9}, {"name": "TikT`, + want: `{"entities": [{"name": "Facebook Ads", "confidence": 0.9}, {"name": "TikT`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := sanitizeJSON(tt.input) + if got != tt.want { + t.Errorf("sanitizeJSON():\n input: %s\n got: %s\n want: %s", tt.input, got, tt.want) + } + }) + } +} + +func TestTruncatedJSONParseFails(t *testing.T) { + truncated := `{"entities": [{"name": "Facebook Ads", "confidence": 0.9}, {"name": "TikT` + sanitized := sanitizeJSON(truncated) + + var result ExtractionResult + err := json.Unmarshal([]byte(sanitized), &result) + if err == nil { + t.Error("expected parse error for truncated JSON, got nil") + } +} + +func TestStripCodeBlock(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + { + name: "no code block", + input: `{"entities": []}`, + want: `{"entities": []}`, + }, + { + name: "json code block", + input: "```json\n{\"entities\": []}\n```", + want: `{"entities": []}`, + }, + { + name: "plain code block", + input: "```\n{\"entities\": []}\n```", + want: `{"entities": []}`, + }, + { + name: "code block with surrounding whitespace", + input: " ```json\n{\"entities\": []}\n``` ", + want: `{"entities": []}`, + }, + { + name: "empty input", + input: "", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := stripCodeBlock(tt.input) + if got != tt.want { + t.Errorf("stripCodeBlock():\n got: %s\n want: %s", got, tt.want) + } + }) + } +}