252 lines
8.4 KiB
Python
252 lines
8.4 KiB
Python
"""
|
||
Agent JSON 解析工具
|
||
从 LLM 响应中安全地解析 JSON，参考 llm/service.py 的实现
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
import logging
|
||
from typing import Dict, Any, List, Optional, Union
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 尝试导入 json-repair 库
|
||
try:
|
||
from json_repair import repair_json
|
||
JSON_REPAIR_AVAILABLE = True
|
||
except ImportError:
|
||
JSON_REPAIR_AVAILABLE = False
|
||
logger.debug("json-repair library not available")
|
||
|
||
|
||
class AgentJsonParser:
|
||
"""Agent 专用的 JSON 解析器"""
|
||
|
||
@staticmethod
|
||
def clean_text(text: str) -> str:
|
||
"""清理文本中的控制字符"""
|
||
if not text:
|
||
return ""
|
||
# 移除 BOM 和零宽字符
|
||
text = text.replace('\ufeff', '').replace('\u200b', '').replace('\u200c', '').replace('\u200d', '')
|
||
return text
|
||
|
||
@staticmethod
|
||
def fix_json_format(text: str) -> str:
|
||
"""修复常见的 JSON 格式问题"""
|
||
text = text.strip()
|
||
# 移除尾部逗号
|
||
text = re.sub(r',(\s*[}\]])', r'\1', text)
|
||
# 修复未转义的换行符（在字符串值中）
|
||
text = re.sub(r':\s*"([^"]*)\n([^"]*)"', r': "\1\\n\2"', text)
|
||
return text
|
||
|
||
@classmethod
|
||
def extract_from_markdown(cls, text: str) -> Dict[str, Any]:
|
||
"""从 markdown 代码块提取 JSON"""
|
||
match = re.search(r'```(?:json)?\s*(\{[\s\S]*?\})\s*```', text)
|
||
if match:
|
||
return json.loads(match.group(1))
|
||
raise ValueError("No markdown code block found")
|
||
|
||
@classmethod
|
||
def extract_json_object(cls, text: str) -> Dict[str, Any]:
|
||
"""智能提取 JSON 对象"""
|
||
start_idx = text.find('{')
|
||
if start_idx == -1:
|
||
raise ValueError("No JSON object found")
|
||
|
||
# 考虑字符串内的花括号和转义字符
|
||
brace_count = 0
|
||
in_string = False
|
||
escape_next = False
|
||
end_idx = -1
|
||
|
||
for i in range(start_idx, len(text)):
|
||
char = text[i]
|
||
|
||
if escape_next:
|
||
escape_next = False
|
||
continue
|
||
|
||
if char == '\\':
|
||
escape_next = True
|
||
continue
|
||
|
||
if char == '"' and not escape_next:
|
||
in_string = not in_string
|
||
continue
|
||
|
||
if not in_string:
|
||
if char == '{':
|
||
brace_count += 1
|
||
elif char == '}':
|
||
brace_count -= 1
|
||
if brace_count == 0:
|
||
end_idx = i + 1
|
||
break
|
||
|
||
if end_idx == -1:
|
||
# 如果找不到完整的 JSON，尝试使用最后一个 }
|
||
last_brace = text.rfind('}')
|
||
if last_brace > start_idx:
|
||
end_idx = last_brace + 1
|
||
else:
|
||
raise ValueError("Incomplete JSON object")
|
||
|
||
json_str = text[start_idx:end_idx]
|
||
# 修复格式问题
|
||
json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
|
||
|
||
return json.loads(json_str)
|
||
|
||
@classmethod
|
||
def fix_truncated_json(cls, text: str) -> Dict[str, Any]:
|
||
"""修复截断的 JSON"""
|
||
start_idx = text.find('{')
|
||
if start_idx == -1:
|
||
raise ValueError("Cannot fix truncated JSON")
|
||
|
||
json_str = text[start_idx:]
|
||
|
||
# 计算缺失的闭合符号
|
||
open_braces = json_str.count('{')
|
||
close_braces = json_str.count('}')
|
||
open_brackets = json_str.count('[')
|
||
close_brackets = json_str.count(']')
|
||
|
||
# 补全缺失的闭合符号
|
||
json_str += ']' * max(0, open_brackets - close_brackets)
|
||
json_str += '}' * max(0, open_braces - close_braces)
|
||
|
||
# 修复格式
|
||
json_str = re.sub(r',(\s*[}\]])', r'\1', json_str)
|
||
return json.loads(json_str)
|
||
|
||
@classmethod
|
||
def repair_with_library(cls, text: str) -> Dict[str, Any]:
|
||
"""使用 json-repair 库修复损坏的 JSON"""
|
||
if not JSON_REPAIR_AVAILABLE:
|
||
raise ValueError("json-repair library not available")
|
||
|
||
start_idx = text.find('{')
|
||
if start_idx == -1:
|
||
raise ValueError("No JSON object found for repair")
|
||
|
||
end_idx = text.rfind('}')
|
||
if end_idx > start_idx:
|
||
json_str = text[start_idx:end_idx + 1]
|
||
else:
|
||
json_str = text[start_idx:]
|
||
|
||
repaired = repair_json(json_str, return_objects=True)
|
||
|
||
if isinstance(repaired, dict):
|
||
return repaired
|
||
|
||
raise ValueError(f"json-repair returned unexpected type: {type(repaired)}")
|
||
|
||
@classmethod
|
||
def parse(cls, text: str, default: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
从 LLM 响应中解析 JSON（增强版）
|
||
|
||
Args:
|
||
text: LLM 响应文本
|
||
default: 解析失败时返回的默认值，如果为 None 则抛出异常
|
||
|
||
Returns:
|
||
解析后的字典
|
||
"""
|
||
if not text or not text.strip():
|
||
if default is not None:
|
||
logger.warning("LLM 响应为空，返回默认值")
|
||
return default
|
||
raise ValueError("LLM 响应内容为空")
|
||
|
||
clean = cls.clean_text(text)
|
||
|
||
# 尝试多种方式解析
|
||
attempts = [
|
||
("直接解析", lambda: json.loads(text)),
|
||
("清理后解析", lambda: json.loads(cls.fix_json_format(clean))),
|
||
("Markdown 提取", lambda: cls.extract_from_markdown(text)),
|
||
("智能提取", lambda: cls.extract_json_object(clean)),
|
||
("截断修复", lambda: cls.fix_truncated_json(clean)),
|
||
("json-repair", lambda: cls.repair_with_library(text)),
|
||
]
|
||
|
||
last_error = None
|
||
for name, attempt in attempts:
|
||
try:
|
||
result = attempt()
|
||
if result and isinstance(result, dict):
|
||
if name != "直接解析":
|
||
logger.debug(f"✅ JSON 解析成功（方法: {name}）")
|
||
return result
|
||
except Exception as e:
|
||
last_error = e
|
||
logger.debug(f"JSON 解析方法 '{name}' 失败: {e}")
|
||
|
||
# 所有尝试都失败
|
||
if default is not None:
|
||
logger.warning(f"JSON 解析失败，返回默认值。原始内容: {text[:200]}...")
|
||
return default
|
||
|
||
logger.error(f"❌ 无法解析 JSON，原始内容: {text[:500]}...")
|
||
raise ValueError(f"无法解析 JSON: {last_error}")
|
||
|
||
@classmethod
|
||
def parse_findings(cls, text: str) -> List[Dict[str, Any]]:
|
||
"""
|
||
专门解析 findings 列表
|
||
|
||
Args:
|
||
text: LLM 响应文本
|
||
|
||
Returns:
|
||
findings 列表（每个元素都是字典）
|
||
"""
|
||
try:
|
||
result = cls.parse(text, default={"findings": []})
|
||
findings = result.get("findings", [])
|
||
|
||
# 确保每个 finding 都是字典
|
||
valid_findings = []
|
||
for f in findings:
|
||
if isinstance(f, dict):
|
||
valid_findings.append(f)
|
||
elif isinstance(f, str):
|
||
# 尝试将字符串解析为 JSON
|
||
try:
|
||
parsed = json.loads(f)
|
||
if isinstance(parsed, dict):
|
||
valid_findings.append(parsed)
|
||
except json.JSONDecodeError:
|
||
logger.warning(f"跳过无效的 finding（字符串）: {f[:100]}...")
|
||
else:
|
||
logger.warning(f"跳过无效的 finding（类型: {type(f)}）")
|
||
|
||
return valid_findings
|
||
|
||
except Exception as e:
|
||
logger.error(f"解析 findings 失败: {e}")
|
||
return []
|
||
|
||
@classmethod
|
||
def safe_get(cls, data: Union[Dict, str, Any], key: str, default: Any = None) -> Any:
|
||
"""
|
||
安全地从数据中获取值
|
||
|
||
Args:
|
||
data: 可能是字典或其他类型
|
||
key: 要获取的键
|
||
default: 默认值
|
||
|
||
Returns:
|
||
获取的值或默认值
|
||
"""
|
||
if isinstance(data, dict):
|
||
return data.get(key, default)
|
||
return default
|