@@ -273,6 +273,8 @@ def get_dynamic_document_sections_description(
273273 class_desc = schema .get ("description" , "No description available" )
274274 table_name = f"document_sections_{ _get_table_suffix (class_name )}
275275 properties = schema .get (SCHEMA_PROPERTIES , {})
276+ # Get $defs for resolving $ref references
277+ defs = schema .get ("$defs" , {})
276278
277279 description += f'**`{ table_name } { class_name } \n '
278280 description += f"- **Description**: { class_desc } \n "
@@ -307,7 +309,7 @@ def get_dynamic_document_sections_description(
307309 column_count = 0 # Reset for each table
308310 prop_list = list (properties .keys ())
309311 for prop_index , (prop_desc_text , columns_added ) in enumerate (
310- _walk_properties_for_columns (properties )
312+ _walk_properties_for_columns (properties , defs = defs )
311313 ):
312314 description += prop_desc_text
313315 column_count += columns_added
@@ -320,8 +322,6 @@ def get_dynamic_document_sections_description(
320322 else :
321323 description += "- **Configuration-Specific Columns**: None configured\n "
322324
323- description += "\n "
324-
325325 description += """### Column Naming Patterns:
326326- **Simple attributes**: `inference_result.{attribute_name_lowercase}` (all strings)
327327- **Group attributes**: `inference_result.{group_name_lowercase}.{sub_attribute_lowercase}` (all strings)
@@ -433,6 +433,7 @@ def _walk_properties_for_columns(
433433 properties : Dict [str , Any ],
434434 parent_path : str = "inference_result" ,
435435 indent : str = " " ,
436+ defs : Optional [Dict [str , Any ]] = None ,
436437) -> Generator [tuple [str , int ], None , None ]:
437438 """
438439 Walk JSON Schema properties and yield (column_description, count) tuples.
@@ -441,24 +442,35 @@ def _walk_properties_for_columns(
441442 properties: JSON Schema properties dict
442443 parent_path: Parent column path
443444 indent: Indentation for formatting
445+ defs: Schema definitions for resolving $ref references
444446
445447 Yields:
446448 Tuples of (description_text, columns_added_count)
447449 """
448450 for prop_name , prop_schema in properties .items ():
451+ # Handle $ref by resolving to the actual definition
452+ if "$ref" in prop_schema and defs :
453+ ref_path = prop_schema ["$ref" ]
454+ # Extract the definition name from the reference (e.g., "#/$defs/employer_info")
455+ if ref_path .startswith ("#/$defs/" ):
456+ def_name = ref_path .replace ("#/$defs/" , "" )
457+ if def_name in defs :
458+ # Merge the referenced definition with any override fields
459+ resolved_schema = {** defs [def_name ], ** prop_schema }
460+ # Remove $ref from the resolved schema
461+ resolved_schema .pop ("$ref" , None )
462+ prop_schema = resolved_schema
463+
449464 prop_type = prop_schema .get (SCHEMA_TYPE )
450465 prop_desc = prop_schema .get (SCHEMA_DESCRIPTION , "" )
451466 column_path = f"{ parent_path } { prop_name .lower ()}
452467
453468 if prop_type == TYPE_OBJECT :
454- # Group - yield header and recurse
469+ # Group - recurse to get leaf columns only (no group header)
470+ # Groups don't become columns themselves - only leaf attributes do
455471 nested_props = prop_schema .get (SCHEMA_PROPERTIES , {})
456- yield (
457- f"{ indent } { prop_name } { len (nested_props )} \n " ,
458- 0 ,
459- )
460472 yield from _walk_properties_for_columns (
461- nested_props , column_path , indent + " "
473+ nested_props , column_path , indent , defs
462474 )
463475
464476 elif prop_type == TYPE_ARRAY :
@@ -657,6 +669,8 @@ def _get_specific_document_sections_table_info(
657669 class_name = matching_schema .get (X_AWS_IDP_DOCUMENT_TYPE , "Unknown" )
658670 class_desc = matching_schema .get ("description" , "No description available" )
659671 properties = matching_schema .get (SCHEMA_PROPERTIES , {})
672+ # Get $defs for resolving $ref references
673+ defs = matching_schema .get ("$defs" , {})
660674
661675 info = f"""## Document Sections Table: { table_name }
662676
@@ -678,7 +692,9 @@ def _get_specific_document_sections_table_info(
678692"""
679693
680694 if properties :
681- for prop_desc_text , _ in _walk_properties_for_columns (properties ):
695+ for prop_desc_text , _ in _walk_properties_for_columns (
696+ properties , defs = defs
697+ ):
682698 info += prop_desc_text
683699 else :
684700 info += "No configuration-specific columns defined.\n "
0 commit comments