open-webui/backend/open_webui/utils/schemas.py

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_}")