247 lines
8.0 KiB
Python
247 lines
8.0 KiB
Python
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
|
|
|
|
|
|
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"]]
|
|
) -> Optional[dict]:
|
|
"""
|
|
Get the JSON schema constraint for the specified tool choice.
|
|
|
|
Args:
|
|
tool_choice: The tool choice specification
|
|
|
|
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:
|
|
return {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"maxItems": 1,
|
|
"items": _get_tool_schema(tool),
|
|
}
|
|
return None
|
|
elif tool_choice == "required":
|
|
json_schema = {
|
|
"type": "array",
|
|
"minItems": 1,
|
|
"items": {
|
|
"type": "object",
|
|
"anyOf": [_get_tool_schema(tool) for tool in tools],
|
|
},
|
|
}
|
|
json_schema_defs = _get_tool_schema_defs(tools)
|
|
if json_schema_defs:
|
|
json_schema["$defs"] = json_schema_defs
|
|
return json_schema
|
|
|
|
return None
|