[Fix] Preserve GLM tool argument types across JSON Schema unions (#39136)

Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
William Hu
2026-09-13 13:47:13 +08:00
committed by GitHub
co-authored by Xinyuan Tong
parent 34b2904741
commit 6d6d42d1c0
2 changed files with 618 additions and 154 deletions
@@ -16,11 +16,7 @@ from sglang.srt.function_call.core_types import (
ToolCallItem,
_GetInfoFunc,
)
from sglang.srt.function_call.utils import (
get_schema_properties,
infer_type_from_json_schema,
safe_literal_eval,
)
from sglang.srt.function_call.utils import safe_literal_eval
logger = logging.getLogger(__name__)
@@ -50,45 +46,170 @@ class StreamState(str, Enum):
IN_VALUE = "IN_VALUE"
def get_argument_type(
func_name: str, arg_key: str, defined_tools: List[Tool]
) -> Optional[str]:
"""Get the expected type of a function argument from tool definitions.
_UNSET = object()
Supports complex JSON Schema definitions including:
- Direct type field (including type arrays)
- anyOf/oneOf: parameter can be any of multiple types
- enum: parameter must be one of enum values
- allOf: parameter must satisfy all type definitions
- properties: inferred as object type
- items: inferred as array type
Args:
func_name: Name of the function/tool
arg_key: Name of the argument
defined_tools: List of available tools
def _json_type(value: Any) -> str:
if isinstance(value, float) and value.is_integer():
return "integer"
return {
type(None): "null",
bool: "boolean",
int: "integer",
float: "number",
str: "string",
list: "array",
dict: "object",
}[type(value)]
Returns:
The type string (e.g., 'string', 'number', 'object') or None if not found
"""
name2tool = {tool.function.name: tool for tool in defined_tools}
# Check if function exists
tool = name2tool.get(func_name)
if not tool:
def _json_equal(left: Any, right: Any) -> bool:
if _json_type(left) != _json_type(right):
return False
if isinstance(left, dict):
return left.keys() == right.keys() and all(
_json_equal(value, right[key]) for key, value in left.items()
)
if isinstance(left, list):
return len(left) == len(right) and all(
_json_equal(a, b) for a, b in zip(left, right)
)
return left == right
def _matches_discriminators(schema: Any, root: Any, arguments: Dict[str, Any]) -> bool:
for name, value in arguments.items():
types = _argument_types(schema, root, name, value=value)
if types is not None and _json_type(value) not in types:
return False
return True
def _argument_types(
schema: Any,
root: Any,
key: Optional[str] = None,
seen: frozenset[int] = frozenset(),
arguments: Optional[Dict[str, Any]] = None,
value: Any = _UNSET,
) -> Optional[set[str]]:
"""Keep every possible type until the argument value disambiguates a union."""
if not isinstance(schema, dict) or id(schema) in seen:
return None
seen = seen | {id(schema)}
constraints = []
ref = schema.get("$ref")
if isinstance(ref, str) and (ref == "#" or ref.startswith("#/")):
target = root
try:
for part in ref[2:].split("/") if ref != "#" else []:
part = part.replace("~1", "/").replace("~0", "~")
target = target[int(part)] if isinstance(target, list) else target[part]
types = _argument_types(target, root, key, seen, arguments, value)
if types is not None:
constraints.append(types)
except (KeyError, IndexError, TypeError, ValueError):
pass
if key is not None:
field = schema.get("properties", {}).get(key)
types = _argument_types(field, root, seen=seen, value=value)
if types is not None:
constraints.append(types)
else:
types = schema.get("type")
if isinstance(types, str):
constraints.append({types, "integer"} if types == "number" else {types})
elif isinstance(types, list):
constraints.append(
set(types) | ({"integer"} if "number" in types else set())
)
values = [schema["const"]] if "const" in schema else schema.get("enum")
if isinstance(values, list):
constraints.append(
{
_json_type(allowed)
for allowed in values
if value is _UNSET or _json_equal(value, allowed)
}
)
for keyword in ("anyOf", "oneOf"):
if isinstance(schema.get(keyword), list):
alternatives = [
_argument_types(branch, root, key, seen, arguments, value)
for branch in schema[keyword]
if key is None or _matches_discriminators(branch, root, arguments or {})
]
if alternatives and all(types is not None for types in alternatives):
constraints.append(set().union(*alternatives))
if isinstance(schema.get("allOf"), list):
for branch in schema["allOf"]:
types = _argument_types(branch, root, key, seen, arguments, value)
if types is not None:
constraints.append(types)
return set.intersection(*constraints) if constraints else None
# Get parameters safely using getattr
params = getattr(tool.function, "parameters", None)
arg_spec = get_schema_properties(params).get(arg_key)
if isinstance(arg_spec, dict):
# Use the new type inference function for complex JSON Schema support
return infer_type_from_json_schema(arg_spec)
def _get_argument_types(
func_name, arg_key, defined_tools, arguments=None, value=_UNSET
):
for tool in defined_tools:
if tool.function.name == func_name:
schema = tool.function.parameters
return _argument_types(
schema, schema, arg_key, arguments=arguments, value=value
)
return None
def get_argument_type(
func_name: str,
arg_key: str,
defined_tools: List[Tool],
arguments: Optional[Dict[str, Any]] = None,
) -> Optional[str]:
types = _get_argument_types(func_name, arg_key, defined_tools, arguments)
if types == {"number", "integer"}:
return "number"
return next(iter(types)) if types and len(types) == 1 else None
def _needs_argument_context(schema: Any) -> bool:
keys, seen = set(), set()
has_union = False
def visit(node):
nonlocal has_union
if not isinstance(node, dict) or id(node) in seen:
return
seen.add(id(node))
keys.update(node.get("properties", {}))
ref = node.get("$ref")
if isinstance(ref, str) and (ref == "#" or ref.startswith("#/")):
target = schema
try:
for part in ref[2:].split("/") if ref != "#" else []:
part = part.replace("~1", "/").replace("~0", "~")
target = (
target[int(part)] if isinstance(target, list) else target[part]
)
visit(target)
except (KeyError, IndexError, TypeError, ValueError):
pass
for keyword in ("anyOf", "oneOf", "allOf"):
branches = node.get(keyword)
if isinstance(branches, list):
has_union |= keyword != "allOf"
for branch in branches:
visit(branch)
visit(schema)
for key in keys if has_union else ():
types = _argument_types(schema, schema, key)
if types is None or len(types) > 1 and types != {"number", "integer"}:
return True
return False
def _convert_to_number(value: str) -> Any:
"""Convert string to appropriate number type (int or float).
@@ -107,6 +228,24 @@ def _convert_to_number(value: str) -> Any:
return value
def _convert_to_integer(value: str) -> Any:
try:
return int(value)
except (ValueError, AttributeError):
return value
def _coerce_numeric_string(value: Any, arg_type: Optional[str]) -> Any:
# A quoted value must not become a type its schema forbids: "1.5" may coerce
# for number but never for integer.
if isinstance(value, str):
if arg_type == "integer":
return _convert_to_integer(value)
if arg_type == "number":
return _convert_to_number(value)
return value
def parse_arguments(
json_value: str, arg_type: Optional[str] = None
) -> Tuple[Any, bool]:
@@ -121,25 +260,14 @@ def parse_arguments(
"""
# Strategy 1: Direct JSON parsing
try:
parsed_value = json.loads(json_value)
# Type coercion for number type
if arg_type == "number" and isinstance(parsed_value, str):
parsed_value = _convert_to_number(parsed_value)
return parsed_value, True
return _coerce_numeric_string(json.loads(json_value), arg_type), True
except (json.JSONDecodeError, ValueError):
pass
# Strategy 2: Unescape and parse
try:
wrapped = json.loads('{"tmp": "' + json_value + '"}')
parsed_value = json.loads(wrapped["tmp"])
if arg_type == "number" and isinstance(parsed_value, str):
parsed_value = _convert_to_number(parsed_value)
return parsed_value, True
return _coerce_numeric_string(json.loads(wrapped["tmp"]), arg_type), True
except (json.JSONDecodeError, ValueError, KeyError):
pass
@@ -212,6 +340,7 @@ class Glm47MoeDetector(BaseFormatDetector):
)
self._tool_call_completed = False # Reset tool call completion status
self._sent_empty_object = False # Reset empty object sent status
self._buffer_arguments = None
def has_tool_call(self, text: str) -> bool:
"""Check if the text contains a glm-4.5 / glm-4.6 format tool call."""
@@ -273,84 +402,15 @@ class Glm47MoeDetector(BaseFormatDetector):
return StreamingParseResult(normal_text=text)
def _get_value_type(self, func_name: str, key: str, tools: List[Tool]) -> str:
"""Get parameter type from tool definition, with fallback to auto-detection.
Args:
func_name: Name of the function
key: Parameter name
tools: List of available tools
Returns:
Type string: 'string', 'number', 'object', 'array', or 'boolean'
"""
arg_type = get_argument_type(func_name, key, tools)
if arg_type:
return arg_type
# Improved auto-detection type from value (best effort)
value_content = self._current_value.strip() if self._current_value else ""
if not value_content:
return "string"
# Try to parse as valid JSON first
try:
parsed = json.loads(value_content)
if isinstance(parsed, dict):
return "object"
elif isinstance(parsed, list):
return "array"
elif isinstance(parsed, bool):
return "boolean"
elif isinstance(parsed, (int, float)):
return "number"
# For string values, check if they look like numbers
elif isinstance(parsed, str):
if parsed.isdigit() or (
parsed.startswith("-") and parsed[1:].isdigit()
):
return "number"
return "string"
except json.JSONDecodeError:
# Not valid JSON, try heuristic detection
first_char = value_content[0] if value_content else ""
if first_char.isdigit() or first_char in ["-", "."]:
return "number"
elif first_char in ["{", "["]:
return "object"
elif first_char in ['"', "'"]:
return "string"
# Default to string (safest fallback)
return "string"
return get_argument_type(func_name, key, tools) or "auto"
def _format_value_complete(self, value: str, value_type: str) -> str:
"""Format complete value based on type.
Args:
value: Raw value string
value_type: Expected type ('string', 'number', 'object')
Returns:
Properly formatted JSON value string
"""
"""Format a value that completed in one chunk; number/integer/auto never
reach here (they are parsed at value close instead)."""
if value_type == "string":
# Ensure proper JSON string formatting with quotes
return json.dumps(value, ensure_ascii=False)
elif value_type == "number":
try:
num = _convert_to_number(value.strip() if value else "")
return str(num)
except (ValueError, AttributeError):
# Fallback to string if not a valid number
logger.warning(
f"Failed to parse '{value}' as number, treating as string"
)
return json.dumps(str(value) if value else "", ensure_ascii=False)
else:
# For object/array types, return as-is (should already be valid JSON)
return value
# object/array/boolean values arrive as JSON already
return value
def _process_xml_to_json_streaming(
self, raw_increment: str, func_name: str, tools: List[Tool]
@@ -397,7 +457,6 @@ class Glm47MoeDetector(BaseFormatDetector):
self._current_value = ""
self._xml_tag_buffer = ""
self._value_started = False
# Determine and cache the value type at the start
self._cached_value_type = self._get_value_type(
func_name, self._current_key, tools
)
@@ -407,11 +466,13 @@ class Glm47MoeDetector(BaseFormatDetector):
final_value = self._xml_tag_buffer[:-12]
self._current_value += final_value
# Use cached value type for consistency
value_type = self._cached_value_type or "string"
if self._value_started:
# Output any remaining content
if value_type in ("auto", "number", "integer"):
parsed = self._parse_argument_pairs(
[(self._current_key, self._current_value)], func_name, tools
)[self._current_key]
json_output += json.dumps(parsed, ensure_ascii=False)
elif self._value_started:
if final_value:
if value_type == "string":
json_output += json.dumps(
@@ -419,11 +480,10 @@ class Glm47MoeDetector(BaseFormatDetector):
)[1:-1]
else:
json_output += final_value
# Always output closing quote for string type when value was started
if value_type == "string":
json_output += '"'
else:
# Value was never started (empty or complete in one chunk)
# Value never started: empty, or completed in one chunk.
json_output += self._format_value_complete(
self._current_value, value_type
)
@@ -432,7 +492,7 @@ class Glm47MoeDetector(BaseFormatDetector):
self._stream_state = StreamState.BETWEEN
self._current_value = ""
self._value_started = False
self._cached_value_type = None # Reset cached type
self._cached_value_type = None
else:
closing_tag = "</arg_value>"
is_potential_closing = len(self._xml_tag_buffer) <= len(
@@ -441,10 +501,12 @@ class Glm47MoeDetector(BaseFormatDetector):
if not is_potential_closing:
content = self._xml_tag_buffer
# Use cached value type for consistency
value_type = self._cached_value_type or "string"
if value_type == "string":
if value_type in ("auto", "number", "integer"):
self._current_value += content
self._xml_tag_buffer = ""
elif value_type == "string":
if not self._value_started:
json_output += '"'
self._value_started = True
@@ -454,15 +516,8 @@ class Glm47MoeDetector(BaseFormatDetector):
]
self._current_value += content
self._xml_tag_buffer = ""
elif value_type == "number":
if content:
if not self._value_started:
self._value_started = True
json_output += content
self._current_value += content
self._xml_tag_buffer = ""
else:
# For object/array types, output as-is
# object/array/boolean values stream through verbatim
if content:
if not self._value_started:
self._value_started = True
@@ -544,13 +599,20 @@ class Glm47MoeDetector(BaseFormatDetector):
"""
current_raw_length = len(func_args_raw)
if self._buffer_arguments is None:
self._buffer_arguments = any(
tool.function.name == func_name
and _needs_argument_context(tool.function.parameters)
for tool in tools
)
if self._buffer_arguments:
return None
if current_raw_length <= self._streamed_raw_length:
return None
# Get new raw XML content
raw_increment = func_args_raw[self._streamed_raw_length :]
# Convert XML to JSON using state machine
json_increment = self._process_xml_to_json_streaming(
raw_increment, func_name, tools
)
@@ -562,7 +624,6 @@ class Glm47MoeDetector(BaseFormatDetector):
if not json_increment:
return None
# Update state
self._last_arguments += json_increment
self.streamed_args_for_tool[self.current_tool_id] += json_increment
@@ -594,9 +655,22 @@ class Glm47MoeDetector(BaseFormatDetector):
"""
calls = []
if self._buffer_arguments:
arguments = self._parse_argument_pairs(
self.func_arg_regex.findall(func_args_raw), func_name, tools
)
serialized = json.dumps(arguments, ensure_ascii=False)
calls.append(
ToolCallItem(
tool_index=self.current_tool_id, name=None, parameters=serialized
)
)
self._last_arguments += serialized
self.streamed_args_for_tool[self.current_tool_id] += serialized
self._sent_empty_object = True
# Handle no-arg function or need to close braces
if self._is_first_param and not self._sent_empty_object:
# No-arg function
calls.append(
ToolCallItem(
tool_index=self.current_tool_id,
@@ -621,7 +695,6 @@ class Glm47MoeDetector(BaseFormatDetector):
self.streamed_args_for_tool[self.current_tool_id] += "}"
self._sent_empty_object = True
# Parse final arguments
if func_args_raw:
try:
pairs = self.func_arg_regex.findall(func_args_raw)
@@ -633,7 +706,6 @@ class Glm47MoeDetector(BaseFormatDetector):
except Exception as e:
logger.debug(f"Failed to parse arguments: {e}", exc_info=True)
# Clean buffer
self._buffer = current_text[match_end_pos:]
# Reset state for next tool call
@@ -786,26 +858,56 @@ class Glm47MoeDetector(BaseFormatDetector):
Dictionary of parsed arguments
"""
arguments = {}
known_arguments = {}
# Values with a certain type are resolved first so completed sibling
# discriminators (e.g. kind=const) can steer root-union branch selection.
if len(pairs) > 1 and any(
tool.function.name == func_name
and _needs_argument_context(tool.function.parameters)
for tool in tools
):
for key, raw in pairs:
key = key.strip()
parsed = self._parse_argument_pairs([(key, raw)], func_name, tools)[key]
native_types = _get_argument_types(func_name, key, tools, value=raw)
if (
isinstance(parsed, str)
or native_types is None
or "string" not in native_types
):
known_arguments[key] = parsed
for arg_key, arg_value in pairs:
arg_key = arg_key.strip()
arg_type = get_argument_type(func_name, arg_key, tools)
parsed_value, is_good_json = parse_arguments(arg_value, arg_type)
arg_type = get_argument_type(func_name, arg_key, tools, known_arguments)
parsed_value, is_good_json = parse_arguments(
arg_value, arg_type or "string"
)
if arg_type == "string":
# Only convert to string if explicitly defined as string type
if isinstance(parsed_value, str):
arguments[arg_key] = parsed_value
elif isinstance(parsed_value, (dict, list)):
# If parsed as dict/list but schema says string, convert to JSON string
arguments[arg_key] = json.dumps(parsed_value, ensure_ascii=False)
else:
arguments[arg_key] = str(parsed_value)
arguments[arg_key] = arg_value
elif arg_type is None:
# If type is not defined, keep the parsed value as-is
# A value whose parsed JSON type is inadmissible stays the raw
# string when the schema allows a string (e.g. enum ["7", 8]).
allowed = _get_argument_types(
func_name, arg_key, tools, known_arguments, parsed_value
)
native = _get_argument_types(
func_name, arg_key, tools, known_arguments, arg_value
)
if (
allowed is not None
and _json_type(parsed_value) not in allowed
and native is not None
and "string" in native
):
parsed_value = arg_value
arguments[arg_key] = parsed_value if is_good_json else arg_value
else:
# For other types (number, object, array, etc.), use parsed value
arguments[arg_key] = parsed_value if is_good_json else arg_value
known_arguments[arg_key] = arguments[arg_key]
return arguments
@@ -0,0 +1,362 @@
import json
import unittest
from sglang.srt.entrypoints.openai.protocol import Function, Tool
from sglang.srt.function_call.glm47_moe_detector import Glm47MoeDetector
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestGlm47SchemaTypes(unittest.TestCase):
def check_arguments(self, schema, expected, raw_values=None):
tools = [
Tool(type="function", function=Function(name="inspect", parameters=schema))
]
pairs = []
for key, value in expected.items():
raw = (raw_values or {}).get(
key, value if isinstance(value, str) else json.dumps(value)
)
pairs.append(f"<arg_key>{key}</arg_key><arg_value>{raw}</arg_value>")
text = "<tool_call>inspect" + "".join(pairs) + "</tool_call>"
for size in (0, 1, 7, len(text)):
with self.subTest(schema=schema, expected=expected, chunk_size=size):
detector = Glm47MoeDetector()
if size == 0:
calls = detector.detect_and_parse(text, tools).calls
else:
calls = []
for offset in range(0, len(text), size):
calls.extend(
detector.parse_streaming_increment(
text[offset : offset + size], tools
).calls
)
actual = json.loads("".join(call.parameters for call in calls))
self.assertEqual(actual, expected)
for key in expected:
self.assertIs(type(actual[key]), type(expected[key]))
def test_field_types(self):
cases = [
({"type": ["integer", "string"]}, "auto"),
({"type": ["string", "integer"]}, 7),
({"type": ["number", "string"]}, "auto"),
({"type": ["object", "string"]}, "auto"),
({"type": ["string", "object"]}, {"ok": False}),
({"type": ["boolean", "string"]}, "auto"),
({"type": ["string", "boolean"]}, False),
({"type": ["string", "null"]}, None),
({"type": ["null"]}, None),
({"type": "null"}, None),
({"enum": ["ok", None]}, None),
({"type": ["string", "null"], "enum": ["ok", None]}, None),
({"anyOf": [{"type": "string", "enum": ["auto"]}, {"type": "integer"}]}, 7),
(
{
"oneOf": [
{"type": "null"},
{"type": "array", "items": {"type": "integer"}},
]
},
[1, 2],
),
({"const": 7}, 7),
({"const": {"ok": False}}, {"ok": False}),
({"const": "null"}, "null"),
({"allOf": [{"type": ["integer", "string"]}, {"type": "integer"}]}, 7),
({}, "123_456"),
({}, {"ok": False}),
]
cases.extend(
({"type": "string"}, value)
for value in (
"null",
"123_456",
"\\d+",
'line one\n"quoted" \\ café',
)
)
for field, value in cases:
self.check_arguments(
{"type": "object", "properties": {"value": field}}, {"value": value}
)
self.check_arguments(
{"type": "object", "properties": {"value": {"type": ["string", "null"]}}},
{"value": "null"},
{"value": '"null"'},
)
def test_local_references_and_cycles(self):
for value, field in [
(7, {"type": "integer"}),
({"ok": False}, {"type": "object"}),
("null", {"type": "string"}),
]:
for defs in ("$defs", "definitions"):
schema = {
defs: {"a/b~c": field},
"type": "object",
"properties": {"value": {"$ref": f"#/{defs}/a~1b~0c"}},
}
self.check_arguments(schema, {"value": value})
self.check_arguments(
{
"$defs": {
"args": {
"type": "object",
"properties": {"value": {"type": "string"}},
}
},
"$ref": "#/$defs/args",
},
{"value": "null"},
)
self.check_arguments(
{
"$defs": {"cycle": {"$ref": "#/$defs/cycle"}},
"properties": {"value": {"$ref": "#/$defs/cycle"}},
},
{"value": "auto"},
)
def test_root_all_of_preserves_value_type(self):
self.check_arguments(
{
"allOf": [
{
"type": "object",
"properties": {
"kind": {"const": "count"},
"value": {"type": "integer"},
},
},
{"required": ["kind", "value"]},
]
},
{"kind": "count", "value": 7},
)
def test_plain_strings_still_stream_before_value_completes(self):
tools = [
Tool(
type="function",
function=Function(
name="inspect",
parameters={
"type": "object",
"properties": {"value": {"type": "string"}},
},
),
)
]
detector = Glm47MoeDetector()
first = detector.parse_streaming_increment(
"<tool_call>inspect<arg_key>value</arg_key><arg_value>hello", tools
)
self.assertIn('"hello', "".join(call.parameters for call in first.calls))
last = detector.parse_streaming_increment(
" world</arg_value></tool_call>", tools
)
self.assertEqual(
json.loads("".join(call.parameters for call in first.calls + last.calls)),
{"value": "hello world"},
)
def test_native_strings_reject_disallowed_json_types_and_enum_values(self):
for field, value in [
({"type": ["string", "object"]}, "7"),
({"type": ["object", "string"]}, "7"),
({"oneOf": [{"type": "string"}, {"type": "object"}]}, "7"),
({"type": ["string", "array"]}, "1e2"),
({"type": ["string", "integer"]}, "true"),
({"type": ["string", "integer"]}, "null"),
({"enum": ["7", None]}, "7"),
({"enum": ["7", 8]}, "7"),
({"enum": ["7", 8]}, 8),
({"enum": ["null", 7]}, "null"),
({"enum": ["7", False]}, "7"),
({"type": ["number", "string"], "enum": [1, "auto"]}, 1),
({"type": ["number", "string"], "enum": [1, "auto"]}, 1.0),
]:
self.check_arguments(
{"type": "object", "properties": {"value": field}}, {"value": value}
)
self.check_arguments(
{
"$defs": {"value": {"type": ["string", "object"]}},
"properties": {"value": {"$ref": "#/$defs/value"}},
},
{"value": "7"},
)
def test_json_looking_string_bytes_are_preserved(self):
for value in ('{"a":1}', '{ "a" : true }'):
for field in ({"type": "string"}, {"const": value}):
self.check_arguments({"properties": {"value": field}}, {"value": value})
def test_nonstream_numeric_strings_with_integer_enum(self):
tools = [
Tool(
type="function",
function=Function(
name="inspect",
parameters={
"properties": {"value": {"type": "number", "enum": [1, 2]}}
},
),
)
]
for raw in ('"1"', r"\"1\""):
text = (
"<tool_call>inspect<arg_key>value</arg_key><arg_value>"
+ raw
+ "</arg_value></tool_call>"
)
calls = Glm47MoeDetector().detect_and_parse(text, tools).calls
self.assertEqual(json.loads(calls[0].parameters), {"value": 1})
def test_numeric_values_are_serialized_after_the_value_closes(self):
for value_type in ("number", "integer"):
schema = {"properties": {"value": {"type": value_type}}}
for raw, expected in (
("123abc", "123abc"),
('"123"', 123),
(r"\"123\"", 123),
("-12", -12),
("1.25e2", 125.0),
("", ""),
):
self.check_arguments(schema, {"value": expected}, {"value": raw})
self.check_arguments(
{"properties": {"value": {"type": "number"}}},
{"value": 1.5},
{"value": '"1.5"'},
)
for raw in ('"1.5"', '"1e2"'):
self.check_arguments(
{"properties": {"value": {"type": "integer"}}},
{"value": json.loads(raw)},
{"value": raw},
)
def test_numeric_prefix_does_not_commit_to_invalid_json(self):
tools = [
Tool(
type="function",
function=Function(
name="inspect",
parameters={"properties": {"value": {"type": "number"}}},
),
)
]
detector = Glm47MoeDetector()
prefix = detector.parse_streaming_increment(
"<tool_call>inspect<arg_key>value</arg_key><arg_value>123", tools
).calls
self.assertNotIn("123", "".join(call.parameters for call in prefix))
suffix = detector.parse_streaming_increment(
"abc</arg_value></tool_call>", tools
).calls
self.assertEqual(
json.loads("".join(call.parameters for call in prefix + suffix)),
{"value": "123abc"},
)
def test_completed_siblings_disambiguate_root_union(self):
for keyword in ("oneOf", "anyOf"):
for discriminator in ("const", "enum"):
branches = [
{
"properties": {
"kind": {
discriminator: [kind]
if discriminator == "enum"
else kind
},
"value": {"type": value_type},
}
}
for kind, value_type in (("text", "string"), ("count", "integer"))
]
for ordered in (branches, list(reversed(branches))):
for schema in (
{keyword: ordered},
{keyword: [{"allOf": [branch]} for branch in ordered]},
{
"$defs": {
str(i): branch for i, branch in enumerate(ordered)
},
keyword: [{"$ref": f"#/$defs/{i}"} for i in range(2)],
},
):
self.check_arguments(schema, {"kind": "text", "value": "7"})
self.check_arguments(schema, {"kind": "count", "value": 7})
self.check_arguments(schema, {"value": "7", "kind": "text"})
self.check_arguments(schema, {"value": 7, "kind": "count"})
def test_completed_arguments_reset_between_tool_calls(self):
schema = {
"oneOf": [
{
"properties": {
"kind": {"const": "text"},
"value": {"type": "string"},
}
},
{
"properties": {
"kind": {"const": "count"},
"value": {"type": "integer"},
}
},
]
}
tools = [
Tool(type="function", function=Function(name="inspect", parameters=schema))
]
tools.append(
Tool(
type="function",
function=Function(
name="plain",
parameters={"properties": {"value": {"type": "string"}}},
),
)
)
detector = Glm47MoeDetector()
outputs = {}
for kind, discriminator_first in (
("text", True),
("count", False),
("text", False),
):
tag = f"<arg_key>kind</arg_key><arg_value>{kind}</arg_value>"
value = "<arg_key>value</arg_key><arg_value>7</arg_value>"
text = (
"<tool_call>inspect"
+ (tag + value if discriminator_first else value + tag)
+ "</tool_call>"
)
for char in text:
for call in detector.parse_streaming_increment(char, tools).calls:
outputs[call.tool_index] = (
outputs.get(call.tool_index, "") + call.parameters
)
self.assertEqual(json.loads(outputs[0]), {"kind": "text", "value": "7"})
self.assertEqual(json.loads(outputs[1]), {"kind": "count", "value": 7})
self.assertEqual(json.loads(outputs[2]), {"kind": "text", "value": "7"})
first = detector.parse_streaming_increment(
"<tool_call>plain<arg_key>value</arg_key><arg_value>hello", tools
)
self.assertIn('"hello', "".join(call.parameters for call in first.calls))
last = detector.parse_streaming_increment("</arg_value></tool_call>", tools)
self.assertEqual(
json.loads("".join(call.parameters for call in first.calls + last.calls)),
{"value": "hello"},
)
if __name__ == "__main__":
unittest.main()