mirror of
https://gitee.com/dromara/mayfly-go
synced 2026-05-17 16:35:19 +08:00
101 lines
1.9 KiB
Go
101 lines
1.9 KiB
Go
package utils
|
||
|
||
import (
|
||
"bytes"
|
||
"errors"
|
||
"mayfly-go/pkg/utils/collx"
|
||
"mayfly-go/pkg/utils/jsonx"
|
||
"regexp"
|
||
"strings"
|
||
)
|
||
|
||
// ParseLLMJSON 尝试从大模型输出中解析 JSON
|
||
func ParseLLMJSON[T any](raw string) (*T, error) {
|
||
candidates := extractJSONCandidates(raw)
|
||
|
||
var lastErr error
|
||
for _, c := range candidates {
|
||
if v, err := jsonx.ToByStr[T](c); err == nil {
|
||
return v, nil
|
||
} else {
|
||
lastErr = err
|
||
}
|
||
}
|
||
|
||
if lastErr == nil {
|
||
lastErr = errors.New("no json candidate found")
|
||
}
|
||
return nil, lastErr
|
||
}
|
||
|
||
// ParseLLMJSON2Map 解析 LLM 返回的JSON为map
|
||
func ParseLLMJSON2Map(raw string) (collx.M, error) {
|
||
if res, err := ParseLLMJSON[collx.M](raw); err != nil {
|
||
return nil, err
|
||
} else {
|
||
return *res, nil
|
||
}
|
||
}
|
||
|
||
func extractJSONCandidates(raw string) []string {
|
||
var results []string
|
||
text := strings.TrimSpace(raw)
|
||
|
||
// 1. 优先提取 code block 中的 JSON（对象 or 数组）
|
||
codeBlockRe := regexp.MustCompile(
|
||
"(?s)```(?:json)?\\s*([\\[{].*?[\\]}])\\s*```",
|
||
)
|
||
matches := codeBlockRe.FindAllStringSubmatch(text, -1)
|
||
for _, m := range matches {
|
||
results = append(results, strings.TrimSpace(m[1]))
|
||
}
|
||
|
||
// 2. 如果没找到 code block，尝试从全文裁剪 JSON
|
||
if len(results) == 0 {
|
||
if clipped := clipJSONValue(text); clipped != "" {
|
||
results = append(results, clipped)
|
||
}
|
||
}
|
||
|
||
return results
|
||
}
|
||
|
||
func clipJSONValue(s string) string {
|
||
objIdx := strings.Index(s, "{")
|
||
arrIdx := strings.Index(s, "[")
|
||
|
||
start := -1
|
||
var open, close byte
|
||
|
||
switch {
|
||
case objIdx != -1 && (arrIdx == -1 || objIdx < arrIdx):
|
||
start = objIdx
|
||
open, close = '{', '}'
|
||
case arrIdx != -1:
|
||
start = arrIdx
|
||
open, close = '[', ']'
|
||
default:
|
||
return ""
|
||
}
|
||
|
||
var buf bytes.Buffer
|
||
depth := 0
|
||
|
||
for i := start; i < len(s); i++ {
|
||
ch := s[i]
|
||
buf.WriteByte(ch)
|
||
|
||
switch ch {
|
||
case open:
|
||
depth++
|
||
case close:
|
||
depth--
|
||
if depth == 0 {
|
||
return buf.String()
|
||
}
|
||
}
|
||
}
|
||
|
||
return ""
|
||
}
|