DataEyond/Agentic-Service-Data-Eyond

1import pandas as pd2from db_pipeline.connector import get_engine3from sqlalchemy import inspect4 5EXCLUDED_TABLES = {"knowledge_chunks"}6 7 8def get_schema(engine=None) -> dict[str, list[dict]]:9    """Returns {table_name: [{name, type}, ...]} for all tables."""10    if engine is None:11        engine = get_engine()12    inspector = inspect(engine)13    schema = {}14    for table_name in inspector.get_table_names():15        if table_name in EXCLUDED_TABLES:16            continue17        cols = inspector.get_columns(table_name)18        schema[table_name] = [{"name": c["name"], "type": str(c["type"])} for c in cols]19    return schema20 21# for now, table level. but later, will change to column level22def profile_table(engine, table_name: str, sample_size: int = 5) -> dict:23    """Returns row_count, null_counts, value_ranges, and sample_rows."""24    df_sample = pd.read_sql(f'SELECT * FROM "{table_name}" LIMIT {sample_size}', engine)25    row_count = pd.read_sql(f'SELECT COUNT(*) FROM "{table_name}"', engine).iloc[0, 0]26    null_counts = df_sample.isnull().sum().to_dict()27    numeric_cols = df_sample.select_dtypes(include="number").columns.tolist()28    value_ranges = {29        col: {"min": df_sample[col].min(), "max": df_sample[col].max()}30        for col in numeric_cols31    }32    return {33        "row_count": row_count,34        "null_counts": null_counts,35        "value_ranges": value_ranges,36        "sample_rows": df_sample.to_dict(orient="records"),37    }38 39 40def build_text(table_name: str, columns: list[dict], profile: dict) -> str:41    col_lines = "\n".join(f"  - {c['name']} ({c['type']})" for c in columns)42    sample_preview = "\n".join(str(row)[:300] for row in profile["sample_rows"][:3])43    return (44        f"Table: {table_name}\n"45        f"Columns:\n{col_lines}\n"46        f"Row count: {profile['row_count']}\n"47        f"Sample rows:\n{sample_preview}"48    )49