diff --git a/internal/common/tools.go b/internal/common/tools.go index 1515e290..783354ff 100644 --- a/internal/common/tools.go +++ b/internal/common/tools.go @@ -96,6 +96,10 @@ func DeduplicateWithScore[T ScoreComparable, K comparable](keyFunc func(T) K, it // {"key": "value"} // ``` // or regular JSON responses directly. +// jsonCodeFenceRE extracts a JSON payload wrapped in a Markdown code fence. +// Compiled once: ParseLLMJsonResponse runs on the graph-extraction path. +var jsonCodeFenceRE = regexp.MustCompile("```(?:json)?\\s*([\\s\\S]*?)```") + func ParseLLMJsonResponse(content string, target interface{}) error { // First, try to parse directly as JSON err := json.Unmarshal([]byte(content), target) @@ -104,8 +108,7 @@ func ParseLLMJsonResponse(content string, target interface{}) error { } // If direct parsing fails, try to extract JSON from code blocks - re := regexp.MustCompile("```(?:json)?\\s*([\\s\\S]*?)```") - matches := re.FindStringSubmatch(content) + matches := jsonCodeFenceRE.FindStringSubmatch(content) if len(matches) >= 2 { // Extract the JSON content within the code block jsonContent := strings.TrimSpace(matches[1]) diff --git a/internal/common/tools_test.go b/internal/common/tools_test.go new file mode 100644 index 00000000..6b75caf4 --- /dev/null +++ b/internal/common/tools_test.go @@ -0,0 +1,47 @@ +package common + +import "testing" + +func TestParseLLMJsonResponse(t *testing.T) { + type payload struct { + Key string `json:"key"` + } + + tests := []struct { + name string + content string + want string + wantErr bool + }{ + {"direct json", `{"key":"value"}`, "value", false}, + {"fenced json", "```json\n{\"key\":\"fenced\"}\n```", "fenced", false}, + {"fenced no lang", "```\n{\"key\":\"plain\"}\n```", "plain", false}, + {"prose around fence", "Here you go:\n```json\n{\"key\":\"wrapped\"}\n```\nThanks", "wrapped", false}, + {"not json", "just some text", "", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var got payload + err := ParseLLMJsonResponse(tt.content, &got) + if (err != nil) != tt.wantErr { + t.Fatalf("err = %v, wantErr %v", err, tt.wantErr) + } + if !tt.wantErr && got.Key != tt.want { + t.Errorf("Key = %q, want %q", got.Key, tt.want) + } + }) + } +} + +func BenchmarkParseLLMJsonResponse_Fenced(b *testing.B) { + const content = "```json\n{\"name\":\"Acme\",\"slug\":\"entity/acme\",\"aliases\":[\"A\",\"B\"]}\n```" + type payload struct { + Name string `json:"name"` + } + b.ReportAllocs() + for i := 0; i < b.N; i++ { + var p payload + _ = ParseLLMJsonResponse(content, &p) + } +}