""" Query interpreter component for causal inference. This module provides functionality to match query concepts to actual dataset variables, identifying treatment, outcome, and covariate variables for causal inference analysis. """ import re from typing import Dict, List, Any, Optional, Union, Tuple import pandas as pd import logging import numpy as np from auto_causal.config import get_llm_client # Import LLM and message types from langchain_core.language_models import BaseChatModel from langchain_core.messages import HumanMessage from langchain_core.exceptions import OutputParserException # Import base Pydantic models needed directly from pydantic import BaseModel, ValidationError from dowhy import CausalModel import json # Import shared Pydantic models from the central location from auto_causal.models import ( LLMSelectedVariable, LLMSelectedCovariates, LLMIVars, LLMRDDVars, LLMRCTCheck, LLMTreatmentReferenceLevel, LLMInteractionSuggestion, LLMEstimand, # LLMDIDCheck, # LLMDiDTemporalVars, # LLMDiDGroupVars, # LLMRDDCheck, # LLMRDDVarsExtended ) # Import the new prompt templates from auto_causal.prompts.method_identification_prompts import ( IV_IDENTIFICATION_PROMPT_TEMPLATE, RDD_IDENTIFICATION_PROMPT_TEMPLATE, RCT_IDENTIFICATION_PROMPT_TEMPLATE, TREATMENT_REFERENCE_IDENTIFICATION_PROMPT_TEMPLATE, INTERACTION_TERM_IDENTIFICATION_PROMPT_TEMPLATE, TREATMENT_VAR_IDENTIFICATION_PROMPT_TEMPLATE, OUTCOME_VAR_IDENTIFICATION_PROMPT_TEMPLATE, COVARIATES_IDENTIFICATION_PROMPT_TEMPLATE, ESTIMAND_PROMPT_TEMPLATE, CONFOUNDER_IDENTIFICATION_PROMPT_TEMPLATE, DID_TERM_IDENTIFICATION_PROMPT_TEMPLATE) # Assume central models are defined elsewhere or keep local definitions for now # from ..models import ... # --- Pydantic models for LLM structured output --- # REMOVED - Now defined in causalscientist/auto_causal/models.py # class LLMSelectedVariable(BaseModel): ... # class LLMSelectedCovariates(BaseModel): ... # class LLMIVars(BaseModel): ... # class LLMRDDVars(BaseModel): ... # class LLMRCTCheck(BaseModel): ... logger = logging.getLogger(__name__) def infer_treatment_variable_type(treatment_variable: str, column_categories: Dict[str, str], dataset_analysis: Dict[str, Any]) -> str: """ Determine treatment variable type from column category and unique value count Args: treatment_variable: name of the treatment variable column_categories: mapping of column names to their categories dataset_analysis: exploratory analysis results Returns: str: type of the treatment variable (e.g., "binary", "continuous", etc """ treatment_variable_type = "unknown" if treatment_variable and treatment_variable in column_categories: category = column_categories[treatment_variable] logger.info(f"Category for treatment '{treatment_variable}' is '{category}'.") if category == "continuous_numeric": treatment_variable_type = "continuous" elif category == "discrete_numeric": num_unique = dataset_analysis.get("column_nunique_counts", {}).get(treatment_variable, -1) if num_unique > 10: logger.info(f"'{treatment_variable}' has {num_unique} unique values, treating as continuous.") treatment_variable_type = "continuous" elif num_unique == 2: logger.info(f"'{treatment_variable}' has 2 unique values, treating as binary.") treatment_variable_type = "binary" elif num_unique > 0: logger.info(f"'{treatment_variable}' has {num_unique} unique values, treating as discrete_multi_value.") treatment_variable_type = "discrete_multi_value" else: logger.info(f"'{treatment_variable}' unique value count unknown or too few.") treatment_variable_type = "discrete_numeric_unknown_cardinality" elif category in ["binary", "binary_categorical"]: treatment_variable_type = "binary" elif category in ["categorical", "categorical_numeric"]: num_unique = dataset_analysis.get("column_nunique_counts", {}).get(treatment_variable, -1) if num_unique == 2: treatment_variable_type = "binary" elif num_unique > 0: treatment_variable_type = "categorical_multi_value" else: treatment_variable_type = "categorical_unknown_cardinality" else: logger.warning(f"Unmapped category '{category}' for '{treatment_variable}', setting as 'other'.") treatment_variable_type = "other" elif treatment_variable: logger.warning(f"'{treatment_variable}' not found in column_categories.") else: logger.info("No treatment variable identified.") logger.info(f"Final Determined Treatment Variable Type: {treatment_variable_type}") return treatment_variable_type def determine_treatment_reference_level(is_rct: Optional[bool], llm: Optional[BaseChatModel], treatment_variable: Optional[str], query_text: str, dataset_description: Optional[str], file_path: Optional[str], columns: List[str]) -> Optional[str]: """ Determines the treatment reference level """ # If LLM didn't explicitly say RCT, default to False or keep None? # Let's default to False if LLM didn't provide a boolean value. if is_rct is None: is_rct = False treatment_reference_level = None if llm and treatment_variable and treatment_variable in columns: treatment_values_sample = [] if file_path: try: df = pd.read_csv(file_path) if treatment_variable in df.columns: unique_vals = df[treatment_variable].unique() treatment_values_sample = [item.item() if hasattr(item, 'item') else item for item in unique_vals][:10] if treatment_values_sample: logger.info(f"Successfully read treatment values sample from dataset at '{file_path}' for variable '{treatment_variable}'.") else: logger.info(f"'{treatment_variable}' in '{file_path}' has no unique values or is empty.") else: logger.warning(f"'{treatment_variable}' not found in dataset columns at '{file_path}'.") except FileNotFoundError: logger.warning(f"File not found at: {file_path}") except pd.errors.EmptyDataError: logger.warning(f"Empty file at: {file_path}") except Exception as e: logger.warning(f"Error reading dataset at '{file_path}' for '{treatment_variable}': {e}") if not treatment_values_sample: logger.warning(f"No unique values found for treatment '{treatment_variable}'. LLM prompt will receive empty list.") else: logger.info(f"Final treatment values sample: {treatment_values_sample}") try: prompt = TREATMENT_REFERENCE_IDENTIFICATION_PROMPT_TEMPLATE.format(query=query_text, description=dataset_description or 'N/A', treatment_variable=treatment_variable, treatment_variable_values=treatment_values_sample) ref_result = _call_llm_for_var(llm, prompt, LLMTreatmentReferenceLevel) if ref_result and ref_result.reference_level: if treatment_values_sample and ref_result.reference_level not in treatment_values_sample: logger.warning(f"LLM reference level '{ref_result.reference_level}' not in sampled values for '{treatment_variable}'.") treatment_reference_level = ref_result.reference_level logger.info(f"LLM identified reference level: {treatment_reference_level} (Reason: {ref_result.reasoning})") elif ref_result: logger.info(f"LLM returned no reference level (Reason: {ref_result.reasoning})") except Exception as e: logger.error(f"LLM error for treatment reference level: {e}") return treatment_reference_level def identify_interaction_term(llm: Optional[BaseChatModel], treatment_variable: Optional[str], covariates: List[str], column_categories: Dict[str, str], query_text: str, dataset_description: Optional[str]) -> Tuple[bool, Optional[str]]: """ Identifies the interaction term based on the query and the dataset information """ interaction_term_suggested, interaction_variable_candidate = False, None if llm and treatment_variable and covariates: try: covariates_list_str = "\n".join([f"- {cov}: {column_categories.get(cov, 'Unknown')}" for cov in covariates]) or "No covariates identified or available." prompt = INTERACTION_TERM_IDENTIFICATION_PROMPT_TEMPLATE.format(query=query_text, description=dataset_description or 'N/A', treatment_variable=treatment_variable, covariates_list_with_types=covariates_list_str) result = _call_llm_for_var(llm, prompt, LLMInteractionSuggestion) if result: interaction_term_suggested = result.interaction_needed if result.interaction_needed is not None else False if interaction_term_suggested and result.interaction_variable: if result.interaction_variable in covariates: interaction_variable_candidate = result.interaction_variable logger.info(f"LLM suggested interaction: needed={interaction_term_suggested}, variable='{interaction_variable_candidate}' (Reason: {result.reasoning})") else: logger.warning(f"LLM suggested variable '{result.interaction_variable}' not in covariates {covariates}. Ignoring.") interaction_term_suggested = False elif interaction_term_suggested: logger.info(f"LLM suggested interaction is needed but no variable provided (Reason: {result.reasoning})") else: logger.info(f"LLM suggested no interaction is needed (Reason: {result.reasoning})") else: logger.warning("LLM returned no result for interaction term suggestion.") except Exception as e: logger.error(f"LLM error during interaction term check: {e}") return interaction_term_suggested, interaction_variable_candidate def interpret_query(query_info: Dict[str, Any], dataset_analysis: Dict[str, Any], dataset_description: Optional[str] = None) -> Dict[str, Any]: """ Interpret query using hybrid heuristic/LLM approach to identify variables. Args: query_info: Information extracted from the user's query (text, hints). dataset_analysis: Information about the dataset structure (columns, types, etc.). dataset_description: Optional textual description of the dataset. llm: Optional language model instance. Returns: Dict containing identified variables (treatment, outcome, covariates, etc., and is_rct). """ logger.info("Interpreting query with hybrid approach...") llm = get_llm_client() query_text = query_info.get("query_text", "") columns = dataset_analysis.get("columns", []) column_categories = dataset_analysis.get("column_categories", {}) file_path = dataset_analysis["dataset_info"]["file_path"] # --- Identify Treatment --- treatment_hints = query_info.get("potential_treatments", []) dataset_treatments = dataset_analysis.get("potential_treatments", []) treatment_variable = _identify_variable_hybrid(role="treatment", query_hints=treatment_hints, dataset_suggestions=dataset_treatments, columns=columns, column_categories=column_categories, prioritize_types=["binary", "binary_categorical", "discrete_numeric","continuous_numeric"], # Prioritize binary/discrete query_text=query_text, dataset_description=dataset_description,llm=llm) logger.info(f"Identified Treatment: {treatment_variable}") treatment_variable_type = infer_treatment_variable_type(treatment_variable, column_categories, dataset_analysis) # --- Identify Outcome --- outcome_hints = query_info.get("outcome_hints", []) dataset_outcomes = dataset_analysis.get("potential_outcomes", []) outcome_variable = _identify_variable_hybrid(role="outcome", query_hints=outcome_hints, dataset_suggestions=dataset_outcomes, columns=columns, column_categories=column_categories, prioritize_types=["continuous_numeric", "discrete_numeric"], # Prioritize numeric exclude_vars=[treatment_variable], # Exclude treatment query_text=query_text, dataset_description=dataset_description, llm=llm) logger.info(f"Identified Outcome: {outcome_variable}") # --- Identify Covariates --- covariate_hints = query_info.get("covariates_hints", []) covariates = _identify_covariates_hybrid("covars", treatment_variable=treatment_variable, outcome_variable=outcome_variable, columns=columns, column_categories=column_categories, query_hints=covariate_hints, query_text=query_text, dataset_description=dataset_description, llm=llm) logger.info(f"Identified Covariates: {covariates}") # --- Identify Confounders --- confounder_hints = query_info.get("covariates_hints", []) confounders = _identify_covariates_hybrid("confounders", treatment_variable=treatment_variable, outcome_variable=outcome_variable, columns=columns, column_categories=column_categories, query_hints=confounder_hints, query_text=query_text, dataset_description=dataset_description, llm=llm) logger.info(f"Identified Confounders: {confounders}") # --- Identify Time/Group (from dataset analysis) --- time_variable = None group_variable = None has_temporal = dataset_analysis.get("temporal_structure", {}).get("has_temporal_structure", False) temporal_structure = dataset_analysis.get("temporal_structure", {}) if temporal_structure.get("has_temporal_structure", False): time_variable = temporal_structure.get("time_column") or temporal_structure.get("temporal_columns", [None])[0] if temporal_structure.get("is_panel_data", False): group_variable = temporal_structure.get("id_column") logger.info(f"Identified Time Var: {time_variable}, Group Var: {group_variable}, temporal structure: {temporal_structure}") # --- Identify IV/RDD/RCT using LLM --- instrument_variable = None running_variable = None cutoff_value = None is_rct = None smd_score = None if llm: try: # Check for RCT prompt_rct = _create_identify_prompt("whether data is from RCT", query_text, dataset_description, columns, column_categories, treatment_variable, outcome_variable) rct_result = _call_llm_for_var(llm, prompt_rct, LLMRCTCheck) is_rct = rct_result.is_rct if rct_result else None logger.info(f"LLM identified RCT: {is_rct}") # Check for IV prompt_iv = _create_identify_prompt("instrumental variable", query_text, dataset_description, columns, column_categories, treatment_variable, outcome_variable) iv_result = _call_llm_for_var(llm, prompt_iv, LLMIVars) instrument_variable = iv_result.instrument_variable if iv_result else None if instrument_variable not in columns: instrument_variable = None logger.info(f"LLM identified IV: {instrument_variable}") # Check for RDD prompt_rdd = _create_identify_prompt("regression discontinuity (running variable and cutoff)", query_text, dataset_description, columns, column_categories, treatment_variable, outcome_variable) rdd_result = _call_llm_for_var(llm, prompt_rdd, LLMRDDVars) if rdd_result: running_variable = rdd_result.running_variable cutoff_value = rdd_result.cutoff_value if running_variable not in columns or cutoff_value is None: running_variable = None cutoff_value = None logger.info(f"LLM identified RDD: Running={running_variable}, Cutoff={cutoff_value}") ## For graph based methods exclude_cols = [treatment_variable, outcome_variable] potential_covariates = [col for col in columns if col not in exclude_cols and col is not None] usable_covariates = [col for col in potential_covariates if column_categories.get(col) not in ["text_or_other"]] logger.info(f"Usable covariates for graph: {usable_covariates}") estimand_prompt = ESTIMAND_PROMPT_TEMPLATE.format(query=query_text,dataset_description=dataset_description, dataset_columns=usable_covariates, treatment=treatment_variable, outcome=outcome_variable) estimand_result = _call_llm_for_var(llm, estimand_prompt, LLMEstimand) estimand = "ate" if "ate" in estimand_result.estimand.strip().lower() else "att" logger.info(f"LLM identified estimand: {estimand}") ## Did Term did_term_prompt = DID_TERM_IDENTIFICATION_PROMPT_TEMPLATE.format(query=query_text, description=dataset_description, column_info=columns, time_variable=time_variable, group_variable=group_variable, column_types=column_categories) did_term_result = _call_llm_for_var(llm, did_term_prompt, LLMRDDVars) did_term_result = did_term_result.did_term if did_term_result in columns else None logger.info(f"LLM identified DiD term: {did_term_result}") #smd_score_all = compute_smd(dataset_analysis.get("data", pd.DataFrame()), treatment_variable, usable_covariates) #smd_score = smd_score_all.get("ate", 0.0) if smd_score_all else 0.0 #logger.info(f"Computed SMD score: {smd_score}") #logger.debug(f"Computed SMD score for {estimand}: {smd_score}") except Exception as e: logger.error(f"Error during LLM checks for IV/RDD/RCT: {e}") # --- Identify Treatment Reference Level --- treatment_reference_level = determine_treatment_reference_level(is_rct=is_rct, llm=llm, treatment_variable=treatment_variable, query_text=query_text, dataset_description=dataset_description, file_path=file_path, columns=columns) # --- Identify Interaction Term Suggestion --- interaction_term_suggested, interaction_variable_candidate = identify_interaction_term(llm=llm, treatment_variable=treatment_variable, covariates=covariates, column_categories=column_categories, query_text=query_text, dataset_description=dataset_description) # --- Consolidate --- return { "treatment_variable": treatment_variable, "treatment_variable_type": treatment_variable_type, "outcome_variable": outcome_variable, "covariates": covariates, "time_variable": time_variable, "group_variable": group_variable, "instrument_variable": instrument_variable, "running_variable": running_variable, "cutoff_value": cutoff_value, "is_rct": is_rct, "treatment_reference_level": treatment_reference_level, "interaction_term_suggested": interaction_term_suggested, "interaction_variable_candidate": interaction_variable_candidate, "confounders": confounders, "did_term": did_term_result } def compute_smd(df: pd.DataFrame, treat, covars_list) -> Dict[str, float]: """ Computed the standardized mean differences (SMD) for the treatment variable Args: df (pd.DataFrame): The dataset. treat (str): Name of the binary treatment column (0/1). covars_list (List[str]): List of covariate names to consider for SMD calculation Returns: Dict{str ->float}: the standardized mean difference (SMD) """ logger.info(f"Computing SMD for treatment variable '{treat}' with covariates: {covars_list}") df_t = df[df[treat] == 1] df_c = df[df[treat] == 0] covariates = covars_list if covars_list else df.columns.tolist() smd_ate = np.zeros(len(covariates)) smd_att = np.zeros(len(covariates)) for i, col in enumerate(covariates): try: m_t, m_c = df_t[col].mean(), df_c[col].mean() s_t, s_c = df_t[col].std(ddof=0), df_c[col].std(ddof=0) pooled = np.sqrt((s_t**2 + s_c**2) / 2) ate_val = 0.0 if pooled == 0 else (m_t - m_c) / pooled att_val = 0.0 if s_t == 0 else (m_t - m_c) / s_t smd_ate.append(ate_val) smd_att.append(att_val) except Exception as e: logger.warning(f"SMD computation failed for column '{col}': {e}") continue avg_ate = np.nanmean(np.abs(smd_ate)) avg_att = np.nanmean(np.abs(smd_att)) return {"ate":avg_ate, "att":avg_att} # --- Helper Functions for Hybrid Identification --- def _identify_variable_hybrid(role: str, query_hints: List[str], dataset_suggestions: List[str], columns: List[str], column_categories: Dict[str, str], prioritize_types: List[str], query_text: str, dataset_description: Optional[str],llm: Optional[BaseChatModel], exclude_vars: Optional[List[str]] = None) -> Optional[str]: """ Used to identify a variable from the avaiable information by prompting the LLM. In case of failure, it will fallback to a programmatic selection (heuristics) Args: role: variable type (treatment or outcome) query_hints: hints from the query for this variable dataset_suggestions: dataset-specific suggestions for this variable columns: list of available columns in the dataset column_categories: mapping of column names to their categories prioritize_types: types to prioritize for this variable query_text: the original query text dataset_description: description of the dataset llm: language model exclude_vars: list of variables to exclude from selection (e.g., treatment for outcome) Returns: str: name of the identified variable, or None if not found """ candidates = set() available_columns = [c for c in columns if c not in (exclude_vars or [])] if not available_columns: return None # 1. Exact matches from hints for hint in query_hints: if hint in available_columns: candidates.add(hint) # 2. Add dataset suggestions for sugg in dataset_suggestions: if sugg in available_columns: candidates.add(sugg) # 3. Programmatic Filtering based on type plausible_candidates = [c for c in candidates if column_categories.get(c) in prioritize_types] if llm: if role == "treatment": prompt_template = TREATMENT_VAR_IDENTIFICATION_PROMPT_TEMPLATE elif role == "outcome": prompt_template = OUTCOME_VAR_IDENTIFICATION_PROMPT_TEMPLATE else: raise ValueError(f"Unsupported role for LLM variable identification: {role}") prompt = prompt_template.format(query=query_text, description=dataset_description, column_info=available_columns) llm_choice = _call_llm_for_var(llm, prompt, LLMSelectedVariable) if llm_choice and llm_choice.variable_name in available_columns: logger.info(f"LLM selected {role}: {llm_choice.variable_name}") return llm_choice.variable_name else: fallback = plausible_candidates[0] if plausible_candidates else None logger.warning(f"LLM failed to select valid {role}. Falling back to: {fallback}") return fallback if plausible_candidates: logger.info(f"No LLM provided. Using first plausible {role}: {plausible_candidates[0]}") return plausible_candidates[0] logger.warning(f"No plausible candidates for {role}. Cannot identify variable.") return None def _identify_covariates_hybrid(role, treatment_variable: Optional[str], outcome_variable: Optional[str], columns: List[str], column_categories: Dict[str, str], query_hints: List[str], query_text: str, dataset_description: Optional[str], llm: Optional[BaseChatModel]) -> List[str]: """ Prompts an LLM to identify the covariates """ # 1. Initial Programmatic Filtering exclude_cols = [treatment_variable, outcome_variable] potential_covariates = [col for col in columns if col not in exclude_cols and col is not None] # Filter out unusable types usable_covariates = [col for col in potential_covariates if column_categories.get(col) not in ["text_or_other"]] logger.debug(f"Initial usable covariates: {usable_covariates}") # 2. LLM Refinement (if LLM available) if llm: logger.info("Using LLM to refine covariate list...") prompt = "" if role == "covars": prompt = COVARIATES_IDENTIFICATION_PROMPT_TEMPLATE.format("covars", query=query_text, description=dataset_description, column_info=", ".join(usable_covariates), treatment=treatment_variable, outcome=outcome_variable) elif role == "confounders": prompt = CONFOUNDER_IDENTIFICATION_PROMPT_TEMPLATE.format(query=query_text, description=dataset_description, column_info=", ".join(usable_covariates), treatment=treatment_variable, outcome=outcome_variable) llm_selection = _call_llm_for_var(llm, prompt, LLMSelectedCovariates) if llm_selection and llm_selection.covariates: # Validate LLM output against available columns valid_llm_covs = [c for c in llm_selection.covariates if c in usable_covariates] if len(valid_llm_covs) < len(llm_selection.covariates): logger.warning("LLM suggested covariates not found in initial usable list.") if valid_llm_covs: # Use LLM selection if it's valid and non-empty logger.info(f"LLM refined covariates to: {valid_llm_covs}") return valid_llm_covs[:10] # Cap at 10 else: logger.warning("LLM refinement failed or returned empty/invalid list. Falling back.") else: logger.warning("LLM refinement call failed or returned no covariates. Falling back.") # 3. Fallback to Programmatic List (Capped) logger.info(f"Using programmatically determined covariates (capped at 10): {usable_covariates[:10]}") return usable_covariates[:10] def _create_identify_prompt(target: str, query: str, description: Optional[str], columns: List[str], categories: Dict[str,str], treatment: Optional[str], outcome: Optional[str]) -> str: """ Creates a prompt to ask LLM to identify specific roles like IV, RDD, or RCT by selecting and formatting a specific template """ column_info = "\n".join([f"- '{c}' (Type: {categories.get(c, 'Unknown')})" for c in columns]) # Select the appropriate detailed prompt template based on the target if "instrumental variable" in target.lower(): template = IV_IDENTIFICATION_PROMPT_TEMPLATE elif "regression discontinuity" in target.lower(): template = RDD_IDENTIFICATION_PROMPT_TEMPLATE elif "rct" in target.lower(): template = RCT_IDENTIFICATION_PROMPT_TEMPLATE else: # Fallback or error? For now, let's raise an error if target is unexpected. logger.error(f"Unsupported target for _create_identify_prompt: {target}") raise ValueError(f"Unsupported target for specific identification prompt: {target}") # Format the selected template with the provided context prompt = template.format(query=query, description=description or 'N/A', column_info=column_info, treatment=treatment or 'N/A', outcome=outcome or 'N/A') return prompt def _call_llm_for_var(llm: BaseChatModel, prompt: str, pydantic_model: BaseModel) -> Optional[BaseModel]: """Helper to call LLM with structured output and handle errors.""" try: messages = [HumanMessage(content=prompt)] structured_llm = llm.with_structured_output(pydantic_model) parsed_result = structured_llm.invoke(messages) return parsed_result except (OutputParserException, ValidationError) as e: logger.error(f"LLM call failed parsing/validation for {pydantic_model.__name__}: {e}") except Exception as e: logger.error(f"LLM call failed unexpectedly for {pydantic_model.__name__}: {e}", exc_info=True) return None