Create the json schema conversion methods

Author: RobinPicard

outlines/types/dsl.py

@@ -22,6 +22,7 @@
 from typing import (
     Any,
     List,
+    Literal,
     Optional as OptionalType,
     Union,
     get_args,
@@ -41,6 +42,11 @@
 
 import outlines.types as types
 from outlines import grammars
+from outlines.types.json_schema_utils import (
+    json_schema_dict_to_pydantic,
+    json_schema_dict_to_typeddict,
+    json_schema_dict_to_dataclass,
+)
 from outlines.types.utils import (
     get_schema_from_signature,
     is_int,
@@ -282,6 +288,9 @@ class JsonSchema(Term):
     genSON schema builder.
 
     """
+    schema: str
+    whitespace_pattern: OptionalType[str]
+
     def __init__(
         self,
         schema: Union[
@@ -329,6 +338,109 @@ def __init__(
     def __post_init__(self):
         jsonschema.Draft7Validator.check_schema(json.loads(self.schema))
 
+    @classmethod
+    def is_json_schema(cls, obj: Any) -> bool:
+        """Check if the object provided is a JSON schema type.
+
+        Parameters
+        ----------
+        obj: Any
+            The object to check
+
+        Returns
+        -------
+        bool
+            True if the object is a JSON schema type, False otherwise
+
+        """
+        return (
+            isinstance(obj, cls)
+            or is_pydantic_model(obj)
+            or is_typed_dict(obj)
+            or is_dataclass(obj)
+            or is_genson_schema_builder(obj)
+        )
+
+    @classmethod
+    def convert_to(
+        cls,
+        schema: Union[
+            "JsonSchema",
+            type[BaseModel],
+            _TypedDictMeta,
+            type,
+            SchemaBuilder,
+        ],
+        target_types: List[Literal[
+            "str",
+            "dict",
+            "pydantic",
+            "typeddict",
+            "dataclass",
+            "genson",
+        ]],
+    ) -> Union[str, dict, type[BaseModel], _TypedDictMeta, type, SchemaBuilder]:
+        """Convert a JSON schema type to a different JSON schema type.
+
+        If the schema provided is already of a type in the target_types, return
+        it unchanged.
+
+        Parameters
+        ----------
+        schema: Union[JsonSchema, type[BaseModel], _TypedDictMeta, type, SchemaBuilder]
+            The schema to convert
+        target_types: List[Literal["str", "dict", "pydantic", "typeddict", "dataclass", "genson"]]
+            The target types to convert to
+
+        """
+        # If the schema provided is already of a type in the target_types,
+        # just return it
+        if isinstance(schema, cls):
+            if "str" in target_types:
+                return schema.schema
+            elif "dict" in target_types:
+                return json.loads(schema.schema)
+        elif is_pydantic_model(schema) and "pydantic" in target_types:
+            return schema
+        elif is_typed_dict(schema) and "typeddict" in target_types:
+            return schema
+        elif is_dataclass(schema) and "dataclass" in target_types:
+            return schema
+        elif is_genson_schema_builder(schema) and "genson" in target_types:
+            return schema
+
+        # Convert the schema to a JSON schema string/dict
+        if isinstance(schema, cls):
+            schema_str = schema.schema
+        else:
+            schema_str = cls(schema).schema
+        schema_dict = json.loads(schema_str)
+
+        for target_type in target_types:
+            try:
+                # Convert the JSON schema string to the target type
+                if target_type == "str":
+                    return schema_str
+                elif target_type == "dict":
+                    return schema_dict
+                elif target_type == "pydantic":
+                    return json_schema_dict_to_pydantic(schema_dict)
+                elif target_type == "typeddict":
+                    return json_schema_dict_to_typeddict(schema_dict)
+                elif target_type == "dataclass":
+                    return json_schema_dict_to_dataclass(schema_dict)
+                # No conversion available for genson
+            except Exception as e:  # pragma: no cover
+                warnings.warn(
+                    f"Cannot convert schema type {type(schema)} to {target_type}: {e}"
+                )
+                continue
+
+        raise ValueError(
+            f"Cannot convert schema type {type(schema)} to any of the target "
+            f"types {target_types}"
+        )
+
     def _display_node(self) -> str:
         return f"JsonSchema('{self.schema}')"
 

outlines/types/json_schema_utils.py

@@ -1 +1,175 @@
-"""Utilities for handling JSON schema compatibility."""
+"""Convert JSON Schema dicts to Python types."""
+
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, create_model
+
+if sys.version_info >= (3, 12):  # pragma: no cover
+    from typing import _TypedDictMeta, TypedDict  # type: ignore
+else:  # pragma: no cover
+    from typing_extensions import _TypedDictMeta, TypedDict  # type: ignore
+
+
+def schema_type_to_python(
+    schema: dict,
+    caller_target_type: Literal["pydantic", "typeddict", "dataclass"]
+) -> Any:
+    """Get a Python type from a JSON Schema dict.
+
+    Parameters
+    ----------
+    schema: dict
+        The JSON Schema dict to convert to a Python type
+    caller_target_type: Literal["pydantic", "typeddict", "dataclass"]
+        The type of the caller
+
+    Returns
+    -------
+    Any
+        The Python type
+
+    """
+    if "enum" in schema:
+        values = schema["enum"]
+        return Literal[tuple(values)]
+
+    t = schema.get("type")
+
+    if t == "string":
+        return str
+    elif t == "integer":
+        return int
+    elif t == "number":
+        return float
+    elif t == "boolean":
+        return bool
+    elif t == "array":
+        items = schema.get("items", {})
+        if items:
+            item_type = schema_type_to_python(items, caller_target_type)
+        else:
+            item_type = Any
+        return List[item_type]  # type: ignore
+    elif t == "object":
+        name = schema.get("title")
+        if caller_target_type == "pydantic":
+            return json_schema_dict_to_pydantic(schema, name)
+        elif caller_target_type == "typeddict":
+            return json_schema_dict_to_typeddict(schema, name)
+        elif caller_target_type == "dataclass":
+            return json_schema_dict_to_dataclass(schema, name)
+
+    return Any
+
+
+def json_schema_dict_to_typeddict(
+    schema: dict,
+    name: Optional[str] = None
+) -> _TypedDictMeta:
+    """Convert a JSON Schema dict into a TypedDict class.
+
+    Parameters
+    ----------
+    schema: dict
+        The JSON Schema dict to convert to a TypedDict
+    name: Optional[str]
+        The name of the TypedDict
+
+    Returns
+    -------
+    _TypedDictMeta
+        The TypedDict class
+
+    """
+    required = set(schema.get("required", []))
+    properties = schema.get("properties", {})
+
+    annotations: Dict[str, Any] = {}
+
+    for property, details in properties.items():
+        typ = schema_type_to_python(details, "typeddict")
+        if property not in required:
+            typ = Optional[typ]
+        annotations[property] = typ
+
+    return TypedDict(name or "AnonymousTypedDict", annotations)  # type: ignore
+
+
+def json_schema_dict_to_pydantic(
+    schema: dict,
+    name: Optional[str] = None
+) -> type[BaseModel]:
+    """Convert a JSON Schema dict into a Pydantic BaseModel class.
+
+    Parameters
+    ----------
+    schema: dict
+        The JSON Schema dict to convert to a Pydantic BaseModel
+    name: Optional[str]
+        The name of the Pydantic BaseModel
+
+    Returns
+    -------
+    type[BaseModel]
+        The Pydantic BaseModel class
+
+    """
+    required = set(schema.get("required", []))
+    properties = schema.get("properties", {})
+
+    field_definitions: Dict[str, Any] = {}
+
+    for property, details in properties.items():
+        typ = schema_type_to_python(details, "pydantic")
+        if property not in required:
+            field_definitions[property] = (Optional[typ], None)
+        else:
+            field_definitions[property] = (typ, ...)
+
+    return create_model(name or "AnonymousPydanticModel", **field_definitions)
+
+
+def json_schema_dict_to_dataclass(
+    schema: dict,
+    name: Optional[str] = None
+) -> type:
+    """Convert a JSON Schema dict into a dataclass.
+
+    Parameters
+    ----------
+    schema: dict
+        The JSON Schema dict to convert to a dataclass
+    name: Optional[str]
+        The name of the dataclass
+
+    Returns
+    -------
+    type
+        The dataclass
+
+    """
+    required = set(schema.get("required", []))
+    properties = schema.get("properties", {})
+
+    annotations: Dict[str, Any] = {}
+    defaults: Dict[str, Any] = {}
+
+    for property, details in properties.items():
+        typ = schema_type_to_python(details, "dataclass")
+        annotations[property] = typ
+
+        if property not in required:
+            defaults[property] = None
+
+    class_dict = {
+        '__annotations__': annotations,
+        '__module__': __name__,
+    }
+
+    for property, default_val in defaults.items():
+        class_dict[property] = field(default=default_val)
+
+    cls = type(name or "AnonymousDataclass", (), class_dict)
+    return dataclass(cls)