File size: 5,285 Bytes
e062359 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | from __future__ import annotations
import inspect
from typing import Any, Literal, Optional, cast
from typing_extensions import assert_never
import pydantic
from ..._utils import is_list
SupportedTypes = Literal[
"object",
"array",
"string",
"integer",
"number",
"boolean",
"null",
]
SupportedStringFormats = {
"date-time",
"time",
"date",
"duration",
"email",
"hostname",
"uri",
"ipv4",
"ipv6",
"uuid",
}
def get_transformed_string(
schema: dict[str, Any],
) -> dict[str, Any]:
"""Transforms a JSON schema of type string to ensure it conforms to the API's expectations.
Specifically, it ensures that if the schema is of type "string" and does not already
specify a "format", it sets the format to "text".
Args:
schema: The original JSON schema.
Returns:
The transformed JSON schema.
"""
if schema.get("type") == "string" and "format" not in schema:
schema["format"] = "text"
return schema
def transform_schema(
json_schema: type[pydantic.BaseModel] | dict[str, Any],
) -> dict[str, Any]:
"""
Transforms a JSON schema to ensure it conforms to the API's expectations.
Args:
json_schema (Dict[str, Any]): The original JSON schema.
Returns:
The transformed JSON schema.
Examples:
>>> transform_schema(
... {
... "type": "integer",
... "minimum": 1,
... "maximum": 10,
... "description": "A number",
... }
... )
{'type': 'integer', 'description': 'A number\n\n{minimum: 1, maximum: 10}'}
"""
if inspect.isclass(json_schema) and issubclass(json_schema, pydantic.BaseModel): # pyright: ignore[reportUnnecessaryIsInstance]
json_schema = json_schema.model_json_schema()
strict_schema: dict[str, Any] = {}
json_schema = {**json_schema}
ref = json_schema.pop("$ref", None)
if ref is not None:
strict_schema["$ref"] = ref
return strict_schema
defs = json_schema.pop("$defs", None)
if defs is not None:
strict_defs: dict[str, Any] = {}
strict_schema["$defs"] = strict_defs
for name, schema in defs.items():
strict_defs[name] = transform_schema(schema)
type_: Optional[SupportedTypes] = json_schema.pop("type", None)
any_of = json_schema.pop("anyOf", None)
one_of = json_schema.pop("oneOf", None)
all_of = json_schema.pop("allOf", None)
if is_list(any_of):
strict_schema["anyOf"] = [transform_schema(cast("dict[str, Any]", variant)) for variant in any_of]
elif is_list(one_of):
strict_schema["anyOf"] = [transform_schema(cast("dict[str, Any]", variant)) for variant in one_of]
elif is_list(all_of):
strict_schema["allOf"] = [transform_schema(cast("dict[str, Any]", variant)) for variant in all_of]
else:
if type_ is None:
raise ValueError("Schema must have a 'type', 'anyOf', 'oneOf', or 'allOf' field.")
strict_schema["type"] = type_
description = json_schema.pop("description", None)
if description is not None:
strict_schema["description"] = description
title = json_schema.pop("title", None)
if title is not None:
strict_schema["title"] = title
if type_ == "object":
strict_schema["properties"] = {
key: transform_schema(prop_schema) for key, prop_schema in json_schema.pop("properties", {}).items()
}
json_schema.pop("additionalProperties", None)
strict_schema["additionalProperties"] = False
required = json_schema.pop("required", None)
if required is not None:
strict_schema["required"] = required
elif type_ == "string":
format = json_schema.pop("format", None)
if format and format in SupportedStringFormats:
strict_schema["format"] = format
elif format:
# add it back so its treated as an extra property and appended to the description
json_schema["format"] = format
elif type_ == "array":
items = json_schema.pop("items", None)
if items is not None:
strict_schema["items"] = transform_schema(items)
min_items = json_schema.pop("minItems", None)
if min_items is not None and min_items == 0 or min_items == 1:
strict_schema["minItems"] = min_items
elif min_items is not None:
# add it back so its treated as an extra property and appended to the description
json_schema["minItems"] = min_items
elif type_ == "boolean" or type_ == "integer" or type_ == "number" or type_ == "null" or type_ is None:
pass
else:
assert_never(type_)
# if there are any propes leftover then they aren't supported, so we add them to the description
# so that the model *might* follow them.
if json_schema:
description = strict_schema.get("description")
strict_schema["description"] = (
(description + "\n\n" if description is not None else "")
+ "{"
+ ", ".join(f"{key}: {value}" for key, value in json_schema.items())
+ "}"
)
return strict_schema
|