diff --git a/app.py b/app.py index b406335..407e3f2 100644 --- a/app.py +++ b/app.py @@ -26,7 +26,7 @@ get_product_schema_description, validate_product_schema, ) -from excel_processing import preview_excel_import, process_excel_file +from excel_processing import estimate_import_impact, preview_excel_import, process_excel_file from guardrails import ( DestructiveActionApprovalRequired, SchemaChangeApprovalRequired, @@ -306,9 +306,21 @@ def _get_cached_import_preview(uploaded_file, db_path): st.write("Column names in the uploaded file:", import_preview["dataframe"].columns.tolist()) st.write("Resolved column mappings:", import_preview["column_mappings"]) if action in {"remove", "modify"}: + impact = estimate_import_impact(import_preview, db_path, action) st.warning( - f"The '{action}' action changes or removes existing inventory rows." + f"The '{action}' action changes or removes existing inventory rows. " + f"File rows: {impact['rows_in_file']}; " + f"matching DB rows: {impact['matched_db_rows']}." ) + if impact["ambiguous_names"]: + st.error( + "Ambiguous NAME matches (include an ID column to target a specific row): " + + ", ".join(impact["ambiguous_names"]) + ) + if impact["missing_identity_rows"]: + st.error( + f"{impact['missing_identity_rows']} file row(s) are missing both ID and NAME." + ) approve_destructive_action = st.checkbox( f"Approve the '{action}' action for this import", key=f"approve_destructive_{action}", diff --git a/excel_processing.py b/excel_processing.py index a903851..f9cd385 100644 --- a/excel_processing.py +++ b/excel_processing.py @@ -4,7 +4,10 @@ This module handles the processing of uploaded Excel files and updates the database accordingly. """ +from __future__ import annotations + import sqlite3 +from typing import Any import pandas as pd @@ -20,12 +23,116 @@ from utils import _normalize_identifier, map_columns +class AmbiguousProductMatchError(ValueError): + """Raised when a NAME-based remove/modify would touch more than one row.""" + + def _read_excel_frame(uploaded_file): if hasattr(uploaded_file, "seek"): uploaded_file.seek(0) return pd.read_excel(uploaded_file) +def _mapped_row_from_excel_row(row: Any, column_mappings: dict[str, str]) -> dict[str, Any]: + unmapped = [str(col) for col in row.keys() if str(col) not in column_mappings] + if unmapped: + raise ValueError( + f"AI column mapping is incomplete — no mapping for: {', '.join(unmapped)}" + ) + return {column_mappings[str(col)]: value for col, value in row.items()} + + +def _has_usable_id(mapped_row: dict[str, Any]) -> bool: + if "ID" not in mapped_row: + return False + value = mapped_row.get("ID") + if value is None: + return False + try: + if isinstance(value, float) and value != value: # NaN + return False + except Exception: + pass + text = str(value).strip() + return bool(text) and text.lower() != "nan" + + +def _resolve_row_identity(mapped_row: dict[str, Any], action: str) -> tuple[str, Any]: + """Return (mode, key) where mode is 'ID' or 'NAME'.""" + + if _has_usable_id(mapped_row): + return "ID", mapped_row["ID"] + + if "NAME" not in mapped_row or mapped_row.get("NAME") is None: + raise ValueError( + f"Row is missing an ID or NAME mapping; cannot determine which product to {action}." + ) + name = mapped_row.get("NAME") + if str(name).strip() == "" or str(name).strip().lower() == "nan": + raise ValueError( + f"Row is missing an ID or NAME mapping; cannot determine which product to {action}." + ) + return "NAME", name + + +def _count_matches(cursor: sqlite3.Cursor, mode: str, key: Any) -> int: + cursor.execute( + f"SELECT COUNT(*) FROM PRODUCT WHERE {quote_identifier(mode)}=?", + (key,), + ) + return int(cursor.fetchone()[0]) + + +def _enforce_unique_name_match(cursor: sqlite3.Cursor, mode: str, key: Any, action: str) -> None: + """Fail closed when NAME matches multiple products for destructive actions.""" + + if mode != "NAME" or action not in {"remove", "modify"}: + return + match_count = _count_matches(cursor, mode, key) + if match_count > 1: + raise AmbiguousProductMatchError( + f"Ambiguous {action}: NAME={key!r} matches {match_count} products. " + "Include an ID column to target a specific row." + ) + + +def estimate_import_impact(preview: dict[str, Any], db_path: str, action: str) -> dict[str, Any]: + """Estimate how many DB rows an import action would touch.""" + + df = preview["dataframe"] + column_mappings = preview["column_mappings"] + rows_in_file = len(df) + matched_rows = 0 + ambiguous_names: list[str] = [] + identity_modes: set[str] = set() + missing_identity = 0 + + with sqlite3.connect(db_path) as conn: + cursor = conn.cursor() + for _, row in df.iterrows(): + try: + mapped_row = _mapped_row_from_excel_row(row, column_mappings) + mode, key = _resolve_row_identity(mapped_row, action) + except ValueError: + missing_identity += 1 + continue + + identity_modes.add(mode) + match_count = _count_matches(cursor, mode, key) + matched_rows += match_count + if mode == "NAME" and match_count > 1 and action in {"remove", "modify"}: + ambiguous_names.append(str(key)) + + return { + "action": action, + "rows_in_file": rows_in_file, + "matched_db_rows": matched_rows, + "identity_modes": sorted(identity_modes), + "ambiguous_names": sorted(set(ambiguous_names)), + "missing_identity_rows": missing_identity, + } + + def preview_excel_import(uploaded_file, db_path, *, emit_audit_event=False): """Return the AI-produced column mapping and any pending schema changes.""" @@ -89,6 +196,7 @@ def process_excel_file( "allow_destructive_actions": allow_destructive_actions, } processed_rows = 0 + affected_db_rows = 0 try: enforce_destructive_action_policy( @@ -120,48 +228,55 @@ def process_excel_file( # Process each row in the Excel file for _, row in df.iterrows(): - unmapped = [str(col) for col in row.keys() if str(col) not in column_mappings] - if unmapped: - raise ValueError( - f"AI column mapping is incomplete — no mapping for: {', '.join(unmapped)}" - ) - mapped_row = {column_mappings[str(col)]: value for col, value in row.items()} - - if "NAME" not in mapped_row: - raise ValueError( - f"Row is missing a NAME mapping; cannot determine which product to {action}." - ) + mapped_row = _mapped_row_from_excel_row(row, column_mappings) + mode, key = _resolve_row_identity(mapped_row, action) + _enforce_unique_name_match(cursor, mode, key, action) if action == "remove": cursor.execute( - f"DELETE FROM PRODUCT WHERE {quote_identifier('NAME')}=?", - (mapped_row.get('NAME'),), + f"DELETE FROM PRODUCT WHERE {quote_identifier(mode)}=?", + (key,), ) + affected_db_rows += cursor.rowcount if cursor.rowcount is not None else 0 elif action == "modify": - set_clause = ", ".join([f"{quote_identifier(col)}=?" for col in mapped_row.keys()]) + set_clause = ", ".join( + [f"{quote_identifier(col)}=?" for col in mapped_row.keys()] + ) values = tuple(mapped_row.values()) cursor.execute( - f"UPDATE PRODUCT SET {set_clause} WHERE {quote_identifier('NAME')}=?", - values + (mapped_row.get('NAME'),), + f"UPDATE PRODUCT SET {set_clause} WHERE {quote_identifier(mode)}=?", + values + (key,), ) + affected_db_rows += cursor.rowcount if cursor.rowcount is not None else 0 else: # add action (or update if product exists) + # Prefer ID when present; otherwise upsert by NAME. cursor.execute( - f"SELECT * FROM PRODUCT WHERE {quote_identifier('NAME')}=?", - (mapped_row.get('NAME'),), + f"SELECT * FROM PRODUCT WHERE {quote_identifier(mode)}=?", + (key,), ) existing_product = cursor.fetchone() if existing_product: - set_clause = ", ".join([f"{quote_identifier(col)}=?" for col in mapped_row.keys()]) + set_clause = ", ".join( + [f"{quote_identifier(col)}=?" for col in mapped_row.keys()] + ) values = tuple(mapped_row.values()) cursor.execute( - f"UPDATE PRODUCT SET {set_clause} WHERE {quote_identifier('NAME')}=?", - values + (mapped_row.get('NAME'),), + f"UPDATE PRODUCT SET {set_clause} WHERE {quote_identifier(mode)}=?", + values + (key,), ) else: - columns = ", ".join(quote_identifier(col) for col in mapped_row.keys()) - placeholders = ", ".join(["?" for _ in mapped_row]) - values = tuple(mapped_row.values()) - cursor.execute(f"INSERT INTO PRODUCT ({columns}) VALUES ({placeholders})", values) + insert_row = dict(mapped_row) + # Avoid inserting an explicit NULL/empty ID into an AUTOINCREMENT PK. + if mode == "NAME" and "ID" in insert_row and not _has_usable_id(insert_row): + insert_row.pop("ID", None) + columns = ", ".join(quote_identifier(col) for col in insert_row.keys()) + placeholders = ", ".join(["?" for _ in insert_row]) + values = tuple(insert_row.values()) + cursor.execute( + f"INSERT INTO PRODUCT ({columns}) VALUES ({placeholders})", + values, + ) + affected_db_rows += 1 processed_rows += 1 append_audit_event( db_path, @@ -169,6 +284,7 @@ def process_excel_file( { **audit_details, "processed_rows": processed_rows, + "affected_db_rows": affected_db_rows, "status": "success", }, ) @@ -179,6 +295,7 @@ def process_excel_file( { **audit_details, "processed_rows": processed_rows, + "affected_db_rows": affected_db_rows, "status": "blocked" if isinstance(exc, GuardrailViolation) else "failed", "error": str(exc), }, diff --git a/prompt.py b/prompt.py index b44ed57..2850f91 100644 --- a/prompt.py +++ b/prompt.py @@ -19,10 +19,6 @@ COLUMN_MAPPING_PROMPT_VERSION = "v1" -def _normalize_token(value: str) -> str: - return re.sub(r"[^A-Z0-9]+", "", value.upper()) - - def _fallback_sql(question: str) -> str: text = question.lower().strip() @@ -58,41 +54,10 @@ def _fallback_column_mapping(prompt: str) -> str: if database_match: database_columns = [value.strip() for value in database_match.group(1).split(",") if value.strip()] - mapping: dict[str, str] = {} - db_lookup = {_normalize_token(column): column for column in database_columns} - - aliases = { - "NAME": ("NAME", "PRODUCTNAME", "ITEMNAME", "PRODUCT"), - "CATEGORY": ("CATEGORY", "TYPE", "GROUP"), - "BRAND": ("BRAND",), - "PRICE": ("PRICE", "COST", "AMOUNT", "RATE"), - "STOCK": ("STOCK", "QUANTITY", "QTY", "INVENTORY"), - "QUANTITY": ("QUANTITY", "STOCK", "QTY", "INVENTORY"), - "COLOR": ("COLOR", "COLOUR"), - "SIZE": ("SIZE",), - "WEIGHT": ("WEIGHT",), - "SPECIFICATIONS": ("SPECIFICATIONS", "SPECIFICATION", "DETAILS", "DESCRIPTION", "SPEC"), - "ID": ("ID", "PRODUCTID"), - } - - for excel_column in excel_columns: - normalized = _normalize_token(excel_column) - if normalized in db_lookup: - mapping[excel_column] = db_lookup[normalized] - continue - - mapped_column = None - for target, candidates in aliases.items(): - if normalized in candidates: - mapped_column = db_lookup.get(_normalize_token(target), target) - break - - if mapped_column is None: - mapped_column = re.sub(r"[^A-Z0-9]+", "_", excel_column.upper()).strip("_") or "COLUMN" - - mapping[excel_column] = mapped_column + # Lazy import avoids a circular dependency (utils imports prompt builders). + from utils import _heuristic_column_mapping - return json.dumps(mapping) + return json.dumps(_heuristic_column_mapping(excel_columns, database_columns)) def get_sql_prompt_metadata() -> dict[str, str]: diff --git a/tests/test_audit.py b/tests/test_audit.py index bcdcca4..b73a59a 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -98,3 +98,22 @@ def fake_response(prompt: str) -> str: # The prompt must contain the columns so the model can produce a mapping. assert "Name" in captured["prompt"] assert "NAME" in captured["prompt"] + + +def test_heuristic_and_prompt_fallback_agree_on_aliases(): + from prompt import _fallback_column_mapping, build_column_mapping_prompt + from utils import _heuristic_column_mapping + + excel_columns = ["Product Name", "Qty", "Colour", "Cost"] + database_columns = ["ID", "NAME", "PRICE", "STOCK", "COLOR"] + + heuristic = _heuristic_column_mapping(excel_columns, database_columns) + fallback = json.loads( + _fallback_column_mapping(build_column_mapping_prompt(excel_columns, database_columns)) + ) + + assert heuristic == fallback + assert heuristic["Product Name"] == "NAME" + assert heuristic["Qty"] == "STOCK" + assert heuristic["Colour"] == "COLOR" + assert heuristic["Cost"] == "PRICE" diff --git a/tests/test_pytest_regressions.py b/tests/test_pytest_regressions.py index 0859cc3..3502cef 100644 --- a/tests/test_pytest_regressions.py +++ b/tests/test_pytest_regressions.py @@ -24,6 +24,9 @@ def __init__(self, rows): self._rows = list(rows) self.columns = FakeColumns(list(self._rows[0].keys()) if self._rows else []) + def __len__(self): + return len(self._rows) + def iterrows(self): if not self._rows: return @@ -421,6 +424,140 @@ def test_process_excel_file_blocks_destructive_actions_without_approval( assert events[-1]["details"]["error"].endswith("Blocked action: REMOVE.") +def test_process_excel_file_remove_by_id_with_duplicate_names( + excel_processing_module, + inventory_db: Path, +): + with sqlite3.connect(inventory_db) as connection: + connection.execute( + """ + INSERT INTO PRODUCT + (NAME, CATEGORY, BRAND, PRICE, STOCK, SIZE, COLOR, WEIGHT, SPECIFICATIONS) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ("Widget", "Gadgets", "Acme", 1.0, 1, "S", "Red", 0.5, "Duplicate name"), + ) + widget_ids = [ + row[0] + for row in connection.execute( + "SELECT ID FROM PRODUCT WHERE NAME = ? ORDER BY ID", + ("Widget",), + ) + ] + target_id = widget_ids[0] + + excel_processing_module.pd.read_excel = lambda uploaded_file: FakeFrame( + [{"Id": target_id, "Name": "Widget"}] + ) + + excel_processing_module.process_excel_file( + object(), + str(inventory_db), + "remove", + allow_destructive_actions=True, + ) + + with sqlite3.connect(inventory_db) as connection: + remaining_widget_ids = [ + row[0] + for row in connection.execute( + "SELECT ID FROM PRODUCT WHERE NAME = ? ORDER BY ID", + ("Widget",), + ) + ] + all_names = [row[0] for row in connection.execute("SELECT NAME FROM PRODUCT ORDER BY NAME")] + + assert target_id not in remaining_widget_ids + assert len(remaining_widget_ids) == 1 + assert "Gizmo" in all_names + + +def test_process_excel_file_remove_rejects_ambiguous_name( + excel_processing_module, + inventory_db: Path, +): + with sqlite3.connect(inventory_db) as connection: + connection.execute( + """ + INSERT INTO PRODUCT + (NAME, CATEGORY, BRAND, PRICE, STOCK, SIZE, COLOR, WEIGHT, SPECIFICATIONS) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ("Widget", "Gadgets", "Acme", 1.0, 1, "S", "Red", 0.5, "Duplicate name"), + ) + + excel_processing_module.pd.read_excel = lambda uploaded_file: FakeFrame([{"Name": "Widget"}]) + + with pytest.raises(excel_processing_module.AmbiguousProductMatchError, match="Ambiguous remove"): + excel_processing_module.process_excel_file( + object(), + str(inventory_db), + "remove", + allow_destructive_actions=True, + ) + + with sqlite3.connect(inventory_db) as connection: + widget_count = connection.execute( + "SELECT COUNT(*) FROM PRODUCT WHERE NAME = ?", + ("Widget",), + ).fetchone()[0] + assert widget_count == 2 + + +def test_process_excel_file_modify_by_id( + excel_processing_module, + inventory_db: Path, +): + with sqlite3.connect(inventory_db) as connection: + widget_id = connection.execute( + "SELECT ID FROM PRODUCT WHERE NAME = ?", + ("Widget",), + ).fetchone()[0] + + excel_processing_module.pd.read_excel = lambda uploaded_file: FakeFrame( + [{"Id": widget_id, "Name": "Widget", "Category": "ID Updated", "Stock": 99}] + ) + + excel_processing_module.process_excel_file( + object(), + str(inventory_db), + "modify", + allow_destructive_actions=True, + ) + + with sqlite3.connect(inventory_db) as connection: + row = connection.execute( + "SELECT ID, NAME, CATEGORY, STOCK FROM PRODUCT WHERE ID = ?", + (widget_id,), + ).fetchone() + + assert row == (widget_id, "Widget", "ID Updated", 99) + + +def test_estimate_import_impact_reports_ambiguous_names( + excel_processing_module, + inventory_db: Path, +): + with sqlite3.connect(inventory_db) as connection: + connection.execute( + """ + INSERT INTO PRODUCT + (NAME, CATEGORY, BRAND, PRICE, STOCK, SIZE, COLOR, WEIGHT, SPECIFICATIONS) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ("Widget", "Gadgets", "Acme", 1.0, 1, "S", "Red", 0.5, "Duplicate name"), + ) + + excel_processing_module.pd.read_excel = lambda uploaded_file: FakeFrame([{"Name": "Widget"}]) + preview = excel_processing_module.preview_excel_import(object(), str(inventory_db)) + impact = excel_processing_module.estimate_import_impact(preview, str(inventory_db), "remove") + + assert impact["rows_in_file"] == 1 + assert impact["matched_db_rows"] == 2 + assert impact["ambiguous_names"] == ["Widget"] + assert impact["identity_modes"] == ["NAME"] + + def test_get_gemini_response_uses_mocked_sdk_when_api_key_is_available(monkeypatch): calls = [] diff --git a/utils.py b/utils.py index b7c27fc..9a45c8d 100644 --- a/utils.py +++ b/utils.py @@ -76,6 +76,27 @@ def _normalize_identifier(value: str) -> str: return re.sub(r"[^A-Z0-9]+", "_", value.upper()).strip("_") +def _normalize_alias_token(value: str) -> str: + """Alphanumeric-only token used for synonym matching (e.g. Product Name → PRODUCTNAME).""" + + return re.sub(r"[^A-Z0-9]+", "", value.upper()) + + +COLUMN_MAPPING_SYNONYMS: dict[str, tuple[str, ...]] = { + "NAME": ("NAME", "PRODUCTNAME", "ITEMNAME", "PRODUCT"), + "CATEGORY": ("CATEGORY", "TYPE", "GROUP"), + "BRAND": ("BRAND",), + "PRICE": ("PRICE", "COST", "AMOUNT", "RATE"), + "STOCK": ("STOCK", "QUANTITY", "QTY", "INVENTORY"), + "QUANTITY": ("QUANTITY", "STOCK", "QTY", "INVENTORY"), + "COLOR": ("COLOR", "COLOUR"), + "SIZE": ("SIZE",), + "WEIGHT": ("WEIGHT",), + "SPECIFICATIONS": ("SPECIFICATIONS", "SPECIFICATION", "DETAILS", "DESCRIPTION", "SPEC"), + "ID": ("ID", "PRODUCTID"), +} + + def _existing_columns(db_path: str) -> list[str]: with sqlite3.connect(db_path) as connection: cursor = connection.cursor() @@ -194,29 +215,21 @@ def _heuristic_column_mapping( existing_lookup = {_normalize_identifier(column): column for column in existing_columns} result: dict[str, str] = {} - synonym_groups = { - "NAME": ("NAME", "PRODUCTNAME", "ITEMNAME", "PRODUCT"), - "CATEGORY": ("CATEGORY", "TYPE", "GROUP"), - "BRAND": ("BRAND",), - "PRICE": ("PRICE", "COST", "AMOUNT", "RATE"), - "STOCK": ("STOCK", "QUANTITY", "QTY", "INVENTORY"), - "QUANTITY": ("QUANTITY", "STOCK", "QTY", "INVENTORY"), - "COLOR": ("COLOR", "COLOUR"), - "SIZE": ("SIZE",), - "WEIGHT": ("WEIGHT",), - "SPECIFICATIONS": ("SPECIFICATIONS", "SPECIFICATION", "DETAILS", "DESCRIPTION", "SPEC"), - "ID": ("ID", "PRODUCTID"), - } - for excel_column in excel_columns: normalized = _normalize_identifier(str(excel_column)) + alias_token = _normalize_alias_token(str(excel_column)) if normalized in existing_lookup: result[str(excel_column)] = existing_lookup[normalized] continue mapped_column = None - for target, aliases in synonym_groups.items(): - if normalized == target or normalized in aliases: + for target, aliases in COLUMN_MAPPING_SYNONYMS.items(): + if ( + normalized == target + or normalized in aliases + or alias_token == target + or alias_token in aliases + ): mapped_column = existing_lookup.get(_normalize_identifier(target), target) break