Spaces:
Running
Running
File size: 3,294 Bytes
dd8990d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
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_}")
|