{filename or 'Uploaded data'}
", unsafe_allow_html=True, ) st.markdown( f""" """, unsafe_allow_html=True, ) def parse_question_answer(raw_answer: Any) -> Tuple[Optional[str], str]: """Extract question and answer segments from assistant output, if present.""" if raw_answer is None: return None, "" text = str(raw_answer).strip() lower_text = text.lower() question_marker = "question:" answer_marker = "answer:" question: Optional[str] = None answer = text if answer_marker in lower_text: answer_idx = lower_text.find(answer_marker) answer_start = answer_idx + len(answer_marker) [truncated — 22437 more characters] ``` ### data_assistant/cli.py ```python #!/usr/bin/env python3 """Command-line interface for the AI Data Assistant.""" import argparse import json import os import tempfile from contextlib import suppress from typing import Dict, List, Optional, Tuple import pandas as pd from llm_integration import DataAssistant from profiler import profile_csv TRUE_VALUES = {"true", "t", "yes", "y", "1"} FALSE_VALUES = {"false", "f", "no", "n", "0"} BOOLEAN_COLUMNS = { "schoolsup", "famsup", "paid", "activities", "nursery", "higher", "internet", "romantic", } NUMERIC_COLUMNS = [ "age", "Medu", "Fedu", "traveltime", "studytime", "failures", "famrel", "freetime", "goout", "Dalc", "Walc", "health", "absences", "G1", "G2", "G3", ] INTEGER_COLUMNS = set(NUMERIC_COLUMNS) DEFAULT_CATEGORY_FILL = "Unknown" SUMMARY_PREVIEW_LIMIT = 4 def _format_numeric_value(value: object) -> str: if isinstance(value, int): return str(value) if isinstance(value, float): if value.is_integer(): return str(int(value)) return f"{value:.2f}" return str(value) def _coerce_boolean_series(series: pd.Series) -> Tuple[pd.Series, Optional[bool], int]: """Normalise boolean-like columns and return (series, fill_value, missing_count).""" normalized = [] for value in series: if pd.isna(value): normalized.append(pd.NA) continue if isinstance(value, bool): normalized.append(value) continue if isinstance(value, (int, float)) and not pd.isna(value): if value == 1: normalized.append(True) continue if value == 0: normalized.append(False) continue if isinstance(value, str): stripped = value.strip() if not stripped: normalized.append(pd.NA) continue lowered = stripped.lower() if lowered in TRUE_VALUES: normalized.append(True) continue if lowered in FALSE_VALUES: normalized.append(False) continue normalized.append(pd.NA) normal_series = pd.Series(normalized, index=series.index, dtype="object") missing_count = int(normal_series.isna().sum()) dropna = normal_series.dropna() if dropna.empty: fill_value = False else: fill_value = bool(dropna.mode().iloc[0]) normal_series.loc[normal_series.isna()] = fill_value return normal_series.astype(bool), fill_value, missing_count def _coerce_numeric_series( series: pd.Series, as_integer: bool ) -> Tuple[pd.Series, float, int]: """Coerce numeric column, fill with median (rounded for ints).""" cleaned = series.apply(lambda x: x.strip() if isinstance(x, str) else x) numeric = pd.to_numeric(cleaned, errors="coerce") missing_count = int(numeric.isna().sum()) median_value = numeric.median(skipna=True) if pd.isna(median_value): median_value = 0.0 if as_integer: fill_value = int(round(float(median_value))) reported_value = fill_value else: fill_value = float(median_value) reported_value = fill_value numeric = numeric.fillna(fill_value) if as_integer: numeric = numeric.round().astype(int) return numeric, reported_value, missing_count def _coerce_categorical_series(series: pd.Series) -> Tuple[pd.Series, str, int]: """Standardise categorical column and fill with mode.""" cleaned = series.astype("string").str.strip() cleaned = cleaned.replace("", pd.NA) missing_count = int(cleaned.isna().sum()) mode_values = cleaned.dropna().mode() if mode_values.empty: fill_value = DEFAULT_CATEGORY_FILL else: fill_value = str(mode_values.iloc[0]) cleaned = cleaned.fillna(fill_value) return cleaned.astype(str), fill_value, missing_count def _summarise_section( title: str, items: Dict[str, Dict[str, object]], value_formatter, ) -> Optional[str]: if not items: return None entries = sorted(items.items(), key=lambda kv: kv[0]) formatted: List[str] = [] for idx, (column, meta) in enumerate(entries): if idx >= SUMMARY_PREVIEW_LIMIT: remaining = len(entries) - SUMMARY_PREVIEW_LIMIT formatted.append(f"... (+{remaining} more)") break formatted.append(f"{column}={value_formatter(meta)}") return f"{title}: " + ", ".join(formatted) def _format_report_lines(report: Dict[str, Dict[str, Dict[str, object]]]) -> List[str]: """Return human-readable summary lines for console output.""" lines: List[str] = [] numeric_line = _summarise_section( "Numeric medians", report.get("numeric", {}), lambda meta: f"{_format_numeric_value(meta['fill'])} (filled {meta['missing']} rows)", ) if numeric_line: lines.append(numeric_line) boolean_line = _summarise_section( "Boolean mode", report.get("boolean", {}), lambda meta: f"{'True' if meta['fill'] else 'False'} (filled {meta['missing']} rows)", ) if boolean_line: lines.append(boolean_line) categorical_line = _summarise_section( "Categorical mode", report.get("categorical", {}), lambda meta: f"{meta['fill']} (filled {meta['missing']} rows)", ) if categorical_line: lines.append(categorical_line) return lines def _prepare_dataset(csv_path: str) -> Tuple[str, Dict[str, Dict[str, Dict[str, object]]]]: """Load CSV, apply imputation rules, and persist to a temporary cleaned file.""" df = pd.read_csv(csv_path) imputation_report: Dict[str, Dict[str, Dict[str, object]]] = { "numeric": {}, "boolean": {}, "categorical": {}, } for column in BOOLEAN_COLUMNS: if column in df.columns: seri [truncated — 7122 more characters] ``` ### data_assistant.py ```python #!/usr/bin/env python3 """Backward-compatible entry point for the AI Data Assistant CLI.""" from data_assistant import main if __name__ == "__main__": main() ``` ### test_llm_integration.py ```python #!/usr/bin/env python3 """ Test script for LLM integration This script tests the LLM integration without requiring API keys """ import json import os from llm_integration import AnthropicSQLGenerator def test_sql_generation(): """Test SQL generation with sample profile""" # Sample profile data sample_profile = { "dataset": { "filename": "test_data.csv", "row_count": 5, "column_count": 4 }, "columns": [ { "name": "name", "duckdb_type": "VARCHAR", "semantic_type": "category", "examples": ["John", "Jane", "Bob"] }, { "name": "age", "duckdb_type": "BIGINT", "semantic_type": "numeric", "examples": [25, 30, 35] }, { "name": "city", "duckdb_type": "VARCHAR", "semantic_type": "category", "examples": ["New York", "San Francisco", "Chicago"] }, { "name": "salary", "duckdb_type": "BIGINT", "semantic_type": "numeric", "examples": [50000, 75000, 60000] } ] } print("🧪 Testing LLM Integration...") # Test 1: Check if API key is configured api_key = os.getenv("ANTHROPIC_API_KEY") if not api_key or api_key == "your_api_key_here": print("❌ Anthropic API key not configured!") print("Please set ANTHROPIC_API_KEY in your .env file") return False print("✅ API key configured") # Test 2: Test SQL generator initialization try: generator = AnthropicSQLGenerator() print("✅ SQL generator initialized") except Exception as e: print(f"❌ Failed to initialize SQL generator: {e}") return False # Test 3: Test fallback schema generation try: fallback_schema = generator._fallback_schema(sample_profile) print("✅ Fallback schema generation works") print("Sample fallback schema:") print(fallback_schema) except Exception as e: print(f"❌ Fallback schema generation failed: {e}") return False print("\n🎉 All tests passed! LLM integration is ready.") print("\nTo use the full features:") print("1. Set your Anthropic API key in .env file") print("2. Run: python3 data_assistant.py --csv your_file.csv") print("3. Or run: python3 profiler.py --csv your_file.csv --out output_dir") return True if __name__ == "__main__": test_sql_generation() ``` ### profiler.py ```python #!/usr/bin/env python3 import argparse import hashlib import json import os import sys import time from typing import Any, Dict, List, Optional import duckdb try: import orjson as fastjson def dumps(obj): return fastjson.dumps(obj, option=fastjson.OPT_INDENT_2).decode() except Exception: def dumps(obj): return json.dumps(obj, indent=2, ensure_ascii=False) W_CLIENT = None EMBED = None # Set up connectivity to a local or remote Weaviate instance and, if requested, # load a lightweight embedding model. This makes the profiler capable of pushing # profiles into a RAG store with vectors, without forcing you to install heavy # ML dependencies unless you actually want embeddings. def init_weaviate(url: Optional[str], api_key: Optional[str], embed_model: Optional[str]): global W_CLIENT, EMBED if not url: return try: import weaviate if api_key: W_CLIENT = weaviate.Client(url=url, auth_client_secret=weaviate.AuthApiKey(api_key)) else: W_CLIENT = weaviate.Client(url=url) except Exception as e: print(f"[warn] Weaviate init failed: {e}", file=sys.stderr) W_CLIENT = None if embed_model: try: from sentence_transformers import SentenceTransformer EMBED = SentenceTransformer(embed_model) except Exception as e: print(f"[warn] embedding model load failed: {e}", file=sys.stderr) EMBED = None # Create the collection (aka class) in Weaviate on the fly if it doesn't exist. # We keep the schema minimal and pragmatic so you can start searching right away # and iterate later without ceremony. def ensure_class(class_name: str): if not W_CLIENT or not class_name: return try: schema = W_CLIENT.schema.get() names = {c.get("class") for c in schema.get("classes", [])} if class_name in names: return except Exception: pass try: class_obj = { "class": class_name, "description": "CleanSQL column profiles for hybrid search", "vectorizer": "text2vec-transformers", "properties": [ {"name": "dataset", "dataType": ["text"]}, {"name": "column", "dataType": ["text"]}, {"name": "duckdb_type", "dataType": ["text"]}, {"name": "semantic_type", "dataType": ["text"]}, {"name": "null_ratio", "dataType": ["number"]}, {"name": "distinct_count", "dataType": ["number"]}, {"name": "search_text", "dataType": ["text"]}, {"name": "profile_json", "dataType": ["text"]} ] } W_CLIENT.schema.create_class(class_obj) except Exception as e: print(f"[warn] Weaviate class create failed: {e}", file=sys.stderr) # Turn a short piece of text into a vector using the local model if available. # If no model is loaded, we simply skip vectors so the flow remains optional # and never blocks your basic profiling. def embed_text(text: str) -> Optional[List[float]]: if EMBED is None: return None try: v = EMBED.encode([text], normalize_embeddings=True) return v[0].tolist() except Exception as e: print(f"[warn] embedding failed: {e}", file=sys.stderr) return None # Quick helper to fetch the file size. Not all filesystems behave the same, # so we guard with a try/except and return -1 if anything looks off. def size_bytes(path: str) -> int: try: return os.path.getsize(path) except Exception: return -1 # Read a CSV with DuckDB and compute a compact but useful profile. The goal is # speed and clarity: types, nulls, distincts, numeric stats with outlier bounds, # categorical top values, and basic datetime ranges. No data is modified—this # is purely descriptive so you can build robust plans later. def profile_csv(csv_path: str, sample_rows: Optional[int], topk: int, approx_distinct: bool, date_try_cast_for_varchar: bool) -> Dict[str, Any]: con = duckdb.connect() con.execute("PRAGMA threads=" + str(os.cpu_count() or 8)) file_lower = csv_path.lower() auto_opts = [] if sample_rows: auto_opts.append(f"SAMPLE_SIZE={sample_rows}") if file_lower.endswith('.parquet'): src_sql = f"SELECT * FROM parquet_scan('{csv_path}')" else: src_sql = f"SELECT * FROM read_csv_auto('{csv_path}', {', '.join(auto_opts)})" if auto_opts else f"SELECT * FROM read_csv_auto('{csv_path}')" con.execute(f"CREATE VIEW v AS {src_sql}") schema_rows = con.execute("PRAGMA table_info('v')").fetchall() columns = [{"name": r[1], "duckdb_type": r[2]} for r in schema_rows] row_count = con.execute("SELECT count(*) FROM v").fetchone()[0] col_count = len(columns) dup_rows = 0 try: hash_expr = "hash(" + ", ".join([f'"{c["name"]}"' for c in columns]) + ")" dup_rows = con.execute(f""" WITH h AS ( SELECT {hash_expr} AS h FROM v ) SELECT COALESCE(SUM(cnt - 1), 0) FROM (SELECT h, COUNT(*) AS cnt FROM h GROUP BY h) t WHERE cnt > 1 """).fetchone()[0] except Exception: dup_rows = None profiles: List[Dict[str, Any]] = [] for col in columns: name = col["name"] dtype = col["duckdb_type"].upper() nulls, non_nulls = con.execute(f""" SELECT SUM(CASE WHEN "{name}" IS NULL THEN 1 ELSE 0 END), SUM(CASE WHEN "{name}" IS NOT NULL THEN 1 ELSE 0 END) FROM v """).fetchone() null_ratio = (nulls or 0) / row_count if row_count else 0.0 if approx_distinct: distinct_count = con.execute(f'SELECT approx_count_distinct("{name}") FROM v').fetchone()[0] else: try: distinct_count = con.execut [truncated — 11755 more characters] ``` ### data_assistant/__init__.py ```python """Exports for the AI Data Assistant package.""" from .cli import main, run __all__ = ["main", "run"] ``` ### llm_integration.py ```python #!/usr/bin/env python3 import json import os import re from typing import Any, Dict, List, Optional import pandas as pd import anthropic from dotenv import load_dotenv load_dotenv() class AnthropicSQLGenerator: def __init__(self): self.client = anthropic.Anthropic( api_key=os.getenv("ANTHROPIC_API_KEY") ) @staticmethod def _normalize_sql_output(raw_text: str) -> str: """Extract a usable SQL statement from LLM output.""" if not raw_text: return raw_text text = raw_text.strip() # Prefer content inside fenced code blocks if present code_blocks = re.findall(r"```(?:sql)?\s*(.*?)```", text, flags=re.IGNORECASE | re.DOTALL) if code_blocks: text = code_blocks[0].strip() # Strip leading explanatory sentences before the actual SQL sql_match = re.search( r"(?is)\b(WITH|SELECT|INSERT|UPDATE|DELETE)\b.*", text ) if sql_match: text = sql_match.group(0).strip() # Remove trailing explanations after the SQL statement end_match = re.search(r";[^;]*$", text, flags=re.DOTALL) if end_match and end_match.start() != -1: tail = text[end_match.start():] if not re.search(r";\s*\Z", tail): text = text[: end_match.end()].strip() return text.strip() def generate_schema_sql(self, profile: Dict[str, Any]) -> str: """Generate SQL schema creation from data profile""" prompt = f""" Based on this CSV data profile, generate a simple SQL schema for DuckDB. Dataset Info: - File: {profile['dataset']['filename']} - Rows: {profile['dataset']['row_count']} - Columns: {profile['dataset']['column_count']} Column Details: {json.dumps(profile['columns'], indent=2)} Generate ONLY a CREATE TABLE statement with appropriate data types. Do NOT include indexes, constraints, or ALTER TABLE statements. Use simple DuckDB-compatible syntax. Return only the CREATE TABLE statement, no explanations. """ try: response = self.client.messages.create( model="claude-sonnet-4-5-20250929", max_tokens=1000, messages=[{"role": "user", "content": prompt}] ) # Clean the response to remove markdown formatting sql_text = response.content[0].text # Remove ```sql and ``` markers sql_text = sql_text.replace('```sql', '').replace('```', '').strip() return sql_text except Exception as e: print(f"Error generating schema: {e}") return self._fallback_schema(profile) @staticmethod def _extract_json_payload(text: str) -> Optional[Dict[str, Any]]: """Find and parse the first JSON object in a text blob.""" if not text: return None text = text.strip() json_matches = re.findall(r"\{.*\}", text, flags=re.DOTALL) for candidate in json_matches: try: return json.loads(candidate) except json.JSONDecodeError: continue try: return json.loads(text) except Exception: return None # --- Begin User-Specified RAG Imputation Pipeline --- def choose_policy(dtype, nullp, profile, df): """Minimal policy chooser (works today)""" if dtype == "numeric": if nullp <= 0.01: return {"method":"drop_rows"} if nullp <= 0.15: key = pick_stable_key(df) return {"method":"median", "by":key} return {"method":"iterative", "max_iter":15, "by":None, "fallback":"median"} if dtype in ("categorical","boolean"): key = pick_stable_key(df) if nullp <= 0.10: return {"method":"mode", "by":key} return {"method":"mode", "by":key, "unknown":"Unknown"} if dtype == "date": return {"method":"ffill_bfill"} # or "median_date" return {"method":"drop_rows"} def pick_stable_key(df): """Heuristic: choose 1–2 categorical columns with low cardinality (≤10), good association, enough support (≥30 rows).""" cats = [c for c in df.columns if df[c].dtype == 'object' and df[c].nunique() <= 10] for c in cats: if df[c].value_counts().min() >= 30: return c return None def apply_policies(df, policies, cols): """Apply imputations only to columns actually used; return (df_imputed, report).""" import numpy as np dfq = df.copy() report = {} for c in cols: pol = policies[c] null_mask = dfq[c].isnull() n_missing = null_mask.sum() method = pol.get("method") if method == "drop_rows": dfq = dfq[~null_mask] report[c] = {"imputed":0, "method":"drop_rows"} elif method == "median": by = pol.get("by") if by: medians = dfq.groupby(by)[c].transform('median') dfq.loc[null_mask, c] = medians[null_mask] else: med = dfq[c].median() dfq.loc[null_mask, c] = med report[c] = {"imputed":int(n_missing), "method":"median", "by":by} elif method == "iterative": # Placeholder for IterativeImputer or KNN dfq.loc[null_mask, c] = dfq[c].median() report[c] = {"imputed":int(n_missing), "method":"iterative->median"} elif method == "mode": by = pol.get("by") if by: modes = dfq.groupby(by)[c].transform(lambda x: x.mode().iloc[0] if not x.mode().empty else "Unknown") dfq.loc[n [truncated — 23945 more characters] ```