From df269086c08e446835e534cef9a36bf555186745 Mon Sep 17 00:00:00 2001 From: Vishakh Pillai Date: Tue, 4 Aug 2026 17:40:37 -0400 Subject: [PATCH] fix: resolve SMA join aliases to physical columns in Pass 1 Join-key validation now reverse-maps canonical keys to physical columns (matching the executor), and Pass 1 can append column_aliases_to_add so JOIN name bridges are not dropped when refining. Co-authored-by: Cursor --- .../manifest/prompts/refine.py | 69 ++++++-- .../manifest/validation.py | 65 +++++-- .../manifest/test_join_alias_resolution.py | 165 ++++++++++++++++++ .../manifest/test_pass1_column_aliases.py | 125 +++++++++++++ 4 files changed, 403 insertions(+), 21 deletions(-) create mode 100644 tests/genai/mapping/schema_mapping_agent/manifest/test_join_alias_resolution.py create mode 100644 tests/genai/mapping/schema_mapping_agent/manifest/test_pass1_column_aliases.py diff --git a/src/edvise/genai/mapping/schema_mapping_agent/manifest/prompts/refine.py b/src/edvise/genai/mapping/schema_mapping_agent/manifest/prompts/refine.py index 47b13811d..2e7e3835e 100644 --- a/src/edvise/genai/mapping/schema_mapping_agent/manifest/prompts/refine.py +++ b/src/edvise/genai/mapping/schema_mapping_agent/manifest/prompts/refine.py @@ -2,8 +2,8 @@ SMA refinement + HITL: prompts, orchestration, and post-parse safety nets (single module). **Pass 1** — refinement + HITL flagging (slim JSON: ``field_statuses``, -``refined_corrections``, ``hitl_flags``; no full manifest): one LLM call per entity -(cohort and course are separate calls). +``refined_corrections``, optional ``column_aliases_to_add``, ``hitl_flags``; no full +manifest): one LLM call per entity (cohort and course are separate calls). **Pass 2** — option generation: one LLM call per entity with all Pass 1 flags for that entity in a single ``items`` array (cohort + course = 2 Pass 2 calls per institution; @@ -180,6 +180,15 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: // only for fields with field_statuses[target_field]="refined_by_llm" or "refined_and_proposed_for_hitl" // omit key entirely if no fields were corrected }, + "column_aliases_to_add": [ + // optional — ColumnAlias objects to append to the entity manifest when a + // JOIN fix needs a name bridge (same shape as Pass 2 option.column_alias). + // Example: course join_keys use term_descr but student physical column is + // term_desc → {table: "student", source_column: "term_desc", + // canonical_column: "term_descr", rationale: "..."}. + // Omit key or use [] when no aliases are needed. + // ColumnAlias: {table: str!, source_column: str!, canonical_column: str!, rationale?: str} + ], "hitl_flags": [ { "item_id": "{institution_id}_{entity_type}_{target_field}_{failure_mode}", @@ -205,6 +214,8 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: - Do not change confidence on any field. - Every proposed_for_hitl or refined_and_proposed_for_hitl field must appear in hitl_flags. - Every refined_by_llm or refined_and_proposed_for_hitl field with a correction must appear in refined_corrections. + - When a JOIN fix renames join_keys to a canonical name that differs on the other + table, ALSO emit column_aliases_to_add (do not only change join_keys). """ _PASS1_OUTPUT_SCHEMA_COMBINED = """ @@ -222,6 +233,10 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: } // omit entity key if no corrections for that entity }, + "column_aliases_to_add_by_entity": { + "": [ /* same shape as column_aliases_to_add in single-entity Pass 1 */ ] + // omit entity key or use [] when no aliases for that entity + }, "hitl_flags_by_entity": { "": [ /* same shape as hitl_flags in single-entity Pass 1 */ ] } @@ -230,7 +245,8 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: CRITICAL: - field_statuses_by_entity, refined_corrections_by_entity, and hitl_flags_by_entity must contain exactly the same entity_type keys as listed in the user message. - - Per-entity rules match single-entity Pass 1 (complete field_statuses, corrections, flags). + - Per-entity rules match single-entity Pass 1 (complete field_statuses, corrections, + column_aliases_to_add, flags). - Do not emit full manifests — slim keys only. """ @@ -280,6 +296,8 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: HIGH CONFIDENCE — validation errors: - Deterministic fix (typo, structural): set field_statuses[target_field]="refined_by_llm", emit deltas in refined_corrections, no hitl_flags entry for that field. + JOIN name mismatches that need a ColumnAlias: also emit column_aliases_to_add + (canonical join_keys alone are not enough when physical names differ). - Fix requires judgment or is ambiguous: set field_statuses[target_field]="proposed_for_hitl" and emit hitl_flags. @@ -457,6 +475,7 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: { "field_statuses": { ...every target_field... }, "refined_corrections": { ...optional — refined_by_llm / refined_and_proposed_for_hitl deltas... }, + "column_aliases_to_add": [ ...optional ColumnAlias objects for JOIN name bridges... ], "hitl_flags": [ ...optional — proposed_for_hitl / refined_and_proposed_for_hitl... ] } @@ -465,7 +484,8 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: - Do not invent columns or tables not present in the schema contract. - Do not change confidence on any field. - Do not emit options in Pass 1. - - Do not output a full manifest — field_statuses + refined_corrections + hitl_flags only. + - Do not output a full manifest — field_statuses + refined_corrections + + column_aliases_to_add + hitl_flags only. """ _PASS1_OUTPUT_FORMAT_COMBINED = """ @@ -479,6 +499,10 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: "": { ...optional per-entity refined_corrections... }, ... }, + "column_aliases_to_add_by_entity": { + "": [ ...optional ColumnAlias objects... ], + ... + }, "hitl_flags_by_entity": { "": [ ...Pass 1 flags for that entity only — no options... ], ... @@ -486,8 +510,9 @@ def _parse_sma_refinement_llm_dict(raw: str) -> dict[str, Any]: } CRITICAL: - - All three top-level objects must contain exactly the same entity_type keys - as listed in the user message (e.g. cohort and course). + - field_statuses_by_entity and hitl_flags_by_entity must contain exactly the same + entity_type keys as listed in the user message (e.g. cohort and course); + refined_corrections_by_entity / column_aliases_to_add_by_entity may omit empty entities. - Each field_statuses_by_entity entry must list every target_field for that entity. - Do not invent columns or tables not present in the schema contract. - Do not change confidence on any field. @@ -763,6 +788,8 @@ def build_refinement_pass1_user_prompt( - If confidence <= {HITL_CONFIDENCE_THRESHOLD} and you made a correction: set field_statuses[target_field]="refined_and_proposed_for_hitl", put deltas in refined_corrections, and add a hitl_flag (correction is option 1 in Pass 2). + - JOIN key renames that need a physical↔canonical name bridge: also emit + column_aliases_to_add (same ColumnAlias shape as Pass 2 option.column_alias). 2. For fields you cannot confidently fix (including low confidence with no correction), set field_statuses[target_field]="proposed_for_hitl" and add a hitl_flag with current_field_mapping @@ -773,7 +800,8 @@ def build_refinement_pass1_user_prompt( Do not include them in refined_corrections or hitl_flags. 4. Return the single JSON object described in your instructions. - Do not output a full manifest — field_statuses + refined_corrections + hitl_flags only. + Do not output a full manifest — field_statuses + refined_corrections + + column_aliases_to_add + hitl_flags only. """ @@ -839,6 +867,8 @@ def build_refinement_combined_pass1_user_prompt( - If confidence <= {HITL_CONFIDENCE_THRESHOLD} and you made a correction: set field_statuses_by_entity[entity][target_field]="refined_and_proposed_for_hitl", put deltas in refined_corrections_by_entity, and add a hitl_flags_by_entity[entity] entry. + - JOIN key renames that need a physical↔canonical name bridge: also emit + column_aliases_to_add_by_entity[entity]. 2. For fields you cannot confidently fix (including low confidence with no correction), set field_statuses_by_entity[entity][target_field]="proposed_for_hitl" and add hitl_flags_by_entity[entity] @@ -850,7 +880,7 @@ def build_refinement_combined_pass1_user_prompt( 4. Return the single combined JSON object described in your instructions. Do not output full manifests — field_statuses_by_entity + refined_corrections_by_entity + - hitl_flags_by_entity only. + column_aliases_to_add_by_entity + hitl_flags_by_entity only. """ @@ -940,9 +970,15 @@ def _apply_pass1_result( ) -> tuple[FieldMappingManifest, list[dict[str, Any]]]: """ Reconstruct full manifest from Pass 1 slim output. - Merges refined_corrections onto input records. + Merges refined_corrections onto input records and appends + ``column_aliases_to_add`` via :func:`add_alias_if_missing`. Sets review_status on every record from field_statuses. """ + from edvise.genai.mapping.schema_mapping_agent.manifest.hitl.schemas import ( + add_alias_if_missing, + ) + from edvise.genai.mapping.schema_mapping_agent.manifest.schemas import ColumnAlias + field_statuses = pass1_result.get("field_statuses") if not isinstance(field_statuses, dict): raise ValueError("Pass 1 output missing or invalid field_statuses") @@ -1000,8 +1036,21 @@ def _apply_pass1_result( entity_type=input_manifest.entity_type, target_schema=input_manifest.target_schema, mappings=updated_mappings, - column_aliases=input_manifest.column_aliases, + column_aliases=list(input_manifest.column_aliases), ) + + aliases_raw = pass1_result.get("column_aliases_to_add") or [] + if aliases_raw is None: + aliases_raw = [] + if not isinstance(aliases_raw, list): + raise ValueError("Pass 1 column_aliases_to_add must be a list") + for raw_alias in aliases_raw: + if not isinstance(raw_alias, dict): + raise ValueError( + "Pass 1 column_aliases_to_add entries must be ColumnAlias objects" + ) + add_alias_if_missing(refined_manifest, ColumnAlias.model_validate(raw_alias)) + return refined_manifest, hitl_flags diff --git a/src/edvise/genai/mapping/schema_mapping_agent/manifest/validation.py b/src/edvise/genai/mapping/schema_mapping_agent/manifest/validation.py index c17577ff8..b5f13b1e2 100644 --- a/src/edvise/genai/mapping/schema_mapping_agent/manifest/validation.py +++ b/src/edvise/genai/mapping/schema_mapping_agent/manifest/validation.py @@ -349,21 +349,55 @@ class ManifestValidationError(BaseModel): # --------------------------------------------------------------------------- # Alias resolution helpers # --------------------------------------------------------------------------- +# JoinConfig.join_keys are *canonical* names (same convention as the field +# executor). ColumnAlias maps physical ``source_column`` → ``canonical_column`` +# on one table. Existence checks must reverse that map (canonical → physical); +# bridge checks must confirm both physicals share the join key as canonical. -def _resolve_column_via_aliases( +def _physical_column_for_join_key( table: str, - column: str, + join_key: str, column_aliases: list[ColumnAlias], ) -> str: """ - Return canonical_column if a ColumnAlias entry exists for (table, column), - otherwise return column unchanged. + Map a canonical join key to the physical column name on ``table``. + + If an alias declares ``canonical_column == join_key`` for ``table``, return + its ``source_column``; otherwise the join key is assumed to be the physical + name (same as :func:`~edvise.genai.mapping.schema_mapping_agent.execution.field_executor._resolve_join_keys`). """ for alias in column_aliases: - if alias.table == table and alias.source_column == column: + if alias.table == table and alias.canonical_column == join_key: + return alias.source_column + return join_key + + +def _canonical_column_for_physical( + table: str, + physical_column: str, + column_aliases: list[ColumnAlias], +) -> str: + """Map a physical column to its canonical join-key name on ``table``.""" + for alias in column_aliases: + if alias.table == table and alias.source_column == physical_column: return alias.canonical_column - return column + return physical_column + + +def _resolve_column_via_aliases( + table: str, + column: str, + column_aliases: list[ColumnAlias], +) -> str: + """ + Resolve a join-key name to the physical column on ``table``. + + ``column`` is a canonical ``join_keys`` entry. Prefer + :func:`_physical_column_for_join_key`; this name is kept for call-site + compatibility with older docs/tests. + """ + return _physical_column_for_join_key(table, column, column_aliases) def _alias_bridges_join( @@ -373,14 +407,23 @@ def _alias_bridges_join( column_aliases: list[ColumnAlias], ) -> bool: """ - Returns True if column_aliases resolves join_key to the same canonical name - across base_table and lookup_table. + Return True when both tables' physical columns for ``join_key`` share that + key as their canonical name (via identity or ColumnAlias). """ - base_resolved = _resolve_column_via_aliases(base_table, join_key, column_aliases) - lookup_resolved = _resolve_column_via_aliases( + base_physical = _physical_column_for_join_key(base_table, join_key, column_aliases) + lookup_physical = _physical_column_for_join_key( lookup_table, join_key, column_aliases ) - return base_resolved == lookup_resolved + if base_physical == lookup_physical: + return True + return ( + _canonical_column_for_physical(base_table, base_physical, column_aliases) + == join_key + and _canonical_column_for_physical( + lookup_table, lookup_physical, column_aliases + ) + == join_key + ) # --------------------------------------------------------------------------- diff --git a/tests/genai/mapping/schema_mapping_agent/manifest/test_join_alias_resolution.py b/tests/genai/mapping/schema_mapping_agent/manifest/test_join_alias_resolution.py new file mode 100644 index 000000000..8f47d2236 --- /dev/null +++ b/tests/genai/mapping/schema_mapping_agent/manifest/test_join_alias_resolution.py @@ -0,0 +1,165 @@ +"""Join-key alias resolution: canonical join_keys ↔ physical columns.""" + +from __future__ import annotations + +from edvise.genai.mapping.schema_mapping_agent.manifest.schemas import ( + ColumnAlias, + FieldMappingManifest, + FieldMappingRecord, + JoinConfig, + RowSelectionConfig, + RowSelectionStrategy, +) +from edvise.genai.mapping.schema_mapping_agent.manifest.validation import ( + ManifestValidationErrorCode, + validate_manifest, +) +from edvise.genai.mapping.shared.schema_contract import ( + parse_enriched_schema_contract_for_sma, +) + + +def _cd(name: str) -> dict: + return { + "original_name": name, + "normalized_name": name, + "null_count": 0, + "null_percentage": 0.0, + "unique_count": 1, + "sample_values": [], + } + + +def _course_student_term_contract(): + """course has term_descr; student has term_desc (real-world naming mismatch).""" + return parse_enriched_schema_contract_for_sma( + { + "school_id": "x", + "school_name": "X", + "datasets": { + "course": { + "normalized_columns": {}, + "dtypes": {}, + "non_null_columns": [], + "unique_keys": [], + "null_tokens": [], + "boolean_map": {}, + "training": { + "file_path": "/c.csv", + "num_rows": 1, + "num_columns": 3, + "column_normalization": {"original_to_normalized": {}}, + "column_details": [ + _cd("learner_id"), + _cd("term_descr"), + _cd("course_grade"), + ], + }, + }, + "student": { + "normalized_columns": {}, + "dtypes": {}, + "non_null_columns": [], + "unique_keys": [], + "null_tokens": [], + "boolean_map": {}, + "training": { + "file_path": "/s.csv", + "num_rows": 1, + "num_columns": 3, + "column_normalization": {"original_to_normalized": {}}, + "column_details": [ + _cd("learner_id"), + _cd("term_desc"), + _cd("ugrd_grad_flag"), + ], + }, + }, + }, + } + ) + + +def _grain() -> FieldMappingRecord: + return FieldMappingRecord( + target_field="learner_id", + source_column="learner_id", + source_table="course", + join=None, + row_selection=RowSelectionConfig(strategy=RowSelectionStrategy.any_row), + confidence=1.0, + rationale="", + ) + + +def _term_degree_join(*, join_keys: list[str]) -> FieldMappingRecord: + return FieldMappingRecord( + target_field="term_degree", + source_column="ugrd_grad_flag", + source_table="student", + join=JoinConfig( + base_table="course", + lookup_table="student", + join_keys=join_keys, + ), + row_selection=RowSelectionConfig(strategy=RowSelectionStrategy.any_row), + confidence=0.9, + rationale="", + ) + + +def test_validate_manifest_accepts_canonical_join_key_with_lookup_alias(): + """ + join_keys use course's physical name (canonical); student alias bridges term_desc. + """ + contract = _course_student_term_contract() + manifest = FieldMappingManifest( + entity_type="course", + target_schema="RawEdviseCourseDataSchema", + mappings=[ + _grain(), + _term_degree_join(join_keys=["learner_id", "term_descr"]), + ], + column_aliases=[ + ColumnAlias( + table="student", + source_column="term_desc", + canonical_column="term_descr", + rationale="student term grain column name differs from course", + ) + ], + ) + errors = validate_manifest(manifest, contract) + join_codes = { + e.error_code + for e in errors + if e.target_field == "term_degree" + and e.error_code + in { + ManifestValidationErrorCode.JOIN_KEY_NOT_IN_BASE_TABLE, + ManifestValidationErrorCode.JOIN_KEY_NOT_IN_LOOKUP_TABLE, + ManifestValidationErrorCode.MISSING_COLUMN_ALIAS, + } + } + assert join_codes == set() + + +def test_validate_manifest_rejects_canonical_join_key_without_alias(): + """Incomplete refinement fix: rename join key only, no column_aliases.""" + contract = _course_student_term_contract() + manifest = FieldMappingManifest( + entity_type="course", + target_schema="RawEdviseCourseDataSchema", + mappings=[ + _grain(), + _term_degree_join(join_keys=["learner_id", "term_descr"]), + ], + column_aliases=[], + ) + errors = validate_manifest(manifest, contract) + assert any( + e.error_code == ManifestValidationErrorCode.JOIN_KEY_NOT_IN_LOOKUP_TABLE + and e.target_field == "term_degree" + and e.offending_value == "term_descr" + for e in errors + ) diff --git a/tests/genai/mapping/schema_mapping_agent/manifest/test_pass1_column_aliases.py b/tests/genai/mapping/schema_mapping_agent/manifest/test_pass1_column_aliases.py new file mode 100644 index 000000000..e0130e3f8 --- /dev/null +++ b/tests/genai/mapping/schema_mapping_agent/manifest/test_pass1_column_aliases.py @@ -0,0 +1,125 @@ +"""Pass 1 ``column_aliases_to_add`` merges into the refined manifest.""" + +from __future__ import annotations + +from edvise.genai.mapping.schema_mapping_agent.manifest.prompts.refine import ( + _apply_pass1_result, +) +from edvise.genai.mapping.schema_mapping_agent.manifest.schemas import ( + FieldMappingManifest, + FieldMappingRecord, + JoinConfig, + RowSelectionConfig, + RowSelectionStrategy, +) + + +def test_apply_pass1_result_merges_column_aliases_to_add(): + input_manifest = FieldMappingManifest( + entity_type="course", + target_schema="RawEdviseCourseDataSchema", + mappings=[ + FieldMappingRecord( + target_field="learner_id", + source_column="learner_id", + source_table="course", + join=None, + row_selection=RowSelectionConfig(strategy=RowSelectionStrategy.any_row), + confidence=1.0, + rationale="", + ), + FieldMappingRecord( + target_field="term_degree", + source_column="ugrd_grad_flag", + source_table="student", + join=JoinConfig( + base_table="course", + lookup_table="student", + join_keys=["learner_id", "term_desc"], + ), + row_selection=RowSelectionConfig(strategy=RowSelectionStrategy.any_row), + confidence=0.9, + rationale="", + ), + ], + column_aliases=[], + ) + pass1 = { + "field_statuses": { + "learner_id": "auto_approved", + "term_degree": "refined_by_llm", + }, + "refined_corrections": { + "term_degree": { + "join": { + "base_table": "course", + "lookup_table": "student", + "join_keys": ["learner_id", "term_descr"], + }, + "validation_notes": "Join key corrected; alias bridges student term_desc.", + } + }, + "column_aliases_to_add": [ + { + "table": "student", + "source_column": "term_desc", + "canonical_column": "term_descr", + "rationale": "student physical term column", + } + ], + "hitl_flags": [], + } + + refined, flags = _apply_pass1_result("inst", input_manifest, pass1) + + assert flags == [] + assert len(refined.column_aliases) == 1 + alias = refined.column_aliases[0] + assert alias.table == "student" + assert alias.source_column == "term_desc" + assert alias.canonical_column == "term_descr" + td = next(m for m in refined.mappings if m.target_field == "term_degree") + assert td.join is not None + assert td.join.join_keys == ["learner_id", "term_descr"] + assert td.review_status == "refined_by_llm" + + +def test_apply_pass1_result_dedupes_existing_alias(): + from edvise.genai.mapping.schema_mapping_agent.manifest.schemas import ColumnAlias + + input_manifest = FieldMappingManifest( + entity_type="course", + target_schema="RawEdviseCourseDataSchema", + mappings=[ + FieldMappingRecord( + target_field="learner_id", + source_column="learner_id", + source_table="course", + join=None, + row_selection=RowSelectionConfig(strategy=RowSelectionStrategy.any_row), + confidence=1.0, + rationale="", + ), + ], + column_aliases=[ + ColumnAlias( + table="student", + source_column="term_desc", + canonical_column="term_descr", + ) + ], + ) + pass1 = { + "field_statuses": {"learner_id": "auto_approved"}, + "refined_corrections": {}, + "column_aliases_to_add": [ + { + "table": "student", + "source_column": "term_desc", + "canonical_column": "term_descr", + } + ], + "hitl_flags": [], + } + refined, _ = _apply_pass1_result("inst", input_manifest, pass1) + assert len(refined.column_aliases) == 1