mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-15 23:31:05 +00:00
227 lines
8 KiB
Python
227 lines
8 KiB
Python
"""MCP Tool Schema definitions for recursive JSON Schema representation."""
|
|
|
|
import json
|
|
from typing import Any, Union, Optional
|
|
|
|
from mcp.types import Tool
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator, field_validator
|
|
|
|
from ..enumeration.json_schema_enum import JsonSchemaEnum
|
|
|
|
|
|
class ToolAttr(BaseModel):
|
|
"""Recursive model representing JSON Schema attributes for tool parameters."""
|
|
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
type: str = Field(default=str(JsonSchemaEnum.STRING), description="The data type of the attribute")
|
|
description: Optional[str] = Field(default=None, description="Description of the attribute")
|
|
required: Optional[list[str]] = Field(default=None, description="Required property names for object types")
|
|
properties: Optional[dict[str, "ToolAttr"]] = Field(default=None, description="Child properties for objects")
|
|
items: Optional[Union[dict[str, Any], "ToolAttr"]] = Field(default=None, description="Schema for array items")
|
|
enum: Optional[list[str]] = Field(default=None, description="Allowed values for the attribute")
|
|
|
|
@field_validator("type")
|
|
@classmethod
|
|
def validate_type_is_valid_enum(cls, v: str) -> str:
|
|
"""Validates that the provided type string exists within JsonSchemaEnum values."""
|
|
valid_types = [str(e) for e in JsonSchemaEnum]
|
|
|
|
if v not in valid_types:
|
|
raise ValueError(f"Invalid type: '{v}'. Must be one of {valid_types}")
|
|
return v
|
|
|
|
def simple_input_dump(self) -> dict:
|
|
"""Serializes the attribute into a standard JSON Schema dictionary."""
|
|
res: dict = {"type": self.type}
|
|
if self.description:
|
|
res["description"] = self.description
|
|
if self.enum:
|
|
res["enum"] = self.enum
|
|
|
|
if self.type == "object" and self.properties is not None:
|
|
res["properties"] = {
|
|
k: v.simple_input_dump() if isinstance(v, ToolAttr) else v for k, v in self.properties.items()
|
|
}
|
|
if self.required is not None:
|
|
res["required"] = self.required
|
|
|
|
if self.type == "array" and self.items is not None:
|
|
res["items"] = self.items.simple_input_dump() if isinstance(self.items, ToolAttr) else self.items
|
|
|
|
return res
|
|
|
|
|
|
# Enable recursive type resolution
|
|
ToolAttr.model_rebuild()
|
|
|
|
|
|
class ToolCall(BaseModel):
|
|
"""
|
|
Model representing a tool definition and its call structure.
|
|
Supports parsing from standard JSON Schema formats and converting to MCP Tool objects.
|
|
input:
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_weather",
|
|
"description": "It is very useful when you want to check the weather of a specified city.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.",
|
|
}
|
|
},
|
|
"required": ["location"]
|
|
}
|
|
}
|
|
}
|
|
output:
|
|
{
|
|
"index": 0,
|
|
"id": "call_6596dafa2a6a46f7a217da",
|
|
"function": {
|
|
"arguments": "{\"location\": \"Beijing\"}",
|
|
"name": "get_current_weather"
|
|
},
|
|
"type": "function",
|
|
}
|
|
"""
|
|
|
|
index: int = 0
|
|
id: str = ""
|
|
type: str = "function"
|
|
name: str = ""
|
|
description: str = ""
|
|
|
|
arguments: str = Field(default="", description="JSON string of tool execution arguments")
|
|
|
|
parameters: ToolAttr = Field(
|
|
default_factory=lambda: ToolAttr(type="object", properties={}, required=[]),
|
|
description="Specification for input parameters",
|
|
)
|
|
|
|
@model_validator(mode="before")
|
|
@classmethod
|
|
def init_tool_call(cls, data: dict) -> dict:
|
|
"""Initializes the model by parsing tool-specific body data."""
|
|
data = data.copy()
|
|
t_type = data.get("type", "function")
|
|
body = data.get(t_type, {})
|
|
|
|
# Extract basic metadata
|
|
data["name"] = body.get("name", data.get("name", ""))
|
|
data["arguments"] = body.get("arguments", data.get("arguments", ""))
|
|
data["description"] = body.get("description", data.get("description", ""))
|
|
|
|
# Handle parameters mapping
|
|
if "parameters" in body:
|
|
params = body["parameters"]
|
|
# If parameters is already a dict, ensure it matches ToolAttr structure
|
|
if isinstance(params, dict):
|
|
data["parameters"] = ToolAttr(**params)
|
|
|
|
# Handle output mapping (if provided in source)
|
|
if "output" in body and isinstance(body["output"], dict):
|
|
data["output"] = ToolAttr(**body["output"])
|
|
|
|
return data
|
|
|
|
def simple_input_dump(self, as_dict: bool = True) -> dict | str:
|
|
"""Returns a standardized tool definition dictionary or JSON string.
|
|
|
|
Args:
|
|
as_dict: If True, returns dict; if False, returns JSON string.
|
|
"""
|
|
result = {
|
|
"type": self.type,
|
|
self.type: {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": self.parameters.simple_input_dump(),
|
|
},
|
|
}
|
|
return result if as_dict else json.dumps(result, ensure_ascii=False)
|
|
|
|
def simple_output_dump(self, as_dict: bool = True, enable_argument_dict: bool = False) -> dict | str:
|
|
"""Convert ToolCall to output format dictionary or JSON string for API responses."""
|
|
result = {
|
|
"index": self.index,
|
|
"id": self.id,
|
|
self.type: {
|
|
"arguments": self.argument_dict if enable_argument_dict else self.arguments,
|
|
"name": self.name,
|
|
},
|
|
"type": self.type,
|
|
}
|
|
return result if as_dict else json.dumps(result, ensure_ascii=False)
|
|
|
|
@property
|
|
def argument_dict(self) -> dict:
|
|
"""Parse and return arguments as a dictionary."""
|
|
return json.loads(self.arguments)
|
|
|
|
def check_argument(self) -> bool:
|
|
"""Check if arguments can be parsed as valid JSON."""
|
|
try:
|
|
_ = self.argument_dict
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def sanitize_and_check_argument(self) -> bool:
|
|
"""
|
|
Attempt to sanitize and validate arguments JSON.
|
|
Common issues from LLM streaming:
|
|
- Extra closing brackets: }]}] -> }]
|
|
- Missing closing brackets
|
|
- Trailing commas
|
|
"""
|
|
if not self.arguments or not self.arguments.strip():
|
|
return False
|
|
|
|
try:
|
|
# First try parsing as-is
|
|
_ = json.loads(self.arguments)
|
|
return True
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Try to fix common issues
|
|
sanitized = self.arguments.strip()
|
|
|
|
# Remove trailing extra brackets/braces
|
|
# Pattern: if it ends with multiple closing chars, try removing extras
|
|
while len(sanitized) > 1:
|
|
try:
|
|
json.loads(sanitized)
|
|
self.arguments = sanitized # Update with sanitized version
|
|
return True
|
|
except json.JSONDecodeError:
|
|
# Try removing last character
|
|
if sanitized[-1] in "]}":
|
|
sanitized = sanitized[:-1].rstrip()
|
|
else:
|
|
break
|
|
|
|
return False
|
|
|
|
@classmethod
|
|
def from_mcp_tool(cls, tool: Tool) -> "ToolCall":
|
|
"""Creates a ToolCall instance from an MCP Tool object."""
|
|
# MCP Tool inputSchema maps directly to our parameters ToolAttr
|
|
return cls(
|
|
name=tool.name,
|
|
description=tool.description or "",
|
|
parameters=ToolAttr(**tool.inputSchema),
|
|
)
|
|
|
|
def to_mcp_tool(self) -> Tool:
|
|
"""Converts the instance back into an MCP Tool object."""
|
|
return Tool(
|
|
name=self.name,
|
|
description=self.description,
|
|
inputSchema=self.parameters.simple_input_dump(),
|
|
)
|