265 lines
12 KiB
Python
265 lines
12 KiB
Python
"""工具参数兼容：模型把列表、对象参数写成 JSON 字符串时，按工具的参数 schema 解析回来
|
||
|
||
部分模型（例如 claude-haiku-5.5）在嵌套较深、内容较长时，常把 apply_actions 的 actions 整个写成 JSON 字符串，
|
||
或把列表里的每一项写成字符串；参数校验报 "Input should be a valid list" 后它往往原样再交，评测里出现过
|
||
连续 10 次直到调用次数用尽。这类问题和哪个工具无关，所以在工具调用入口统一处理，对本地工具和 MCP 工具都生效
|
||
|
||
- 合法的输入不动：schema 里任何一处接受字符串（包括联合类型的某个成员、数组某个成员的 items）时不解析；
|
||
布尔 schema、解析不了的 $ref 这类看不准的情况按"什么都接受"处理，交给原有的参数校验
|
||
- 只处理两层：参数本身，以及列表参数里的每一项
|
||
- 写成字符串的长 JSON 多数只是结尾的括号多写或少写了一个：原文以括号结尾、模型输出没有因为长度被截断时，
|
||
按还没闭合的括号重新补齐结尾；截断时丢掉的内容补不回来，补齐反而会得到一个意思不同的合法动作，所以不补
|
||
- 仍然解析不出来时不执行，用中文说明错在哪里：参数校验的英文报错只说"应该是列表"，模型会照旧写成字符串
|
||
- 解析或补齐后执行时，在工具结果末尾说明：补齐只能猜结尾，评测里模型把中间的 option 提前闭合、结尾又少一个括号，
|
||
补齐后 series 落到了 props 下，执行结果对不上而它看不出原因，反复原样提交
|
||
"""
|
||
|
||
import json
|
||
from dataclasses import replace
|
||
from typing import Any
|
||
|
||
from langchain.agents.middleware import wrap_tool_call
|
||
from langchain_core.messages import AIMessage, ToolMessage
|
||
from langchain_core.tools import BaseTool
|
||
from langgraph.types import Command
|
||
from loguru import logger
|
||
|
||
STRUCTURED = {"array", "object"}
|
||
ANY = frozenset({"string", "number", "integer", "boolean", "null", "array", "object"})
|
||
CLOSERS = {"[": "]", "{": "}"}
|
||
# 各家接口表示"输出达到长度上限被截断"的写法
|
||
TRUNCATED = {"length", "max_tokens"}
|
||
|
||
|
||
def _resolve(schema: Any, defs: dict) -> Any:
|
||
"""展开 $ref；指向的定义找不到时按 True（什么都接受）处理"""
|
||
while isinstance(schema, dict) and "$ref" in schema:
|
||
name = schema["$ref"].rsplit("/", 1)[-1]
|
||
if name not in defs:
|
||
return True
|
||
schema = defs[name]
|
||
return schema
|
||
|
||
|
||
def _types(schema: Any, defs: dict) -> set[str]:
|
||
"""schema 接受的 JSON 类型：anyOf、oneOf 取并集，allOf 取交集，没有类型约束时什么都接受"""
|
||
schema = _resolve(schema, defs)
|
||
if schema is False:
|
||
return set()
|
||
if not isinstance(schema, dict):
|
||
return set(ANY)
|
||
kind = schema.get("type")
|
||
if kind:
|
||
types = set(kind) if isinstance(kind, list) else {kind}
|
||
elif "properties" in schema:
|
||
types = {"object"}
|
||
elif "items" in schema:
|
||
types = {"array"}
|
||
else:
|
||
types = set(ANY)
|
||
for member in schema.get("allOf", []):
|
||
types &= _types(member, defs)
|
||
if members := [*schema.get("anyOf", []), *schema.get("oneOf", [])]:
|
||
types &= set().union(*(_types(member, defs) for member in members))
|
||
return types
|
||
|
||
|
||
def _items(schema: Any, defs: dict) -> list:
|
||
"""各个数组成员的 items schema；有数组成员没有限定 items 时，其中有一个 True（什么都接受）"""
|
||
schema = _resolve(schema, defs)
|
||
if not isinstance(schema, dict):
|
||
return [True] if schema is True else []
|
||
kind = schema.get("type")
|
||
own = set(kind) if isinstance(kind, list) else {kind}
|
||
found = []
|
||
if "prefixItems" in schema:
|
||
# 每个位置各有 schema 的元组，不逐项解析
|
||
found.append(True)
|
||
elif "items" in schema:
|
||
found.append(schema["items"])
|
||
elif "array" in own:
|
||
found.append(True)
|
||
for key in ("anyOf", "oneOf", "allOf"):
|
||
for member in schema.get(key, []):
|
||
found.extend(_items(member, defs))
|
||
return found
|
||
|
||
|
||
def _reclose(text: str) -> str | None:
|
||
"""去掉结尾连续的 ] 和 }，按还没闭合的括号重新补齐；结尾停在字符串中间时补不了"""
|
||
body = text.rstrip().rstrip("]} \t\r\n")
|
||
stack: list[str] = []
|
||
in_string = escaped = False
|
||
for char in body:
|
||
if in_string:
|
||
if escaped:
|
||
escaped = False
|
||
elif char == "\\":
|
||
escaped = True
|
||
elif char == '"':
|
||
in_string = False
|
||
elif char == '"':
|
||
in_string = True
|
||
elif char in CLOSERS:
|
||
stack.append(char)
|
||
elif char in "]}" and (not stack or CLOSERS[stack.pop()] != char):
|
||
# 中间就有括号对不上，不是结尾的问题，补不了
|
||
return None
|
||
if in_string or not stack:
|
||
return None
|
||
return body + "".join(CLOSERS[item] for item in reversed(stack))
|
||
|
||
|
||
def _matches(value: Any, types: set[str]) -> bool:
|
||
return (isinstance(value, list) and "array" in types) or (
|
||
isinstance(value, dict) and "object" in types
|
||
)
|
||
|
||
|
||
def _parse(value: Any, types: set[str], repair: bool) -> tuple[Any, str | None, bool]:
|
||
"""schema 只接受列表或对象、却收到字符串时解析成 JSON；返回 (值, 解析不了时的说明, 是否补了结尾括号)"""
|
||
if not isinstance(value, str) or "string" in types or not types & STRUCTURED:
|
||
return value, None, False
|
||
expected = "数组" if "array" in types else "对象"
|
||
patched = False
|
||
try:
|
||
parsed = json.loads(value)
|
||
except json.JSONDecodeError as error:
|
||
# 原文不以括号结尾多半是被截断了，补齐结尾会丢掉后面的内容，不补
|
||
ends_closed = value.rstrip().endswith(("]", "}"))
|
||
repaired = _reclose(value) if repair and ends_closed else None
|
||
try:
|
||
parsed, patched = (json.loads(repaired) if repaired else None), True
|
||
except json.JSONDecodeError:
|
||
parsed = None
|
||
if not _matches(parsed, types):
|
||
problem = (
|
||
f"应该直接写成 JSON {expected}，不要写成字符串；"
|
||
f"这个字符串也不是合法的 JSON（第 {error.pos} 个字符附近：{error.msg}）"
|
||
)
|
||
return value, problem, False
|
||
if not _matches(parsed, types):
|
||
return value, f"应该直接写成 JSON {expected}，不要写成字符串", False
|
||
return parsed, None, patched
|
||
|
||
|
||
def normalize_args(
|
||
tool: BaseTool, args: dict, *, repair: bool = True
|
||
) -> tuple[dict, list[str], list[str]]:
|
||
"""按工具的参数 schema 把写成字符串的列表、对象参数解析回来
|
||
|
||
返回新的参数（不修改传入的）、解析不了的参数说明、补了结尾括号的参数名；有说明时不该执行这次调用
|
||
repair=False 时不补括号，用于模型输出被截断的情况
|
||
"""
|
||
schema = tool.tool_call_schema
|
||
schema = schema if isinstance(schema, dict) else schema.model_json_schema()
|
||
defs = schema.get("$defs", {})
|
||
properties = schema.get("properties", {})
|
||
fixed, problems, repaired = {}, [], []
|
||
for name, value in args.items():
|
||
prop = properties.get(name)
|
||
if prop is None:
|
||
fixed[name] = value
|
||
continue
|
||
value, problem, patched = _parse(value, _types(prop, defs), repair)
|
||
if problem:
|
||
problems.append(f"{name} {problem}")
|
||
item_schemas = _items(prop, defs)
|
||
if isinstance(value, list) and item_schemas:
|
||
item_types = set().union(*(_types(item, defs) for item in item_schemas))
|
||
parsed_items = []
|
||
for index, item in enumerate(value):
|
||
item, problem, item_patched = _parse(item, item_types, repair)
|
||
if problem:
|
||
problems.append(f"{name} 的第 {index + 1} 项{problem}")
|
||
patched = patched or item_patched
|
||
parsed_items.append(item)
|
||
value = parsed_items
|
||
if patched:
|
||
repaired.append(name)
|
||
fixed[name] = value
|
||
return fixed, problems, repaired
|
||
|
||
|
||
def _truncated(state: Any, call_id: str) -> bool:
|
||
"""发起这次调用的模型输出是不是因为长度上限被截断"""
|
||
messages = state.get("messages", []) if isinstance(state, dict) else []
|
||
message = next(
|
||
(
|
||
item
|
||
for item in reversed(messages)
|
||
if isinstance(item, AIMessage)
|
||
and any(call["id"] == call_id for call in item.tool_calls)
|
||
),
|
||
None,
|
||
)
|
||
metadata = message.response_metadata if message is not None else {}
|
||
return bool(
|
||
{metadata.get("finish_reason"), metadata.get("stop_reason")} & TRUNCATED
|
||
)
|
||
|
||
|
||
@wrap_tool_call
|
||
async def normalize_tool_args(request, handler):
|
||
"""在参数校验前统一修正参数；没有改动时原样执行，解析不了时返回中文说明，不执行"""
|
||
tool = request.tool
|
||
args = request.tool_call.get("args")
|
||
if tool is None or not isinstance(args, dict):
|
||
return await handler(request)
|
||
truncated = _truncated(request.state, request.tool_call["id"])
|
||
fixed, problems, repaired = normalize_args(tool, args, repair=not truncated)
|
||
if problems:
|
||
logger.info(
|
||
f"工具 {tool.name} 的参数写成了不合法的 JSON 字符串"
|
||
f"{'（模型输出被截断）' if truncated else ''}: {problems}"
|
||
)
|
||
return ToolMessage(
|
||
"参数格式错误，这次调用没有执行：\n"
|
||
+ "\n".join(f"- {item}" for item in problems)
|
||
+ ("\n- 输出达到了长度上限被截断，分成几批提交" if truncated else ""),
|
||
tool_call_id=request.tool_call["id"],
|
||
name=tool.name,
|
||
status="error",
|
||
)
|
||
if fixed == args:
|
||
return await handler(request)
|
||
changed = [name for name in args if fixed[name] != args[name]]
|
||
strict = [name for name in changed if name not in repaired]
|
||
notes = []
|
||
if strict:
|
||
logger.info(
|
||
f"工具 {tool.name} 的参数 {'、'.join(strict)} 写成了 JSON 字符串，已解析"
|
||
)
|
||
notes.append(f"{'、'.join(strict)} 写成了 JSON 字符串，已解析后执行")
|
||
if repaired:
|
||
logger.info(
|
||
f"工具 {tool.name} 的参数 {'、'.join(repaired)} 写成了 JSON 字符串，结尾括号对不上，补齐后解析"
|
||
)
|
||
# 补的是结尾，括号多写或少写在中间时，补齐后层级是错的，但它是合法的 JSON，执行时不一定报错
|
||
notes.append(
|
||
f"{'、'.join(repaired)} 写成了 JSON 字符串，结尾的括号数量不对，已按还没闭合的括号补齐后执行；"
|
||
"如果中间的括号写错了，补齐后的层级也是错的，对照执行结果核对"
|
||
)
|
||
result = await handler(
|
||
request.override(tool_call={**request.tool_call, "args": fixed})
|
||
)
|
||
# 不告诉模型时它会一直写成字符串，补齐猜错层级时它也看不出原因
|
||
return _with_note(
|
||
result, "参数说明：" + "；".join(notes) + "；以后直接写成 JSON，不要写成字符串"
|
||
)
|
||
|
||
|
||
def _with_note(result: Any, note: str) -> Any:
|
||
"""在工具返回的消息末尾加一句说明；工具返回 Command 时改其中的工具消息"""
|
||
if isinstance(result, ToolMessage):
|
||
return result.model_copy(update={"content": f"{result.content}\n{note}"})
|
||
if isinstance(result, Command) and isinstance(result.update, dict):
|
||
messages = [
|
||
item.model_copy(update={"content": f"{item.content}\n{note}"})
|
||
if isinstance(item, ToolMessage)
|
||
else item
|
||
for item in result.update.get("messages", [])
|
||
]
|
||
return replace(result, update={**result.update, "messages": messages})
|
||
return result
|