Spaces:
Running
Running
from ast import literal_eval | |
from typing import Any, Literal, Optional, Type | |
from pydantic import BaseModel, Field, create_model | |
def json_schema_to_model(tool_dict: dict[str, Any]) -> Type[BaseModel]: | |
""" | |
Converts a JSON schema to a Pydantic BaseModel class. | |
Args: | |
json_schema: The JSON schema to convert. | |
Returns: | |
A Pydantic BaseModel class. | |
""" | |
# Extract the model name from the schema title. | |
model_name = tool_dict["name"] | |
schema = tool_dict["parameters"] | |
# Extract the field definitions from the schema properties. | |
field_definitions = { | |
name: json_schema_to_pydantic_field(name, prop, schema.get("required", [])) | |
for name, prop in schema.get("properties", {}).items() | |
} | |
# Create the BaseModel class using create_model(). | |
return create_model(model_name, **field_definitions) | |
def json_schema_to_pydantic_field( | |
name: str, json_schema: dict[str, Any], required: list[str] | |
) -> Any: | |
""" | |
Converts a JSON schema property to a Pydantic field definition. | |
Args: | |
name: The field name. | |
json_schema: The JSON schema property. | |
Returns: | |
A Pydantic field definition. | |
""" | |
# Get the field type. | |
type_ = json_schema_to_pydantic_type(json_schema) | |
# Get the field description. | |
description = json_schema.get("description") | |
# Get the field examples. | |
examples = json_schema.get("examples") | |
# Create a Field object with the type, description, and examples. | |
# The 'required' flag will be set later when creating the model. | |
return ( | |
type_, | |
Field( | |
description=description, | |
examples=examples, | |
default=... if name in required else None, | |
), | |
) | |
def json_schema_to_pydantic_type(json_schema: dict[str, Any]) -> Any: | |
""" | |
Converts a JSON schema type to a Pydantic type. | |
Args: | |
json_schema: The JSON schema to convert. | |
Returns: | |
A Pydantic type. | |
""" | |
type_ = json_schema.get("type") | |
if type_ == "string" or type_ == "str": | |
return str | |
elif type_ == "integer" or type_ == "int": | |
return int | |
elif type_ == "number" or type_ == "float": | |
return float | |
elif type_ == "boolean" or type_ == "bool": | |
return bool | |
elif type_ == "array" or type_ == "list": | |
items_schema = json_schema.get("items") | |
if items_schema: | |
item_type = json_schema_to_pydantic_type(items_schema) | |
return list[item_type] | |
else: | |
return list | |
elif type_ == "object": | |
# Handle nested models. | |
properties = json_schema.get("properties") | |
if properties: | |
nested_model = json_schema_to_model(json_schema) | |
return nested_model | |
else: | |
return dict | |
elif type_ == "null": | |
return Optional[Any] # Use Optional[Any] for nullable fields | |
elif type_ == "literal": | |
return Literal[literal_eval(json_schema.get("enum"))] | |
elif type_ == "optional": | |
inner_schema = json_schema.get("items", {"type": "string"}) | |
inner_type = json_schema_to_pydantic_type(inner_schema) | |
return Optional[inner_type] | |
else: | |
raise ValueError(f"Unsupported JSON schema type: {type_}") | |