malepati/custom_template_working
0
1from copy import deepcopy2from typing import Any, List3 4from openai import NOT_GIVEN5 6from utils.dict_utils import (7 get_dict_paths_with_key,8 get_dict_at_path,9 has_more_than_n_keys,10)11 12supported_string_formats = [13 "date-time",14 "time",15 "date",16 "duration",17 "email",18 "hostname",19 "ipv4",20 "ipv6",21 "uuid",22]23 24 25def remove_fields_from_schema(schema: dict, fields_to_remove: List[str]):26 schema = deepcopy(schema)27 properties_paths = get_dict_paths_with_key(schema, "properties")28 for path in properties_paths:29 parent_obj = get_dict_at_path(schema, path)30 if "properties" in parent_obj and isinstance(parent_obj["properties"], dict):31 for field in fields_to_remove:32 if field in parent_obj["properties"]:33 del parent_obj["properties"][field]34 35 required_paths = get_dict_paths_with_key(schema, "required")36 for path in required_paths:37 parent_obj = get_dict_at_path(schema, path)38 if "required" in parent_obj and isinstance(parent_obj["required"], list):39 parent_obj["required"] = [40 field41 for field in parent_obj["required"]42 if field not in fields_to_remove43 ]44 45 return schema46 47 48def add_field_in_schema(schema: dict, field: dict, required: bool = False) -> dict:49 50 if not isinstance(field, dict) or len(field) != 1:51 raise ValueError(52 "`field` must be a dict with exactly one entry: {name: schema_dict}"53 )54 55 field_name, field_schema = next(iter(field.items()))56 if not isinstance(field_name, str):57 raise TypeError("Field name must be a string")58 if not isinstance(field_schema, dict):59 raise TypeError("Field schema must be a dictionary")60 61 updated_schema: dict = deepcopy(schema)62 63 root_properties = updated_schema.get("properties")64 if not isinstance(root_properties, dict):65 updated_schema["properties"] = {}66 root_properties = updated_schema["properties"]67 68 root_properties[field_name] = field_schema69 70 # Update root-level required based on the flag71 existing_required = updated_schema.get("required")72 if not isinstance(existing_required, list):73 existing_required = []74 75 if required:76 if field_name not in existing_required:77 existing_required.append(field_name)78 else:79 if field_name in existing_required:80 existing_required = [name for name in existing_required if name != field_name]81 82 if existing_required:83 updated_schema["required"] = existing_required84 else:85 updated_schema.pop("required", None)86 87 return updated_schema88 89 90# From OpenAI91def ensure_strict_json_schema(92 json_schema: object,93 *,94 path: tuple[str, ...],95 root: dict[str, object],96) -> dict[str, Any]:97 """Mutates the given JSON schema to ensure it conforms to the `strict` standard98 that the API expects.99 """100 if not isinstance(json_schema, dict):101 raise TypeError(f"Expected {json_schema} to be a dictionary; path={path}")102 103 defs = json_schema.get("$defs")104 if isinstance(defs, dict):105 for def_name, def_schema in defs.items():106 ensure_strict_json_schema(107 def_schema, path=(*path, "$defs", def_name), root=root108 )109 110 definitions = json_schema.get("definitions")111 if isinstance(definitions, dict):112 for definition_name, definition_schema in definitions.items():113 ensure_strict_json_schema(114 definition_schema,115 path=(*path, "definitions", definition_name),116 root=root,117 )118 119 typ = json_schema.get("type")120 if typ == "object" and "additionalProperties" not in json_schema:121 json_schema["additionalProperties"] = False122 123 # object types124 # { 'type': 'object', 'properties': { 'a': {...} } }125 properties = json_schema.get("properties")126 if isinstance(properties, dict):127 json_schema["required"] = [prop for prop in properties.keys()]128 json_schema["properties"] = {129 key: ensure_strict_json_schema(130 prop_schema, path=(*path, "properties", key), root=root131 )132 for key, prop_schema in properties.items()133 }134 135 # arrays136 # { 'type': 'array', 'items': {...} }137 items = json_schema.get("items")138 if isinstance(items, dict):139 json_schema["items"] = ensure_strict_json_schema(140 items, path=(*path, "items"), root=root141 )142 143 # unions144 any_of = json_schema.get("anyOf")145 if isinstance(any_of, list):146 json_schema["anyOf"] = [147 ensure_strict_json_schema(variant, path=(*path, "anyOf", str(i)), root=root)148 for i, variant in enumerate(any_of)149 ]150 151 # intersections152 all_of = json_schema.get("allOf")153 if isinstance(all_of, list):154 if len(all_of) == 1:155 json_schema.update(156 ensure_strict_json_schema(157 all_of[0], path=(*path, "allOf", "0"), root=root158 )159 )160 json_schema.pop("allOf")161 else:162 json_schema["allOf"] = [163 ensure_strict_json_schema(164 entry, path=(*path, "allOf", str(i)), root=root165 )166 for i, entry in enumerate(all_of)167 ]168 169 # string170 if typ == "string":171 if "format" in json_schema:172 if json_schema["format"] not in supported_string_formats:173 del json_schema["format"]174 175 # strip `None` defaults as there's no meaningful distinction here176 # the schema will still be `nullable` and the model will default177 # to using `None` anyway178 if json_schema.get("default", NOT_GIVEN) is None:179 json_schema.pop("default")180 181 # we can't use `$ref`s if there are also other properties defined, e.g.182 # `{"$ref": "...", "description": "my description"}`183 #184 # so we unravel the ref185 # `{"type": "string", "description": "my description"}`186 ref = json_schema.get("$ref")187 if ref and has_more_than_n_keys(json_schema, 1):188 assert isinstance(ref, str), f"Received non-string $ref - {ref}"189 190 resolved = resolve_ref(root=root, ref=ref)191 if not isinstance(resolved, dict):192 raise ValueError(193 f"Expected `$ref: {ref}` to resolved to a dictionary but got {resolved}"194 )195 196 # properties from the json schema take priority over the ones on the `$ref`197 json_schema.update({**resolved, **json_schema})198 json_schema.pop("$ref")199 # Since the schema expanded from `$ref` might not have `additionalProperties: false` applied,200 # we call `_ensure_strict_json_schema` again to fix the inlined schema and ensure it's valid.201 return ensure_strict_json_schema(json_schema, path=path, root=root)202 203 return json_schema204 205 206def resolve_ref(*, root: dict[str, object], ref: str) -> object:207 if not ref.startswith("#/"):208 raise ValueError(f"Unexpected $ref format {ref!r}; Does not start with #/")209 210 path = ref[2:].split("/")211 resolved = root212 for key in path:213 value = resolved[key]214 assert isinstance(215 value, dict216 ), f"encountered non-dictionary entry while resolving {ref} - {resolved}"217 resolved = value218 219 return resolved220 221 222# Flattens a JSON schema by inlining all $ref references and removing $defs/definitions223def flatten_json_schema(schema: dict) -> dict:224 root_schema = deepcopy(schema)225 226 def _flatten(node: Any) -> Any:227 if isinstance(node, dict):228 # If node is a pure $ref (or combined with extra fields), inline it229 if "$ref" in node:230 ref_value = node["$ref"]231 assert isinstance(232 ref_value, str233 ), f"Received non-string $ref - {ref_value}"234 resolved = resolve_ref(root=root_schema, ref=ref_value)235 assert isinstance(236 resolved, dict237 ), f"Expected `$ref: {ref_value}` to resolve to a dictionary but got {type(resolved)}"238 # Merge: referenced first, then overlay current (excluding $ref)239 merged: dict[str, Any] = deepcopy(resolved)240 for key, value in node.items():241 if key == "$ref":242 continue243 merged[key] = value244 return _flatten(merged)245 246 flattened: dict[str, Any] = {}247 for key, value in node.items():248 # Drop defs/definitions in output249 if key in ("$defs", "definitions"):250 continue251 if key == "properties" and isinstance(value, dict):252 flattened[key] = {253 prop_key: _flatten(prop_val)254 for prop_key, prop_val in value.items()255 }256 elif key in ("items", "contains", "additionalProperties", "not"):257 if isinstance(value, dict):258 flattened[key] = _flatten(value)259 elif isinstance(value, list):260 flattened[key] = [_flatten(v) for v in value]261 else:262 flattened[key] = value263 elif key in ("allOf", "anyOf", "oneOf", "prefixItems") and isinstance(264 value, list265 ):266 flattened[key] = [_flatten(v) for v in value]267 else:268 flattened[key] = (269 _flatten(value) if isinstance(value, (dict, list)) else value270 )271 return flattened272 if isinstance(node, list):273 return [_flatten(v) for v in node]274 return node275 276 result = _flatten(schema)277 # Ensure top-level cleanup just in case278 if isinstance(result, dict):279 result.pop("$defs", None)280 result.pop("definitions", None)281 return result282 283 284def remove_titles_from_schema(schema: dict) -> dict[str, Any]:285 286 def _strip_titles(node: Any) -> Any:287 if isinstance(node, dict):288 rebuilt: dict[str, Any] = {}289 for key, value in node.items():290 # Preserve properties named "title" under the JSON Schema "properties" mapping291 if key == "properties" and isinstance(value, dict):292 rebuilt[key] = {293 prop_name: _strip_titles(prop_schema)294 for prop_name, prop_schema in value.items()295 }296 continue297 298 # Remove schema metadata field "title" elsewhere299 if key == "title":300 continue301 302 rebuilt[key] = _strip_titles(value)303 return rebuilt304 if isinstance(node, list):305 return [_strip_titles(item) for item in node]306 return node307 308 return _strip_titles(deepcopy(schema))309 310 311# ? Not used312def generate_constraint_sentences(schema: dict) -> str:313 """314 Generate human-readable constraint sentences from a JSON schema.315 316 Args:317 schema: JSON schema dictionary318 319 Returns:320 String containing constraint sentences separated by newlines321 """322 constraints = []323 324 def extract_constraints_recursive(obj, prefix=""):325 if isinstance(obj, dict):326 if "properties" in obj:327 properties = obj["properties"]328 for prop_name, prop_def in properties.items():329 current_path = f"{prefix}.{prop_name}" if prefix else prop_name330 331 if isinstance(prop_def, dict):332 prop_type = prop_def.get("type")333 334 # Handle string constraints335 if prop_type == "string":336 min_length = prop_def.get("minLength")337 max_length = prop_def.get("maxLength")338 339 if min_length is not None and max_length is not None:340 constraints.append(341 f" - {current_path} should be less than {max_length} characters and greater than {min_length} characters"342 )343 elif max_length is not None:344 constraints.append(345 f" - {current_path} should be less than {max_length} characters"346 )347 elif min_length is not None:348 constraints.append(349 f" - {current_path} should be greater than {min_length} characters"350 )351 352 # Handle array constraints353 elif prop_type == "array":354 min_items = prop_def.get("minItems")355 max_items = prop_def.get("maxItems")356 357 if min_items is not None and max_items is not None:358 constraints.append(359 f" - {current_path} should have more than {min_items} items and less than {max_items} items"360 )361 elif max_items is not None:362 constraints.append(363 f" - {current_path} should have less than {max_items} items"364 )365 elif min_items is not None:366 constraints.append(367 f" - {current_path} should have more than {min_items} items"368 )369 370 # Recurse into nested objects371 if prop_type == "object" or "properties" in prop_def:372 extract_constraints_recursive(prop_def, current_path)373 374 # Handle array items if they have properties375 if prop_type == "array" and "items" in prop_def:376 items_def = prop_def["items"]377 if isinstance(items_def, dict) and (378 "properties" in items_def379 or items_def.get("type") == "object"380 ):381 extract_constraints_recursive(382 items_def, f"{current_path}[*]"383 )384 385 # Also recurse into other nested structures386 for key, value in obj.items():387 if key not in [388 "properties",389 "type",390 "minLength",391 "maxLength",392 "minItems",393 "maxItems",394 ] and isinstance(value, dict):395 extract_constraints_recursive(value, prefix)396 397 # Start extraction from the root schema398 extract_constraints_recursive(schema)399 400 return "\n".join(constraints)401 402 403def remove_length_constraints_from_schema(schema: dict) -> dict:404 """405 Recursively removes maxLength and maxItems constraints from a JSON schema.406 This allows the LLM to generate content without arbitrary length limits.407 """408 schema = deepcopy(schema)409 410 def _recursive_remove(obj):411 if isinstance(obj, dict):412 # Remove keys413 if "maxLength" in obj:414 del obj["maxLength"]415 if "maxItems" in obj:416 del obj["maxItems"]417 418 # Recurse419 for key, value in obj.items():420 _recursive_remove(value)421 elif isinstance(obj, list):422 for item in obj:423 _recursive_remove(item)424 425 _recursive_remove(schema)426 return schema427