""" SQLite-backed database layer for the Hugging Face Space demo. Mirrors the original Aurora PostgreSQL API surface (`get_metrics_distinct_values`, `get_metrics_from_db_filtered`, `load_param_from_db`) so the agents and tools can be re-used unchanged, but talks to a single-file `demo.sqlite` instead of a cloud DB. """ from __future__ import annotations import os import re import sqlite3 from pathlib import Path from typing import Iterable import pandas as pd HERE = Path(__file__).resolve().parent DB_PATH = Path(os.environ.get("DEMO_DB_PATH", HERE / "demo.sqlite")) # Column order matches the original sentinel_metrics CSV/Postgres schema. METRICS_TABLE_COLUMNS = [ "Period", "Region", "Level", "Product", "Country", "Cluster", "TA Market", "Class", "Current_Value", "Current_Value_MS", "Value_MS_Change", "Current_Value_Rank", "Value_Rank_Change", "Value_Growth", "Value_Trend_Reversal", "Current_Volume", "Current_Volume_MS", "Volume_MS_Change", "Current_Volume_Rank", "Volume_Rank_Change", "Volume_Growth", "Volume_Trend_Reversal", "ASP_Growth", "Calculation_Type", ] # SQLite uses lowercase identifiers (no quoting headache) and we rename back when # returning DataFrames to the agents. _DF_TO_DB_COL: dict[str, str] = { "Period": "period", "Region": "region", "Level": "level", "Product": "product", "Country": "country", "Cluster": "cluster", "TA Market": "ta_market", "Class": "class", "Current_Value": "current_value", "Current_Value_MS": "current_value_ms", "Value_MS_Change": "value_ms_change", "Current_Value_Rank": "current_value_rank", "Value_Rank_Change": "value_rank_change", "Value_Growth": "value_growth", "Value_Trend_Reversal": "value_trend_reversal", "Current_Volume": "current_volume", "Current_Volume_MS": "current_volume_ms", "Volume_MS_Change": "volume_ms_change", "Current_Volume_Rank": "current_volume_rank", "Volume_Rank_Change": "volume_rank_change", "Volume_Growth": "volume_growth", "Volume_Trend_Reversal": "volume_trend_reversal", "ASP_Growth": "asp_growth", "Calculation_Type": "calculation_type", } _DB_TO_DF_COL: dict[str, str] = {v: k for k, v in _DF_TO_DB_COL.items()} _METRICS_DDL = """ CREATE TABLE IF NOT EXISTS sentinel_metrics ( period TEXT, region TEXT, level TEXT, product TEXT, country TEXT, cluster TEXT, ta_market TEXT, class TEXT, current_value REAL, current_value_ms REAL, value_ms_change REAL, current_value_rank INTEGER, value_rank_change INTEGER, value_growth REAL, value_trend_reversal REAL, current_volume REAL, current_volume_ms REAL, volume_ms_change REAL, current_volume_rank INTEGER, volume_rank_change INTEGER, volume_growth REAL, volume_trend_reversal REAL, asp_growth REAL, calculation_type TEXT ); """ _METRICS_INDEXES = ( "CREATE INDEX IF NOT EXISTS idx_metrics_region ON sentinel_metrics(region);", "CREATE INDEX IF NOT EXISTS idx_metrics_period ON sentinel_metrics(period);", "CREATE INDEX IF NOT EXISTS idx_metrics_cluster ON sentinel_metrics(cluster);", "CREATE INDEX IF NOT EXISTS idx_metrics_product ON sentinel_metrics(product);", "CREATE INDEX IF NOT EXISTS idx_metrics_calc ON sentinel_metrics(calculation_type);", "CREATE INDEX IF NOT EXISTS idx_metrics_class ON sentinel_metrics(class);", ) # Parameter tables — same key names as the original code so callers don't need to change. PARAM_TABLE_MAP: dict[str, str] = { "cluster_mapping": "sentinel_cluster_mapping", "market_summary_mapping": "sentinel_market_summary_mapping", "threshold_matrix": "sentinel_threshold_matrix", "country_region_mapping": "sentinel_country_region_mapping", "merck_products": "sentinel_company_products", # legacy key kept for parameter_tools compatibility "company_products": "sentinel_company_products", } _PARAM_DDL: dict[str, str] = { "sentinel_cluster_mapping": """ CREATE TABLE IF NOT EXISTS sentinel_cluster_mapping ( country TEXT, market TEXT, cluster TEXT, region TEXT ); """, "sentinel_market_summary_mapping": """ CREATE TABLE IF NOT EXISTS sentinel_market_summary_mapping ( region TEXT, market TEXT, value_summary TEXT, volume_summary TEXT, timeframe_summary TEXT ); """, "sentinel_threshold_matrix": """ CREATE TABLE IF NOT EXISTS sentinel_threshold_matrix ( cluster TEXT, market TEXT, value_ms_change REAL, value_growth REAL, volume_ms_change REAL, volume_growth REAL, asp_growth REAL, region TEXT ); """, "sentinel_country_region_mapping": """ CREATE TABLE IF NOT EXISTS sentinel_country_region_mapping ( country_code TEXT, country_name TEXT, region TEXT ); """, "sentinel_company_products": """ CREATE TABLE IF NOT EXISTS sentinel_company_products ( product TEXT ); """, } def _connect() -> sqlite3.Connection: """Open a SQLite connection with row factory and foreign keys enabled.""" DB_PATH.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(DB_PATH)) conn.execute("PRAGMA journal_mode = WAL;") return conn def create_all_tables() -> None: """Create the metrics table, indexes, and all parameter tables if missing.""" with _connect() as conn: conn.execute(_METRICS_DDL) for stmt in _METRICS_INDEXES: conn.execute(stmt) for ddl in _PARAM_DDL.values(): conn.execute(ddl) conn.commit() def db_is_seeded() -> bool: """Return True when the metrics table has at least one row.""" if not DB_PATH.exists(): return False try: with _connect() as conn: row = conn.execute("SELECT COUNT(1) FROM sentinel_metrics;").fetchone() return bool(row and row[0] > 0) except sqlite3.OperationalError: return False # --------------------------------------------------------------------------- # Metrics queries # --------------------------------------------------------------------------- def _rename_db_to_df(df: pd.DataFrame) -> pd.DataFrame: """Rename SQLite (lowercase) columns back to the canonical DataFrame names.""" return df.rename(columns=_DB_TO_DF_COL) def _period_sort_key(period: str) -> tuple[int, int]: match = re.match(r"^\s*(\d{2})\s*Q([1-4])\s*$", str(period or ""), re.IGNORECASE) return (int(match.group(1)), int(match.group(2))) if match else (-1, -1) def get_metrics_distinct_values() -> dict: """ Return distinct values per filterable column in sentinel_metrics. Also returns ``_region_latest_period``: a dict mapping each region to its most recent period, used by the filter extractor to default smartly when the user doesn't specify a period. """ if not DB_PATH.exists(): return {} with _connect() as conn: cols_lower = ["region", "period", "cluster", "product", "calculation_type", "ta_market", "class", "level"] result: dict = {} for db_col in cols_lower: df_col = _DB_TO_DF_COL[db_col] rows = conn.execute( f"SELECT DISTINCT {db_col} FROM sentinel_metrics " f"WHERE {db_col} IS NOT NULL ORDER BY {db_col};" ).fetchall() result[df_col] = [str(r[0]) for r in rows if r[0] is not None] region_period_rows = conn.execute( "SELECT region, period FROM sentinel_metrics " "WHERE region IS NOT NULL AND period IS NOT NULL " "GROUP BY region, period;" ).fetchall() region_periods: dict[str, list[str]] = {} for region, period in region_period_rows: region_periods.setdefault(str(region), []).append(str(period)) region_latest: dict[str, str] = { r: sorted(ps, key=_period_sort_key)[-1] for r, ps in region_periods.items() if ps } result["_region_latest_period"] = region_latest return result _FILTER_COL_TO_DB: dict[str, str] = { "Region": "region", "Period": "period", "Cluster": "cluster", "Product": "product", "Calculation_Type": "calculation_type", "TA Market": "ta_market", "Class": "class", "Level": "level", } def get_metrics_from_db_filtered(filters: dict) -> pd.DataFrame: """ Return the sentinel_metrics rows that match *filters*. *filters* keys must be DataFrame column names (e.g. ``Region``, ``Period``, ``Cluster``, ``Product``, ``Calculation_Type``, ``TA Market``, ``Class``). Unrecognised keys are silently ignored. Empty *filters* returns all rows. """ where_clauses: list[str] = [] params: list = [] for col, filter_values in (filters or {}).items(): db_col = _FILTER_COL_TO_DB.get(col) if not db_col or not filter_values: continue clean = [str(v).strip() for v in filter_values if str(v).strip()] if col == "Period": clean = [v.replace(" ", "") for v in clean] if not clean: continue placeholders = ", ".join(["?"] * len(clean)) where_clauses.append(f"{db_col} IN ({placeholders})") params.extend(clean) where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else "" query = f"SELECT * FROM sentinel_metrics {where_sql};" with _connect() as conn: df = pd.read_sql_query(query, conn, params=params if params else None) df = _rename_db_to_df(df) # Synthesize Cluster from Country/Region when missing (kept for backward compat # with the original Postgres helper, even though our schema always populates it). if "Cluster" not in df.columns: if {"Level", "Country", "Region"}.issubset(df.columns): df["Cluster"] = df.apply( lambda r: r["Country"] if r["Level"] == "Country" else r["Region"], axis=1 ) elif "Country" in df.columns: df["Cluster"] = df["Country"].fillna("").astype(str) return df # --------------------------------------------------------------------------- # Parameter tables # --------------------------------------------------------------------------- def load_param_from_db(key: str) -> pd.DataFrame: """ Load a parameter table by its logical key. Keys: ``cluster_mapping``, ``market_summary_mapping``, ``threshold_matrix``, ``country_region_mapping``, ``merck_products`` (alias of ``company_products``). Returns a DataFrame matching the original CSV column casing. """ if key not in PARAM_TABLE_MAP: raise ValueError( f"Unknown parameter key: {key!r}. Valid keys: {sorted(set(PARAM_TABLE_MAP))}" ) table = PARAM_TABLE_MAP[key] with _connect() as conn: df = pd.read_sql_query(f"SELECT * FROM {table};", conn) if key == "threshold_matrix": df = df.rename(columns={ "cluster": "Cluster", "value_ms_change": "Value_MS_Change", "value_growth": "Value_Growth", "volume_ms_change": "Volume_MS_Change", "volume_growth": "Volume_Growth", "asp_growth": "ASP_Growth", }) return df # --------------------------------------------------------------------------- # Bulk insert helpers (used by seed_data.py) # --------------------------------------------------------------------------- def replace_param_table(key: str, df: pd.DataFrame) -> int: """Replace the contents of a parameter table with *df*. Returns rows inserted.""" if key not in PARAM_TABLE_MAP: raise ValueError(f"Unknown parameter key: {key!r}") table = PARAM_TABLE_MAP[key] df = df.copy() df.columns = [c.lower() for c in df.columns] with _connect() as conn: conn.execute(f"DELETE FROM {table};") df.to_sql(table, conn, if_exists="append", index=False) conn.commit() return len(df) def insert_metrics(df: pd.DataFrame) -> int: """Insert metrics rows (DataFrame uses mixed-case column names).""" if df is None or df.empty: return 0 db_df = df.rename(columns=_DF_TO_DB_COL) keep_cols = [c for c in _DF_TO_DB_COL.values() if c in db_df.columns] db_df = db_df[keep_cols] with _connect() as conn: db_df.to_sql("sentinel_metrics", conn, if_exists="append", index=False) conn.commit() return len(db_df) def replace_metrics(df: pd.DataFrame) -> int: """Truncate the metrics table and insert *df*. Returns rows inserted.""" with _connect() as conn: conn.execute("DELETE FROM sentinel_metrics;") conn.commit() return insert_metrics(df)