llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/strong_typing/classdef.py

@@ -0,0 +1,440 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import copy
+import dataclasses
+import datetime
+import decimal
+import enum
+import ipaddress
+import math
+import re
+import sys
+import types
+import typing
+import uuid
+from dataclasses import dataclass
+from typing import Any, Dict, List, Literal, Optional, Tuple, Type, TypeVar, Union
+
+from .auxiliary import (
+    Alias,
+    Annotated,
+    MaxLength,
+    Precision,
+    float32,
+    float64,
+    int16,
+    int32,
+    int64,
+)
+from .core import JsonType, Schema
+from .docstring import Docstring, DocstringParam
+from .inspection import TypeLike
+from .serialization import json_to_object, object_to_json
+
+T = TypeVar("T")
+
+
+@dataclass
+class JsonSchemaNode:
+    title: Optional[str]
+    description: Optional[str]
+
+
+@dataclass
+class JsonSchemaType(JsonSchemaNode):
+    type: str
+    format: Optional[str]
+
+
+@dataclass
+class JsonSchemaBoolean(JsonSchemaType):
+    type: Literal["boolean"]
+    const: Optional[bool]
+    default: Optional[bool]
+    examples: Optional[List[bool]]
+
+
+@dataclass
+class JsonSchemaInteger(JsonSchemaType):
+    type: Literal["integer"]
+    const: Optional[int]
+    default: Optional[int]
+    examples: Optional[List[int]]
+    enum: Optional[List[int]]
+    minimum: Optional[int]
+    maximum: Optional[int]
+
+
+@dataclass
+class JsonSchemaNumber(JsonSchemaType):
+    type: Literal["number"]
+    const: Optional[float]
+    default: Optional[float]
+    examples: Optional[List[float]]
+    minimum: Optional[float]
+    maximum: Optional[float]
+    exclusiveMinimum: Optional[float]
+    exclusiveMaximum: Optional[float]
+    multipleOf: Optional[float]
+
+
+@dataclass
+class JsonSchemaString(JsonSchemaType):
+    type: Literal["string"]
+    const: Optional[str]
+    default: Optional[str]
+    examples: Optional[List[str]]
+    enum: Optional[List[str]]
+    minLength: Optional[int]
+    maxLength: Optional[int]
+
+
+@dataclass
+class JsonSchemaArray(JsonSchemaType):
+    type: Literal["array"]
+    items: "JsonSchemaAny"
+
+
+@dataclass
+class JsonSchemaObject(JsonSchemaType):
+    type: Literal["object"]
+    properties: Optional[Dict[str, "JsonSchemaAny"]]
+    additionalProperties: Optional[bool]
+    required: Optional[List[str]]
+
+
+@dataclass
+class JsonSchemaRef(JsonSchemaNode):
+    ref: Annotated[str, Alias("$ref")]
+
+
+@dataclass
+class JsonSchemaAllOf(JsonSchemaNode):
+    allOf: List["JsonSchemaAny"]
+
+
+@dataclass
+class JsonSchemaAnyOf(JsonSchemaNode):
+    anyOf: List["JsonSchemaAny"]
+
+
+@dataclass
+class Discriminator:
+    propertyName: str
+    mapping: Dict[str, str]
+
+
+@dataclass
+class JsonSchemaOneOf(JsonSchemaNode):
+    oneOf: List["JsonSchemaAny"]
+    discriminator: Optional[Discriminator]
+
+
+JsonSchemaAny = Union[
+    JsonSchemaRef,
+    JsonSchemaBoolean,
+    JsonSchemaInteger,
+    JsonSchemaNumber,
+    JsonSchemaString,
+    JsonSchemaArray,
+    JsonSchemaObject,
+    JsonSchemaOneOf,
+]
+
+
+@dataclass
+class JsonSchemaTopLevelObject(JsonSchemaObject):
+    schema: Annotated[str, Alias("$schema")]
+    definitions: Optional[Dict[str, JsonSchemaAny]]
+
+
+def integer_range_to_type(min_value: float, max_value: float) -> type:
+    if min_value >= -(2**15) and max_value < 2**15:
+        return int16
+    elif min_value >= -(2**31) and max_value < 2**31:
+        return int32
+    else:
+        return int64
+
+
+def enum_safe_name(name: str) -> str:
+    name = re.sub(r"\W", "_", name)
+    is_dunder = name.startswith("__")
+    is_sunder = name.startswith("_") and name.endswith("_")
+    if is_dunder or is_sunder:  # provide an alternative for dunder and sunder names
+        name = f"v{name}"
+    return name
+
+
+def enum_values_to_type(
+    module: types.ModuleType,
+    name: str,
+    values: Dict[str, Any],
+    title: Optional[str] = None,
+    description: Optional[str] = None,
+) -> Type[enum.Enum]:
+    enum_class: Type[enum.Enum] = enum.Enum(name, values)  # type: ignore
+
+    # assign the newly created type to the same module where the defining class is
+    enum_class.__module__ = module.__name__
+    enum_class.__doc__ = str(Docstring(short_description=title, long_description=description))
+    setattr(module, name, enum_class)
+
+    return enum.unique(enum_class)
+
+
+def schema_to_type(schema: Schema, *, module: types.ModuleType, class_name: str) -> TypeLike:
+    """
+    Creates a Python type from a JSON schema.
+
+    :param schema: The JSON schema that the types would correspond to.
+    :param module: The module in which to create the new types.
+    :param class_name: The name assigned to the top-level class.
+    """
+
+    top_node = typing.cast(JsonSchemaTopLevelObject, json_to_object(JsonSchemaTopLevelObject, schema))
+    if top_node.definitions is not None:
+        for type_name, type_node in top_node.definitions.items():
+            type_def = node_to_typedef(module, type_name, type_node)
+            if type_def.default is not dataclasses.MISSING:
+                raise TypeError("disallowed: `default` for top-level type definitions")
+
+            type_def.type.__module__ = module.__name__
+            setattr(module, type_name, type_def.type)
+
+    return node_to_typedef(module, class_name, top_node).type
+
+
+@dataclass
+class TypeDef:
+    type: TypeLike
+    default: Any = dataclasses.MISSING
+
+
+def json_to_value(target_type: TypeLike, data: JsonType) -> Any:
+    if data is not None:
+        return json_to_object(target_type, data)
+    else:
+        return dataclasses.MISSING
+
+
+def node_to_typedef(module: types.ModuleType, context: str, node: JsonSchemaNode) -> TypeDef:
+    if isinstance(node, JsonSchemaRef):
+        match_obj = re.match(r"^#/definitions/(\w+)$", node.ref)
+        if not match_obj:
+            raise ValueError(f"invalid reference: {node.ref}")
+
+        type_name = match_obj.group(1)
+        return TypeDef(getattr(module, type_name), dataclasses.MISSING)
+
+    elif isinstance(node, JsonSchemaBoolean):
+        if node.const is not None:
+            return TypeDef(Literal[node.const], dataclasses.MISSING)
+
+        default = json_to_value(bool, node.default)
+        return TypeDef(bool, default)
+
+    elif isinstance(node, JsonSchemaInteger):
+        if node.const is not None:
+            return TypeDef(Literal[node.const], dataclasses.MISSING)
+
+        integer_type: TypeLike
+        if node.format == "int16":
+            integer_type = int16
+        elif node.format == "int32":
+            integer_type = int32
+        elif node.format == "int64":
+            integer_type = int64
+        else:
+            if node.enum is not None:
+                integer_type = integer_range_to_type(min(node.enum), max(node.enum))
+            elif node.minimum is not None and node.maximum is not None:
+                integer_type = integer_range_to_type(node.minimum, node.maximum)
+            else:
+                integer_type = int
+
+        default = json_to_value(integer_type, node.default)
+        return TypeDef(integer_type, default)
+
+    elif isinstance(node, JsonSchemaNumber):
+        if node.const is not None:
+            return TypeDef(Literal[node.const], dataclasses.MISSING)
+
+        number_type: TypeLike
+        if node.format == "float32":
+            number_type = float32
+        elif node.format == "float64":
+            number_type = float64
+        else:
+            if (
+                node.exclusiveMinimum is not None
+                and node.exclusiveMaximum is not None
+                and node.exclusiveMinimum == -node.exclusiveMaximum
+            ):
+                integer_digits = round(math.log10(node.exclusiveMaximum))
+            else:
+                integer_digits = None
+
+            if node.multipleOf is not None:
+                decimal_digits = -round(math.log10(node.multipleOf))
+            else:
+                decimal_digits = None
+
+            if integer_digits is not None and decimal_digits is not None:
+                number_type = Annotated[
+                    decimal.Decimal,
+                    Precision(integer_digits + decimal_digits, decimal_digits),
+                ]
+            else:
+                number_type = float
+
+        default = json_to_value(number_type, node.default)
+        return TypeDef(number_type, default)
+
+    elif isinstance(node, JsonSchemaString):
+        if node.const is not None:
+            return TypeDef(Literal[node.const], dataclasses.MISSING)
+
+        string_type: TypeLike
+        if node.format == "date-time":
+            string_type = datetime.datetime
+        elif node.format == "uuid":
+            string_type = uuid.UUID
+        elif node.format == "ipv4":
+            string_type = ipaddress.IPv4Address
+        elif node.format == "ipv6":
+            string_type = ipaddress.IPv6Address
+
+        elif node.enum is not None:
+            string_type = enum_values_to_type(
+                module,
+                context,
+                {enum_safe_name(e): e for e in node.enum},
+                title=node.title,
+                description=node.description,
+            )
+
+        elif node.maxLength is not None:
+            string_type = Annotated[str, MaxLength(node.maxLength)]
+        else:
+            string_type = str
+
+        default = json_to_value(string_type, node.default)
+        return TypeDef(string_type, default)
+
+    elif isinstance(node, JsonSchemaArray):
+        type_def = node_to_typedef(module, context, node.items)
+        if type_def.default is not dataclasses.MISSING:
+            raise TypeError("disallowed: `default` for array element type")
+        list_type = List[(type_def.type,)]  # type: ignore
+        return TypeDef(list_type, dataclasses.MISSING)
+
+    elif isinstance(node, JsonSchemaObject):
+        if node.properties is None:
+            return TypeDef(JsonType, dataclasses.MISSING)
+
+        if node.additionalProperties is None or node.additionalProperties is not False:
+            raise TypeError("expected: `additionalProperties` equals `false`")
+
+        required = node.required if node.required is not None else []
+
+        class_name = context
+
+        fields: List[Tuple[str, Any, dataclasses.Field]] = []
+        params: Dict[str, DocstringParam] = {}
+        for prop_name, prop_node in node.properties.items():
+            type_def = node_to_typedef(module, f"{class_name}__{prop_name}", prop_node)
+            if prop_name in required:
+                prop_type = type_def.type
+            else:
+                prop_type = Union[(None, type_def.type)]
+            fields.append((prop_name, prop_type, dataclasses.field(default=type_def.default)))
+            prop_desc = prop_node.title or prop_node.description
+            if prop_desc is not None:
+                params[prop_name] = DocstringParam(prop_name, prop_desc)
+
+        fields.sort(key=lambda t: t[2].default is not dataclasses.MISSING)
+        if sys.version_info >= (3, 12):
+            class_type = dataclasses.make_dataclass(class_name, fields, module=module.__name__)
+        else:
+            class_type = dataclasses.make_dataclass(class_name, fields, namespace={"__module__": module.__name__})
+        class_type.__doc__ = str(
+            Docstring(
+                short_description=node.title,
+                long_description=node.description,
+                params=params,
+            )
+        )
+        setattr(module, class_name, class_type)
+        return TypeDef(class_type, dataclasses.MISSING)
+
+    elif isinstance(node, JsonSchemaOneOf):
+        union_defs = tuple(node_to_typedef(module, context, n) for n in node.oneOf)
+        if any(d.default is not dataclasses.MISSING for d in union_defs):
+            raise TypeError("disallowed: `default` for union member type")
+        union_types = tuple(d.type for d in union_defs)
+        return TypeDef(Union[union_types], dataclasses.MISSING)
+
+    raise NotImplementedError()
+
+
+@dataclass
+class SchemaFlatteningOptions:
+    qualified_names: bool = False
+    recursive: bool = False
+
+
+def flatten_schema(schema: Schema, *, options: Optional[SchemaFlatteningOptions] = None) -> Schema:
+    top_node = typing.cast(JsonSchemaTopLevelObject, json_to_object(JsonSchemaTopLevelObject, schema))
+    flattener = SchemaFlattener(options)
+    obj = flattener.flatten(top_node)
+    return typing.cast(Schema, object_to_json(obj))
+
+
+class SchemaFlattener:
+    options: SchemaFlatteningOptions
+
+    def __init__(self, options: Optional[SchemaFlatteningOptions] = None) -> None:
+        self.options = options or SchemaFlatteningOptions()
+
+    def flatten(self, source_node: JsonSchemaObject) -> JsonSchemaObject:
+        if source_node.type != "object":
+            return source_node
+
+        source_props = source_node.properties or {}
+        target_props: Dict[str, JsonSchemaAny] = {}
+
+        source_reqs = source_node.required or []
+        target_reqs: List[str] = []
+
+        for name, prop in source_props.items():
+            if not isinstance(prop, JsonSchemaObject):
+                target_props[name] = prop
+                if name in source_reqs:
+                    target_reqs.append(name)
+                continue
+
+            if self.options.recursive:
+                obj = self.flatten(prop)
+            else:
+                obj = prop
+            if obj.properties is not None:
+                if self.options.qualified_names:
+                    target_props.update((f"{name}.{n}", p) for n, p in obj.properties.items())
+                else:
+                    target_props.update(obj.properties.items())
+            if obj.required is not None:
+                if self.options.qualified_names:
+                    target_reqs.extend(f"{name}.{n}" for n in obj.required)
+                else:
+                    target_reqs.extend(obj.required)
+
+        target_node = copy.copy(source_node)
+        target_node.properties = target_props or None
+        target_node.additionalProperties = False
+        target_node.required = target_reqs or None
+        return target_node

llama_stack/strong_typing/core.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+"""
+Type-safe data interchange for Python data classes.
+
+:see: https://github.com/hunyadi/strong_typing
+"""
+
+from typing import Dict, List, Union
+
+
+class JsonObject:
+    "Placeholder type for an unrestricted JSON object."
+
+
+class JsonArray:
+    "Placeholder type for an unrestricted JSON array."
+
+
+# a JSON type with possible `null` values
+JsonType = Union[
+    None,
+    bool,
+    int,
+    float,
+    str,
+    Dict[str, "JsonType"],
+    List["JsonType"],
+]
+
+# a JSON type that cannot contain `null` values
+StrictJsonType = Union[
+    bool,
+    int,
+    float,
+    str,
+    Dict[str, "StrictJsonType"],
+    List["StrictJsonType"],
+]
+
+# a meta-type that captures the object type in a JSON schema
+Schema = Dict[str, JsonType]