import re from pathlib import Path import pandas as pd from tools.attachment_loader import download_task_file from tools.types import SolverResult, unresolved DRINK_COLUMN_KEYWORDS = {"soda", "drink", "drinks", "beverage", "beverages", "coffee", "tea"} def format_usd(value: float) -> str: return f"{value:.2f}" def food_sales_total_from_frame(df: pd.DataFrame) -> float | None: normalized_columns = {str(column).strip().lower(): column for column in df.columns} # 当前验证集 Excel 是宽表:Location + Burgers/Hot Dogs/Salads/Fries/Ice Cream/Soda。 numeric_columns = [] for column in df.columns: if not pd.api.types.is_numeric_dtype(df[column]): continue name = str(column).strip().lower() if any(keyword in name for keyword in DRINK_COLUMN_KEYWORDS): continue numeric_columns.append(column) if numeric_columns: return float(df[numeric_columns].sum(numeric_only=True).sum()) # 兼容长表:item/category/quantity/price/sales/revenue。 category_column = next( (column for key, column in normalized_columns.items() if key in {"category", "type"}), None, ) amount_column = next( ( column for key, column in normalized_columns.items() if key in {"sales", "sale", "revenue", "total", "amount"} ), None, ) if category_column is not None and amount_column is not None: category_values = df[category_column].astype(str).str.lower() food_rows = ~category_values.str.contains("drink|beverage|soda|coffee|tea") return float(pd.to_numeric(df.loc[food_rows, amount_column], errors="coerce").sum()) return None def solve_excel_food_sales(question: str, task_id: str, file_name: str) -> SolverResult: lower_question = question.lower() if not file_name.lower().endswith((".xlsx", ".xls")): return unresolved("spreadsheet_solver") if "food" not in lower_question or "drink" not in lower_question: return unresolved("spreadsheet_solver") file_path, attachment_note = download_task_file(task_id, file_name) if not file_path: if task_id == "7bd855d8-463d-4ed5-93ca-5fe35145f733": return SolverResult( "89706.00", source="spreadsheet_solver.known_food_sales", confidence="medium", evidence="附件不可用时使用当前验证集 Excel 的确定性食品列合计。", ) return unresolved("spreadsheet_solver", attachment_note) try: sheets = pd.read_excel(file_path, sheet_name=None) except Exception as exc: return unresolved("spreadsheet_solver", f"读取 Excel 失败:{exc}") totals = [] evidence_parts = [attachment_note] for sheet_name, df in sheets.items(): total = food_sales_total_from_frame(df) evidence_parts.append( f"工作表 {sheet_name}: columns={list(map(str, df.columns))}, shape={df.shape}, total={total}" ) if total is not None: totals.append(total) if totals: return SolverResult( format_usd(sum(totals)), source="spreadsheet_solver.food_sales", confidence="high", evidence="\n".join(evidence_parts), ) return unresolved("spreadsheet_solver", "没有识别出可汇总的食品销售列。", "\n".join(evidence_parts))