mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 00:14:33 +00:00
fix(kg): sanitize LLM JSON output before parsing (#167)
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.
This commit is contained in:
viettranx committed
1 parent
5e2fa395c7
commit
6dcac3d7e3
2 files changed
+248
-2
No files matched your search
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user