Spaces:
Runtime error
Runtime error
Download src/loader.py from BinbomJared/DM-Agent: direct link, hf CLI and curl.
- Browser
- Download file 18.5 kB
-
https://huggingface.co/spaces/BinbomJared/DM-Agent/resolve/main/src/loader.py
- Command line
-
hf download hf://spaces/BinbomJared/DM-Agent/src/loader.py
-
curl -L -o loader.py https://huggingface.co/spaces/BinbomJared/DM-Agent/resolve/main/src/loader.py
18.5 kB
| from __future__ import annotations | |
| """ | |
| 数据加载:支持 Google Sheets(优先)和本地 Excel 两种数据源。 | |
| 处理编码异常列名、过滤汇总行、从 Date 派生 year/week/month。 | |
| """ | |
| import logging | |
| import os | |
| from pathlib import Path | |
| from typing import Any | |
| import pandas as pd | |
| import yaml | |
| log = logging.getLogger("ad-agent") | |
| def _project_root() -> Path: | |
| return Path(__file__).resolve().parent.parent | |
| def load_config(config_path: Path | None = None) -> dict[str, Any]: | |
| path = config_path or (_project_root() / "config.yaml") | |
| with open(path, "r", encoding="utf-8") as f: | |
| return yaml.safe_load(f) | |
| # ───────────────────────────────────────────── | |
| # Google Sheets 数据源 | |
| # ───────────────────────────────────────────── | |
| def _build_google_creds(config: dict[str, Any]): | |
| """ | |
| 构建 Google API 凭证,支持三种认证方式(按优先级): | |
| 1. GOOGLE_CREDENTIALS 环境变量(服务账号 JSON 字符串,云端部署推荐) | |
| 2. GOOGLE_OAUTH_REFRESH_TOKEN 环境变量(OAuth refresh token,GitHub Actions 等场景) | |
| 3. credentials.json 本地文件(本地开发) | |
| """ | |
| from google.oauth2.service_account import Credentials as SACredentials | |
| from google.oauth2.credentials import Credentials as OAuthCredentials | |
| import json | |
| scopes = [ | |
| "https://www.googleapis.com/auth/spreadsheets", | |
| "https://www.googleapis.com/auth/drive", | |
| ] | |
| creds_json = os.environ.get("GOOGLE_CREDENTIALS", "").strip() | |
| if creds_json: | |
| info = json.loads(creds_json) | |
| creds = SACredentials.from_service_account_info(info, scopes=scopes) | |
| log.info("使用环境变量 GOOGLE_CREDENTIALS(服务账号)认证") | |
| return creds | |
| refresh_token = os.environ.get("GOOGLE_OAUTH_REFRESH_TOKEN", "").strip() | |
| client_id = os.environ.get("GOOGLE_OAUTH_CLIENT_ID", "").strip() | |
| client_secret = os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET", "").strip() | |
| if refresh_token and client_id and client_secret: | |
| creds = OAuthCredentials( | |
| token=None, | |
| refresh_token=refresh_token, | |
| client_id=client_id, | |
| client_secret=client_secret, | |
| token_uri="https://oauth2.googleapis.com/token", | |
| ) | |
| log.info("使用 OAuth refresh token 认证") | |
| return creds | |
| gs_cfg = config.get("google_sheets", {}) | |
| creds_file = gs_cfg.get("credentials_file", "credentials.json") | |
| creds_path = Path(creds_file) | |
| if not creds_path.is_absolute(): | |
| creds_path = _project_root() / creds_path | |
| if creds_path.exists(): | |
| creds = SACredentials.from_service_account_file(str(creds_path), scopes=scopes) | |
| log.info(f"使用凭证文件 {creds_path} 认证") | |
| return creds | |
| return None | |
| def _load_from_google_sheets(config: dict[str, Any]) -> pd.DataFrame | None: | |
| """ | |
| 从 Google Sheets 分平台表(Meta/Google/Criteo/Bing)读取 API 同步数据。 | |
| 单一数据源,不读 Export 表,避免多源合并导致的重复。 | |
| """ | |
| gs_cfg = config.get("google_sheets", {}) | |
| if not gs_cfg.get("enabled"): | |
| return None | |
| url = gs_cfg.get("spreadsheet_url", "").strip() | |
| if not url: | |
| log.warning("Google Sheets 已启用但未配置 spreadsheet_url,回退到本地 Excel") | |
| return None | |
| try: | |
| import gspread | |
| creds = _build_google_creds(config) | |
| if creds is None: | |
| log.warning("无可用 Google 认证凭证,回退到本地 Excel") | |
| return None | |
| gc = gspread.authorize(creds) | |
| sh = gc.open_by_url(url) | |
| platform_sheets = gs_cfg.get("platform_sheets", ["Meta", "Google", "Criteo", "Bing"]) | |
| col_cfg = config.get("columns", {}) | |
| idx_rename = { | |
| col_cfg.get("country_idx", 0): "country", | |
| col_cfg.get("platform_idx", 1): "platform", | |
| col_cfg.get("campaign_type_idx", 2): "campaign_type", | |
| col_cfg.get("campaign_name_idx", 3): "campaign_name", | |
| col_cfg.get("date_idx", 4): "date", | |
| } | |
| def _read_one_sheet(ws_name: str) -> pd.DataFrame | None: | |
| try: | |
| ws = sh.worksheet(ws_name) | |
| records = ws.get_all_values() | |
| if len(records) < 2: | |
| return None | |
| header = records[0] | |
| seen: dict[str, int] = {} | |
| unique_header = [] | |
| for col in header: | |
| if col in seen: | |
| seen[col] += 1 | |
| unique_header.append(f"{col}_{seen[col]}") | |
| else: | |
| seen[col] = 0 | |
| unique_header.append(col) | |
| sdf = pd.DataFrame(records[1:], columns=unique_header) | |
| rename = {} | |
| for idx, name in idx_rename.items(): | |
| if idx < len(sdf.columns): | |
| rename[sdf.columns[idx]] = name | |
| sdf = sdf.rename(columns=rename) | |
| log.info(f"从 Sheet [{ws_name}] 读取 {len(sdf)} 行") | |
| return sdf | |
| except Exception: | |
| return None | |
| all_dfs: list[pd.DataFrame] = [] | |
| for ps in platform_sheets: | |
| ps_df = _read_one_sheet(ps) | |
| if ps_df is not None and not ps_df.empty: | |
| all_dfs.append(ps_df) | |
| if not all_dfs: | |
| log.warning("Google Sheets 所有分平台工作表均无数据") | |
| return None | |
| df = pd.concat(all_dfs, ignore_index=True) | |
| log.info(f"从 Google Sheets 分平台表共读取 {len(df)} 行数据({', '.join(platform_sheets)})") | |
| return df | |
| except Exception as e: | |
| log.error(f"Google Sheets 读取失败: {e},回退到本地 Excel") | |
| return None | |
| # ───────────────────────────────────────────── | |
| # 本地 Excel 数据源 | |
| # ───────────────────────────────────────────── | |
| def _load_from_excel(config: dict[str, Any]) -> pd.DataFrame: | |
| """从 data 目录读取最新的 .xlsx 文件。""" | |
| root = _project_root() | |
| data_dir = Path(config.get("data_dir", "data")) | |
| if not data_dir.is_absolute(): | |
| data_dir = root / data_dir | |
| files = sorted( | |
| (f for f in data_dir.glob("*.xlsx") if not f.name.startswith("~$")), | |
| key=lambda p: p.stat().st_mtime, reverse=True, | |
| ) | |
| if not files: | |
| raise FileNotFoundError(f"在 {data_dir} 中未找到 .xlsx 文件") | |
| sheet = config.get("excel", {}).get("sheet", 0) | |
| df = pd.read_excel(files[0], sheet_name=sheet, engine="openpyxl") | |
| log.info(f"从本地 Excel 读取: {files[0].name}({len(df)} 行)") | |
| return df | |
| # ───────────────────────────────────────────── | |
| # 写回 Google Sheets | |
| # ───────────────────────────────────────────── | |
| def _get_spreadsheet(config: dict[str, Any]): | |
| """获取 gspread spreadsheet 对象(读写权限)。""" | |
| import gspread | |
| gs_cfg = config.get("google_sheets", {}) | |
| url = gs_cfg.get("spreadsheet_url", "").strip() | |
| if not url: | |
| raise RuntimeError("未配置 spreadsheet_url") | |
| creds = _build_google_creds(config) | |
| if creds is None: | |
| raise RuntimeError("无可用 Google 认证凭证(需要 GOOGLE_CREDENTIALS、GOOGLE_OAUTH_REFRESH_TOKEN 或 credentials.json)") | |
| gc = gspread.authorize(creds) | |
| return gc.open_by_url(url) | |
| def _get_or_create_worksheet(sh, sheet_name: str, header: list[str]): | |
| """获取或创建指定名称的 worksheet,确保有表头。""" | |
| import gspread | |
| try: | |
| ws = sh.worksheet(sheet_name) | |
| except gspread.WorksheetNotFound: | |
| ws = sh.add_worksheet(title=sheet_name, rows=1000, cols=len(header)) | |
| ws.append_row(header, value_input_option="RAW") | |
| log.info(f"创建新 Sheet: {sheet_name}") | |
| return ws | |
| def _build_header(config: dict[str, Any]) -> list[str]: | |
| """根据 config 构建标准表头。""" | |
| col_cfg = config.get("columns", {}) | |
| dim_count = 5 | |
| if col_cfg.get("ad_id_idx") is not None: | |
| dim_count = max(dim_count, col_cfg["ad_id_idx"] + 1) | |
| dims = [""] * dim_count | |
| dims[col_cfg.get("country_idx", 0)] = "Country" | |
| dims[col_cfg.get("platform_idx", 1)] = "Platform" | |
| dims[col_cfg.get("campaign_type_idx", 2)] = "Campaign Type" | |
| dims[col_cfg.get("campaign_name_idx", 3)] = "Campaign Name" | |
| dims[col_cfg.get("date_idx", 4)] = "Date" | |
| if col_cfg.get("adset_name_idx") is not None: | |
| dims[col_cfg["adset_name_idx"]] = "Ad Set Name" | |
| if col_cfg.get("ad_name_idx") is not None: | |
| dims[col_cfg["ad_name_idx"]] = "Ad Name" | |
| if col_cfg.get("ad_id_idx") is not None: | |
| dims[col_cfg["ad_id_idx"]] = "Ad ID" | |
| return dims + config.get("metrics", []) | |
| def write_api_data_to_sheets( | |
| df: pd.DataFrame, | |
| config_path: Path | None = None, | |
| ) -> dict[str, int]: | |
| """ | |
| 将 API 拉取的数据按平台写入不同的 Google Sheets worksheet。 | |
| 每个平台一个 sheet(如 Meta、Google、Criteo),自动创建不存在的 sheet。 | |
| 按 (country, date, campaign_name, ad_name) 去重,只追加新行。 | |
| """ | |
| config = load_config(config_path) | |
| sh = _get_spreadsheet(config) | |
| col_cfg = config.get("columns", {}) | |
| country_idx = col_cfg.get("country_idx", 0) | |
| date_idx = col_cfg.get("date_idx", 4) | |
| campaign_name_idx = col_cfg.get("campaign_name_idx", 3) | |
| metrics_list = config.get("metrics", []) | |
| header = _build_header(config) | |
| key_indices = [country_idx, date_idx, campaign_name_idx] | |
| platforms = df["platform"].unique() if "platform" in df.columns else [] | |
| total_added = 0 | |
| for platform in platforms: | |
| p_df = df[df["platform"] == platform] | |
| ws = _get_or_create_worksheet(sh, platform, header) | |
| all_values = ws.get_all_values() | |
| col_order = all_values[0] if all_values else header | |
| existing_rows = all_values[1:] if len(all_values) > 1 else [] | |
| existing_keys = set() | |
| for row in existing_rows: | |
| if len(row) > max(key_indices): | |
| key = tuple(row[i].strip() for i in key_indices) | |
| existing_keys.add(key) | |
| new_rows = [] | |
| for _, row_data in p_df.iterrows(): | |
| dt = row_data.get("date") | |
| date_str = dt.strftime("%Y-%m-%d") if hasattr(dt, "strftime") else str(dt).split(" ")[0] | |
| campaign_name = str(row_data.get("campaign_name", "")) | |
| ad_name = str(row_data.get("ad_name", "")) | |
| country = str(row_data.get("country", "")) | |
| key = (country, date_str, campaign_name, ad_name) | |
| if key in existing_keys: | |
| continue | |
| out_row = [""] * len(col_order) | |
| dim_map = { | |
| col_cfg.get("country_idx", 0): str(row_data.get("country", "")), | |
| col_cfg.get("platform_idx", 1): platform, | |
| col_cfg.get("campaign_type_idx", 2): str(row_data.get("campaign_type", "")), | |
| campaign_name_idx: campaign_name, | |
| date_idx: date_str, | |
| } | |
| adset_name_idx = col_cfg.get("adset_name_idx") | |
| ad_name_idx = col_cfg.get("ad_name_idx") | |
| ad_id_idx = col_cfg.get("ad_id_idx") | |
| if adset_name_idx is not None: | |
| dim_map[adset_name_idx] = str(row_data.get("adset_name", "")) | |
| if ad_name_idx is not None: | |
| dim_map[ad_name_idx] = ad_name | |
| if ad_id_idx is not None: | |
| dim_map[ad_id_idx] = str(row_data.get("ad_id", "")) | |
| for idx, val in dim_map.items(): | |
| if idx < len(out_row): | |
| out_row[idx] = val | |
| for col_name in col_order: | |
| if col_name in metrics_list and col_name in row_data.index: | |
| val = row_data[col_name] | |
| if val is None or (isinstance(val, float) and val != val): | |
| continue | |
| col_idx = col_order.index(col_name) | |
| v = float(val) | |
| out_row[col_idx] = str(int(v)) if v == int(v) else str(v) | |
| new_rows.append(out_row) | |
| existing_keys.add(key) | |
| if new_rows: | |
| ws.append_rows(new_rows, value_input_option="USER_ENTERED") | |
| log.info(f"写入 Sheet [{platform}]: {len(new_rows)} 行新数据") | |
| total_added += len(new_rows) | |
| return {"added": total_added} | |
| # ───────────────────────────────────────────── | |
| # 统一加载入口 | |
| # ───────────────────────────────────────────── | |
| _data_source: str = "unknown" | |
| def load_raw(config_path: Path | None = None) -> pd.DataFrame: | |
| """优先 Google Sheets 分平台表,回退本地 Excel。""" | |
| global _data_source | |
| config = load_config(config_path) | |
| df = _load_from_google_sheets(config) | |
| if df is not None: | |
| _data_source = "sheets" | |
| else: | |
| df = _load_from_excel(config) | |
| _data_source = "excel" | |
| return df | |
| def normalize_columns(df: pd.DataFrame, config: dict[str, Any]) -> pd.DataFrame: | |
| """用列索引将编码异常的列重命名为可读字段名。""" | |
| col_cfg = config.get("columns", {}) | |
| idx_map = { | |
| col_cfg.get("country_idx", 0): "country", | |
| col_cfg.get("platform_idx", 1): "platform", | |
| col_cfg.get("campaign_type_idx", 2): "campaign_type", | |
| col_cfg.get("campaign_name_idx", 3): "campaign_name", | |
| col_cfg.get("date_idx", 4): "date", | |
| } | |
| rename = {} | |
| for idx, name in idx_map.items(): | |
| if idx < len(df.columns): | |
| rename[df.columns[idx]] = name | |
| return df.rename(columns=rename) | |
| def _clean_numeric(s: pd.Series) -> pd.Series: | |
| """ | |
| 清洗 Google Sheets 格式化字符串。 | |
| 带 % 的值先标记,去掉符号后除以 100 还原为小数。 | |
| """ | |
| raw = s.astype(str) | |
| is_pct = raw.str.contains('%', na=False) | |
| cleaned = raw.str.replace(r'[$,%\s]', '', regex=True) | |
| values = pd.to_numeric(cleaned, errors="coerce") | |
| values[is_pct] = values[is_pct] / 100.0 | |
| return values | |
| def load_daily_data(config_path: Path | None = None) -> pd.DataFrame: | |
| """ | |
| 加载并清洗数据,只保留日度明细行。 | |
| 自动从 date 列派生 year、week(ISO week)、month。 | |
| 兼容本地 Excel(数值)和 Google Sheets(格式化字符串)。 | |
| """ | |
| config = load_config(config_path) | |
| df = load_raw(config_path) | |
| df = normalize_columns(df, config) | |
| # Excel 层级结构中,子行的维度列为空字符串,需要前向填充 | |
| # API 写入的 Sheets 分平台表每行都有完整字段,不需要 ffill | |
| if _data_source == "excel": | |
| for dim in ["country", "platform", "campaign_type", "campaign_name"]: | |
| if dim in df.columns: | |
| df[dim] = df[dim].replace("", pd.NA).ffill() | |
| df["date"] = pd.to_datetime(df["date"], format="mixed", errors="coerce") | |
| df = df[df["date"].notna()].copy() | |
| df = df[df["country"].notna() & (~df["country"].astype(str).str.contains("Total|Applied filters", na=False))].copy() | |
| df = df[df["platform"].notna() & (df["platform"].astype(str).str.strip() != "Total")].copy() | |
| df = df[df["campaign_type"].notna() & (df["campaign_type"].astype(str).str.strip() != "Total")].copy() | |
| df = df[df["campaign_name"].notna() & (df["campaign_name"].astype(str).str.strip() != "Total")].copy() | |
| df["year"] = df["date"].dt.isocalendar().year.astype(int) | |
| df["week"] = df["date"].dt.isocalendar().week.astype(int) | |
| df["month"] = df["date"].dt.month | |
| metrics = config.get("metrics", []) | |
| for col in metrics: | |
| if col in df.columns: | |
| df[col] = _clean_numeric(df[col]) | |
| df = df.reset_index(drop=True) | |
| df = _fix_campaign_type_from_name(df) | |
| dedup_cols = ["date", "platform", "campaign_name", "country"] | |
| if all(c in df.columns for c in dedup_cols): | |
| before = len(df) | |
| df = df.drop_duplicates(subset=dedup_cols, keep="last") | |
| dropped = before - len(df) | |
| if dropped: | |
| log.info(f"去重删除 {dropped} 行重复数据") | |
| df = df.reset_index(drop=True) | |
| return df | |
| _VALID_FUNNELS = {"Awareness", "Consideration", "Conversion"} | |
| def _fix_campaign_type_from_name(df: pd.DataFrame) -> pd.DataFrame: | |
| """根据 campaign_name 中的第二段修正 campaign_type(如名称含明确漏斗标识)。""" | |
| if "campaign_name" not in df.columns or "campaign_type" not in df.columns: | |
| return df | |
| parts = df["campaign_name"].astype(str).str.split("_", n=2) | |
| name_funnel = parts.str[1] | |
| mask = name_funnel.isin(_VALID_FUNNELS) & (name_funnel != df["campaign_type"]) | |
| if mask.any(): | |
| log.info(f"根据 campaign_name 修正 {mask.sum()} 行 campaign_type") | |
| df.loc[mask, "campaign_type"] = name_funnel[mask] | |
| return df | |
| # 兼容旧接口 | |
| load_weekly_data = load_daily_data | |
| def get_available_weeks(df: pd.DataFrame) -> list[tuple[int, int]]: | |
| """返回可用的 (year, week) 列表,按时间排序。""" | |
| pairs = df[["year", "week"]].drop_duplicates().sort_values(["year", "week"]) | |
| return list(pairs.itertuples(index=False, name=None)) | |
| def get_latest_week(df: pd.DataFrame) -> tuple[int, int]: | |
| """返回数据中最新一周的 (year, week)。""" | |
| weeks = get_available_weeks(df) | |
| if not weeks: | |
| raise ValueError("数据中无可用周数据") | |
| return weeks[-1] | |
| def get_available_months(df: pd.DataFrame) -> list[tuple[int, int]]: | |
| """返回可用的 (year, month) 列表,按时间排序。""" | |
| cal_year = df["date"].dt.year | |
| pairs = pd.DataFrame({"year": cal_year, "month": df["month"]}).drop_duplicates().sort_values(["year", "month"]) | |
| return list(pairs.itertuples(index=False, name=None)) | |