from __future__ import annotations from pathlib import Path import pandas as pd from langchain_core.tools import tool from tools.config import PARAMETER_FILE_MAP, PARAMETERS_DIR _FILENAME_TO_KEY: dict[str, str] = { "cluster_mapping.csv": "cluster_mapping", "market_summary_mapping.csv": "market_summary_mapping", "threshold_matrix.csv": "threshold_matrix", "country_region_mapping.csv": "country_region_mapping", "company_products.csv": "company_products", } def _resolve_parameter_keys(user_query: str) -> list[str]: query = (user_query or "").lower() selected: set[str] = set() for keyword, filename in PARAMETER_FILE_MAP.items(): if keyword in query: key = _FILENAME_TO_KEY.get(filename) if key: selected.add(key) if not selected: selected.update(["cluster_mapping", "country_region_mapping"]) return sorted(selected) def _load_param_df(key: str) -> pd.DataFrame: """Load a parameter table from the SQLite DB; fall back to CSV if needed.""" try: from db import load_param_from_db return load_param_from_db(key) except Exception: pass csv_name = { "cluster_mapping": "cluster_mapping.csv", "market_summary_mapping": "market_summary_mapping.csv", "threshold_matrix": "threshold_matrix.csv", "country_region_mapping": "country_region_mapping.csv", "company_products": "company_products.csv", }.get(key, "") if csv_name: path = Path(PARAMETERS_DIR) / csv_name if path.exists(): return pd.read_csv(path, keep_default_na=False) return pd.DataFrame() def _row_score(row_values: list[str], query_tokens: set[str]) -> int: joined = " ".join(row_values).lower() return sum(1 for token in query_tokens if token and token in joined) @tool def parameter_reader(user_query: str, max_rows_per_file: int = 25) -> list[dict]: """Read relevant parameter rows from the parameter tables (DB or CSV fallback).""" rows: list[dict] = [] keys = _resolve_parameter_keys(user_query) query_tokens = {token for token in (user_query or "").lower().split() if len(token) > 2} for key in keys: try: df = _load_param_df(key) except Exception: continue if df.empty: continue records = df.to_dict(orient="records") if query_tokens: scored = [] for record in records: score = _row_score([str(v) for v in record.values()], query_tokens) if score > 0: scored.append((score, record)) scored.sort(key=lambda item: item[0], reverse=True) selected_records = [record for _, record in scored[:max_rows_per_file]] else: selected_records = records[:max_rows_per_file] source_label = f"{key} (db)" for record in selected_records: rows.append({"_source": source_label, **record}) return rows