From 97da33494adbe63852544b5fe2cce6cb1f71ed0d Mon Sep 17 00:00:00 2001 From: Cole McIntosh <82463175+colesmcintosh@users.noreply.github.com> Date: Tue, 24 Jun 2025 16:38:48 -0600 Subject: [PATCH] Refactor unpack_defs to use iterative approach instead of recursion (#12017) * Refactor unpack_defs to use iterative approach instead of recursion - Replace recursive depth-first traversal with iterative queue-based approach - Add collections.deque import for efficient queue operations - Avoid potential stack overflow issues with deeply nested schemas - Maintain same functionality while improving performance and safety * Remove unused import of Set in common_utils.py * Enhance type hinting for queue in unpack_defs function in common_utils.py * Enhance unpack_defs function to handle key validation for parent structures in common_utils.py - Added checks to ensure that the parent is a dictionary or list and that the key is of the appropriate type (string for dicts, integer for lists) before assigning the resolved schema. - This improves the robustness of the unpack_defs function when dealing with various schema structures. --- .../prompt_templates/common_utils.py | 72 +++++++++---------- 1 file changed, 35 insertions(+), 37 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 6b70f690f9..626c8b7f29 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -16,7 +16,6 @@ from typing import ( Optional, Union, cast, - Set, ) from litellm.types.llms.openai import ( @@ -508,6 +507,7 @@ def unpack_defs(schema: dict, defs: dict) -> None: """ import copy + from collections import deque # Combine the defs handed down by the caller with defs/definitions found on # the current node. Local keys shadow parent keys to match JSON-schema @@ -518,14 +518,17 @@ def unpack_defs(schema: dict, defs: dict) -> None: **schema.get("definitions", {}), } - def _walk_and_resolve(node: Any, active_defs: dict, seen: Set[int]): # type: ignore[name-defined] - """Depth-first resolver that replaces ``{"$ref": "#/defs/Foo"}`` with - the *actual* ``Foo`` schema. - """ - - # Avoid infinite recursion on self-referential schemas + # Use iterative approach with queue to avoid recursion + # Each item in queue is (node, parent_container, key/index, active_defs, seen_ids) + queue: deque[tuple[Any, Union[dict, list, None], Union[str, int, None], dict, set]] = deque([(schema, None, None, root_defs, set())]) + + while queue: + node, parent, key, active_defs, seen = queue.popleft() + + # Avoid infinite loops on self-referential schemas if id(node) in seen: - return node + continue + seen = seen.copy() # Create new set for this branch seen.add(id(node)) # ----------------------------- dict ----------------------------- @@ -536,7 +539,7 @@ def unpack_defs(schema: dict, defs: dict) -> None: target_schema = active_defs.get(ref_name) # Unknown reference – leave untouched if target_schema is None: - return node + continue # Merge defs from the target to capture nested definitions child_defs = { @@ -545,12 +548,24 @@ def unpack_defs(schema: dict, defs: dict) -> None: **target_schema.get("definitions", {}), } - # Recursively resolve the target *copy* to avoid mutating the - # shared definition map. - resolved = _walk_and_resolve(copy.deepcopy(target_schema), child_defs, seen) - return resolved + # Replace the reference with resolved copy + resolved = copy.deepcopy(target_schema) + if parent is not None and key is not None: + if isinstance(parent, dict) and isinstance(key, str): + parent[key] = resolved + elif isinstance(parent, list) and isinstance(key, int): + parent[key] = resolved + else: + # This is the root schema itself + schema.clear() + schema.update(resolved) + resolved = schema + + # Add resolved node to queue for further processing + queue.append((resolved, parent, key, child_defs, seen)) + continue - # --- Case 2: regular dict – recurse into its values --- + # --- Case 2: regular dict – process its values --- # Update defs with any nested $defs/definitions present *here*. current_defs = { **active_defs, @@ -558,32 +573,15 @@ def unpack_defs(schema: dict, defs: dict) -> None: **node.get("definitions", {}), } - for key, val in list(node.items()): - node[key] = _walk_and_resolve(val, current_defs, seen) - return node + # Add all dict values to queue + for k, v in node.items(): + queue.append((v, node, k, current_defs, seen)) # ---------------------------- list ------------------------------ - if isinstance(node, list): + elif isinstance(node, list): + # Add all list items to queue for idx, item in enumerate(node): - node[idx] = _walk_and_resolve(item, active_defs, seen) - return node - - # -------------------------- primitive --------------------------- - return node - - # Kick off traversal - resolved_root = _walk_and_resolve(schema, root_defs, set()) - # If the resolver returned a *different* dict (e.g., the root itself was a - # $ref), mirror the changes back into the original object so that callers - # holding a reference to ``schema`` see the updated structure. - if resolved_root is not schema: - schema.clear() - if isinstance(resolved_root, dict): - schema.update(resolved_root) - else: - # In the very unlikely case the root was resolved to a non-dict - # (e.g., a primitive), replace in-place via a sentinel key. - schema["__resolved_value__"] = resolved_root # type: ignore + queue.append((item, node, idx, active_defs, seen)) def _get_image_mime_type_from_url(url: str) -> Optional[str]: