add pydantic model from json

This commit is contained in:
Michael Poluektov 2024-08-19 16:27:38 +01:00
parent 5edc211392
commit 9d7037b730
2 changed files with 107 additions and 0 deletions

104
backend/utils/schemas.py Normal file
View File

@ -0,0 +1,104 @@
from pydantic import BaseModel, Field, create_model
from typing import Any, Optional, Type
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":
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
else:
raise ValueError(f"Unsupported JSON schema type: {type_}")

View File

@ -6,6 +6,8 @@ from apps.webui.models.tools import Tools
from apps.webui.models.users import UserModel
from apps.webui.utils import load_toolkit_module_by_id
from utils.schemas import json_schema_to_model
log = logging.getLogger(__name__)
@ -70,6 +72,7 @@ def get_tools(
"toolkit_id": tool_id,
"callable": callable,
"spec": spec,
"pydantic_model": json_schema_to_model(spec),
"file_handler": hasattr(module, "file_handler") and module.file_handler,
"citation": hasattr(module, "citation") and module.citation,
}