chore: import upstream snapshot with attribution

python/sglang/srt/function_call/utils.py

@@ -0,0 +1,415 @@
+from json import JSONDecodeError, JSONDecoder
+from json.decoder import WHITESPACE
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import orjson
+import partial_json_parser
+from partial_json_parser.core.options import Allow
+
+from sglang.srt.entrypoints.openai.protocol import Tool, ToolChoice
+
+_STANDARD_JSON_SCHEMA_TYPES = {
+    "null",
+    "boolean",
+    "object",
+    "array",
+    "number",
+    "string",
+    "integer",
+}
+
+# Non-standard ``type`` values commonly emitted by DB/ORM-driven tool-schema
+# generators. Mapped to the closest JSON Schema 2020-12 primitive so that
+# ``Draft202012Validator.check_schema`` does not reject an otherwise-usable
+# tool definition.
+_JSON_SCHEMA_TYPE_ALIASES: Dict[str, str] = {
+    "str": "string",
+    "text": "string",
+    "varchar": "string",
+    "char": "string",
+    "enum": "string",
+    "uuid": "string",
+    "date": "string",
+    "datetime": "string",
+    "time": "string",
+    "timestamp": "string",
+    "binary": "string",
+    "blob": "string",
+    "bytea": "string",
+    "bytes": "string",
+    "varbinary": "string",
+    "bool": "boolean",
+    "bigint": "integer",
+    "smallint": "integer",
+    "tinyint": "integer",
+    "double": "number",
+    "decimal": "number",
+    "real": "number",
+    "numeric": "number",
+    "arr": "array",
+    "tuple": "array",
+    "set": "array",
+    "map": "object",
+}
+
+# Prefix-based matching so that parameterised names like ``int32`` /
+# ``float64`` / ``list[str]`` / ``dict[str, int]`` resolve. A prefix only
+# matches when it spans the entire token or is followed by a non-identifier
+# char, so "int" does not swallow "internal" and "list" does not swallow
+# "list_price".
+_PREFIX_BOUNDARY_CHARS = frozenset("0123456789[<( \t")
+_PREFIX_RULES: Tuple[Tuple[Tuple[str, ...], str], ...] = (
+    (("int", "uint", "long", "short", "unsigned"), "integer"),
+    (("num", "float"), "number"),
+    (("list",), "array"),
+    (("dict",), "object"),
+)
+
+
+def _matches_type_prefix(base: str, prefixes: Tuple[str, ...]) -> bool:
+    for p in prefixes:
+        if base == p:
+            return True
+        if (
+            len(base) > len(p)
+            and base.startswith(p)
+            and base[len(p)] in _PREFIX_BOUNDARY_CHARS
+        ):
+            return True
+    return False
+
+
+def _normalize_single_type(raw: Any) -> Any:
+    if not isinstance(raw, str):
+        return raw
+    if raw in _STANDARD_JSON_SCHEMA_TYPES:
+        return raw
+    # ``split("(", 1)[0]`` strips parenthesized params like ``varchar(255)``
+    # or ``decimal(10,2)`` without the overhead of a regex per call.
+    base = raw.split("(", 1)[0].strip().lower()
+    if base in _STANDARD_JSON_SCHEMA_TYPES:
+        return base
+    mapped = _JSON_SCHEMA_TYPE_ALIASES.get(base)
+    if mapped is not None:
+        return mapped
+    for prefixes, target in _PREFIX_RULES:
+        if _matches_type_prefix(base, prefixes):
+            return target
+    return raw
+
+
+def _normalize_type_list(raw_items: List[Any]) -> List[Any]:
+    normalized_items: List[Any] = []
+    for item in raw_items:
+        normalized_item = _normalize_single_type(item)
+        if normalized_item not in normalized_items:
+            normalized_items.append(normalized_item)
+    return normalized_items
+
+
+def normalize_json_schema_types(schema: Any) -> None:
+    """
+    Walk a JSON Schema in place and rewrite non-standard ``"type"`` values
+    (e.g. ``"varchar"``, ``"enum"``, ``"int"``) to their standard JSON Schema
+    equivalents.
+
+    Acts as a compatibility layer for tool ``parameters`` schemas exported
+    from database / ORM tooling, which often uses DB type names rather than
+    JSON Schema types. Unknown types are left untouched so that downstream
+    validation can still surface genuine errors.
+
+    Mutates the input dict in place; the rewritten schema is also what gets
+    rendered into the model prompt, so e.g. a user-supplied ``"varchar"``
+    reaches the model as ``"string"``. ``$ref`` values are not resolved;
+    callers pass tree-shaped schemas (HTTP JSON input is always a tree).
+    """
+    if isinstance(schema, list):
+        for item in schema:
+            normalize_json_schema_types(item)
+        return
+    if not isinstance(schema, dict):
+        return
+
+    if "type" in schema:
+        t = schema["type"]
+        if isinstance(t, str):
+            schema["type"] = _normalize_single_type(t)
+        elif isinstance(t, list):
+            schema["type"] = _normalize_type_list(t)
+
+    for key in (
+        "properties",
+        "patternProperties",
+        "$defs",
+        "definitions",
+        "dependentSchemas",
+    ):
+        nested = schema.get(key)
+        if isinstance(nested, dict):
+            for v in nested.values():
+                normalize_json_schema_types(v)
+
+    for key in ("anyOf", "oneOf", "allOf", "prefixItems"):
+        nested = schema.get(key)
+        if isinstance(nested, list):
+            for v in nested:
+                normalize_json_schema_types(v)
+
+    for key in (
+        "items",
+        "additionalProperties",
+        "not",
+        "if",
+        "then",
+        "else",
+        "contains",
+        "propertyNames",
+        "unevaluatedItems",
+        "unevaluatedProperties",
+    ):
+        if key in schema:
+            normalize_json_schema_types(schema[key])
+
+
+def _find_common_prefix(s1: str, s2: str) -> str:
+    prefix = ""
+    min_length = min(len(s1), len(s2))
+    for i in range(0, min_length):
+        if s1[i] == s2[i]:
+            prefix += s1[i]
+        else:
+            break
+    return prefix
+
+
+def _partial_json_loads(input_str: str, flags: Allow) -> Tuple[Any, int]:
+    """
+    Parse incomplete or partial JSON strings commonly encountered during streaming.
+
+    Args:
+        input_str (str): The potentially incomplete JSON string to parse.
+        flags (Allow): Bitwise flags controlling what types of partial data are allowed.
+            Common flags include:
+            - Allow.STR: Allow partial strings (e.g., '"hello wo' -> 'hello wo')
+            - Allow.OBJ: Allow partial objects (e.g., '{"key":' -> {'key': None})
+            - Allow.ARR: Allow partial arrays (e.g., '[1, 2,' -> [1, 2])
+            - Allow.ALL: Allow all types of partial data
+
+    Returns:
+        Tuple[Any, int]: A tuple containing:
+            - parsed_object: The Python object parsed from the JSON
+            - consumed_length: Number of characters consumed from input_str
+    """
+    try:
+        return (partial_json_parser.loads(input_str, flags), len(input_str))
+    except (JSONDecodeError, IndexError) as e:
+        msg = getattr(e, "msg", str(e))
+        if "Extra data" in msg or "pop from empty list" in msg:
+            start = WHITESPACE.match(input_str, 0).end()
+            obj, end = JSONDecoder().raw_decode(input_str, start)
+            return obj, end
+        raise
+
+
+def _is_complete_json(input_str: str) -> bool:
+    try:
+        orjson.loads(input_str)
+        return True
+    except JSONDecodeError:
+        return False
+
+
+def _get_tool_schema_defs(tools: List[Tool]) -> dict:
+    """
+    Get consolidated $defs from all tools, validating for conflicts.
+
+    Args:
+        tools: List of tools to process
+
+    Returns:
+        Dictionary of consolidated $defs from all tools
+
+    Raises:
+        ValueError: If conflicting $defs are found
+    """
+    all_defs = {}
+    for tool in tools:
+        if tool.function.parameters is None:
+            continue
+        defs = tool.function.parameters.get("$defs", {})
+        for def_name, def_schema in defs.items():
+            if def_name in all_defs and all_defs[def_name] != def_schema:
+                raise ValueError(
+                    f"Tool definition '{def_name}' has "
+                    "multiple schemas, which is not "
+                    "supported."
+                )
+            else:
+                all_defs[def_name] = def_schema
+    return all_defs
+
+
+def _get_tool_schema(tool: Tool) -> dict:
+    return {
+        "properties": {
+            "name": {"type": "string", "enum": [tool.function.name]},
+            "parameters": (
+                tool.function.parameters
+                if tool.function.parameters
+                else {"type": "object", "properties": {}}
+            ),
+        },
+        "required": ["name", "parameters"],
+    }
+
+
+def infer_type_from_json_schema(schema: Dict[str, Any]) -> Optional[str]:
+    """
+    Infer the primary type of a parameter from JSON Schema.
+
+    Supports complex JSON Schema structures 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:
+        schema: JSON Schema definition
+
+    Returns:
+        Inferred type ('string', 'number', 'object', 'array', etc.) or None
+    """
+    if not isinstance(schema, dict):
+        return None
+
+    # Priority 1: Direct type field (including type arrays)
+    if "type" in schema:
+        type_value = schema["type"]
+        if isinstance(type_value, str):
+            return type_value
+        elif isinstance(type_value, list) and type_value:
+            # Handle type arrays: return first non-null type
+            non_null_types = [t for t in type_value if t != "null"]
+            if non_null_types:
+                return non_null_types[0]
+            return "string"  # If only null, default to string
+
+    # Priority 2: Handle anyOf/oneOf
+    if "anyOf" in schema or "oneOf" in schema:
+        schemas = schema.get("anyOf") or schema.get("oneOf")
+        types = []
+
+        if isinstance(schemas, list):
+            for sub_schema in schemas:
+                inferred_type = infer_type_from_json_schema(sub_schema)
+                if inferred_type:
+                    types.append(inferred_type)
+
+            if types:
+                # If all types are the same, return unified type
+                if len(set(types)) == 1:
+                    return types[0]
+                # When types differ, prioritize string (safest)
+                if "string" in types:
+                    return "string"
+                # Otherwise return first type
+                return types[0]
+
+    # Priority 3: Handle enum (infer type from enum values)
+    if "enum" in schema and isinstance(schema["enum"], list):
+        if not schema["enum"]:
+            return "string"
+
+        # Infer type from enum values
+        enum_types = set()
+        for value in schema["enum"]:
+            if value is None:
+                enum_types.add("null")
+            elif isinstance(value, bool):
+                enum_types.add("boolean")
+            elif isinstance(value, int):
+                enum_types.add("integer")
+            elif isinstance(value, float):
+                enum_types.add("number")
+            elif isinstance(value, str):
+                enum_types.add("string")
+            elif isinstance(value, list):
+                enum_types.add("array")
+            elif isinstance(value, dict):
+                enum_types.add("object")
+
+        # If type is uniform, return that type
+        if len(enum_types) == 1:
+            return enum_types.pop()
+        # Mixed types, prioritize string
+        return "string"
+
+    # Priority 4: Handle allOf (must satisfy all types)
+    if "allOf" in schema and isinstance(schema["allOf"], list):
+        schemas = schema["allOf"]
+        for sub_schema in schemas:
+            inferred_type = infer_type_from_json_schema(sub_schema)
+            if inferred_type and inferred_type != "string":
+                return inferred_type
+        return "string"
+
+    # Priority 5: Infer object type
+    if "properties" in schema:
+        return "object"
+
+    # Priority 6: Infer array type
+    if "items" in schema:
+        return "array"
+
+    return None
+
+
+def get_json_schema_constraint(
+    tools: List[Tool],
+    tool_choice: Union[ToolChoice, Literal["required"]],
+    parallel_tool_calls: bool = True,
+) -> Optional[dict]:
+    """
+    Get the JSON schema constraint for the specified tool choice.
+
+    Args:
+        tool_choice: The tool choice specification
+        parallel_tool_calls: If False, constrain to exactly one tool call (maxItems=1)
+
+    Returns:
+        JSON schema dict, or None if no valid tools found
+    """
+
+    if isinstance(tool_choice, ToolChoice):
+        # For specific function choice, return the user's parameters schema directly
+        fn_name = tool_choice.function.name
+        for tool in tools:
+            if tool.function.name == fn_name:
+                schema = {
+                    "type": "array",
+                    "minItems": 1,
+                    "items": _get_tool_schema(tool),
+                }
+                if not parallel_tool_calls:
+                    schema["maxItems"] = 1
+                return schema
+        return None
+    elif tool_choice == "required":
+        json_schema = {
+            "type": "array",
+            "minItems": 1,
+            "items": {
+                "type": "object",
+                "anyOf": [_get_tool_schema(tool) for tool in tools],
+            },
+        }
+        if not parallel_tool_calls:
+            json_schema["maxItems"] = 1
+        json_schema_defs = _get_tool_schema_defs(tools)
+        if json_schema_defs:
+            json_schema["$defs"] = json_schema_defs
+        return json_schema
+
+    return None