contextsql-garvit 0.1.0__py3-none-any.whl

contextsql/config.py

@@ -0,0 +1,17 @@
+"""Local storage: everything lives in ~/.contextsql (override with CONTEXTSQL_HOME)."""
+import json, os
+from pathlib import Path
+
+HOME = Path(os.getenv("CONTEXTSQL_HOME", Path.home() / ".contextsql"))
+CONFIG, KNOW, GLOSS, LOG = (HOME / "config.json", HOME / "knowledge.json",
+                            HOME / "glossary.json", HOME / "query_log.jsonl")
+
+def load_json(path, default=None):
+    try:
+        return json.loads(Path(path).read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        return default
+
+def save_json(path, obj):
+    HOME.mkdir(parents=True, exist_ok=True)
+    Path(path).write_text(json.dumps(obj, indent=2, ensure_ascii=False), encoding="utf-8")

contextsql/db.py

@@ -0,0 +1,13 @@
+from sqlalchemy import create_engine
+from sqlalchemy.engine import URL
+
+DRIVERS = {"mysql": "mysql+pymysql", "postgresql": "postgresql+psycopg2"}
+
+def make_url(dialect, host=None, port=None, user=None, password=None, database=None):
+    if dialect == "sqlite":
+        return f"sqlite:///{database}"
+    return URL.create(DRIVERS[dialect], username=user, password=password, host=host,
+                      port=port, database=database).render_as_string(hide_password=False)
+
+def get_engine(url):
+    return create_engine(url, pool_pre_ping=True)

contextsql/feedback.py

@@ -0,0 +1,35 @@
+"""Saved question->SQL pairs (confirmed good or user-corrected). Reused as examples for similar questions."""
+import json
+from .config import HOME
+from .retrieve import toks
+
+FILE = HOME / "examples.jsonl"
+
+def save(question, sql, tables):
+    HOME.mkdir(parents=True, exist_ok=True)
+    with open(FILE, "a", encoding="utf-8") as f:
+        f.write(json.dumps({"question": question, "sql": sql, "tables": list(tables)}) + "\n")
+
+def load():
+    try:
+        return [json.loads(l) for l in open(FILE, encoding="utf-8") if l.strip()]
+    except FileNotFoundError:
+        return []
+
+def similar(question, tables, k=2, thresh=0.5):
+    q = set(toks(question))
+    scored = []
+    for ex in load():
+        e = set(toks(ex["question"]))
+        if not q or not e or not set(ex.get("tables", [])) <= set(tables):
+            continue
+        j = len(q & e) / len(q | e)
+        if j >= thresh:
+            scored.append((j, ex))
+    scored.sort(key=lambda x: -x[0])
+    out, seen = [], set()
+    for _, ex in scored:
+        if ex["sql"] not in seen:
+            seen.add(ex["sql"]); out.append((ex["question"], ex["sql"]))
+        if len(out) >= k: break
+    return out

contextsql/glossary.py

@@ -0,0 +1,27 @@
+"""User-editable business glossary: term -> definition + SQL snippet + tables."""
+from .config import GLOSS, load_json, save_json
+
+def load():
+    return load_json(GLOSS, {})
+
+def add(term, definition, sql, tables=()):
+    g = load()
+    g[term.lower()] = {"definition": definition, "sql": sql, "tables": list(tables)}
+    save_json(GLOSS, g)
+
+def remove(term):
+    g = load()
+    found = g.pop(term.lower(), None) is not None
+    save_json(GLOSS, g)
+    return found
+
+def merge_imported(rows):
+    """rows: (term, definition, sql, tables) - never overwrites entries the user already has."""
+    g, added = load(), 0
+    for term, definition, sql, tables in rows:
+        if term.lower() not in g:
+            g[term.lower()] = {"definition": definition, "sql": sql,
+                               "tables": [t.strip() for t in (tables or "").split(",") if t.strip()]}
+            added += 1
+    save_json(GLOSS, g)
+    return added

contextsql/hints.py

@@ -0,0 +1,5 @@
+"""Retry hints - used identically by the agent and by the training-data builder."""
+FAIL = "the previous attempt failed ({err}). Previous SQL: {sql} Write a corrected query."
+EMPTY = "the previous attempt returned no rows. Previous SQL: {sql} Check the filter values against the data notes and fix the query."
+MISS = "the answer must also include the columns: {cols}. Previous SQL: {sql} Add them to the SELECT list (and GROUP BY)."
+UNANSWERABLE = "-- UNANSWERABLE:"

contextsql/introspect.py

@@ -0,0 +1,121 @@
+"""Read schema + profile the data (null rates, categorical values, near-duplicate spellings)."""
+import re, difflib
+from .names import singular, PREFIXES
+from sqlalchemy import inspect
+
+TEXT_HINTS = ("CHAR", "TEXT", "ENUM", "STRING", "CLOB")
+SAMPLE_ABOVE, SAMPLE_SIZE = 200_000, 100_000
+
+def clean_mysql_ddl(raw):
+    raw = raw.replace("`", "")
+    lines = raw.split("\n")
+    head, body = lines[0], [l.strip().rstrip(",") for l in lines[1:-1]]
+    body = [l for l in body if not l.startswith("KEY ") and not l.startswith("UNIQUE KEY")]
+    body = [re.sub(r"\s+(CHARACTER SET \w+|COLLATE \w+)", "", l) for l in body]
+    body = [re.sub(r"\s+DEFAULT NULL", "", l) for l in body]
+    return head + "\n  " + ",\n  ".join(body) + "\n);"
+
+def col_type(engine, col):
+    try:
+        return str(col["type"].compile(dialect=engine.dialect))
+    except Exception:
+        return "UNKNOWN"
+
+def ddl_for(engine, insp, table):
+    if engine.dialect.name == "mysql":
+        with engine.connect() as c:
+            raw = c.exec_driver_sql(f"SHOW CREATE TABLE `{table}`").fetchone()[1]
+        return clean_mysql_ddl(raw)
+    lines = []
+    for col in insp.get_columns(table):
+        lines.append(f"{col['name']} {col_type(engine, col)}" + ("" if col.get("nullable", True) else " NOT NULL"))
+    pk = insp.get_pk_constraint(table).get("constrained_columns") or []
+    if pk: lines.append(f"PRIMARY KEY ({', '.join(pk)})")
+    for fk in insp.get_foreign_keys(table):
+        lines.append(f"FOREIGN KEY ({', '.join(fk['constrained_columns'])}) REFERENCES "
+                     f"{fk['referred_table']} ({', '.join(fk['referred_columns'])})")
+    return f"CREATE TABLE {table} (\n  " + ",\n  ".join(lines) + "\n);"
+
+def similar_groups(values, threshold=0.9):
+    """Find values that look like spelling/case variants of each other."""
+    flagged = set()
+    for i, a in enumerate(values):
+        for b in values[i + 1:]:
+            if a.lower() == b.lower() or difflib.SequenceMatcher(None, a.lower(), b.lower()).ratio() >= threshold:
+                flagged |= {a, b}
+    return flagged
+
+def profile(engine, tables, max_distinct=30):
+    """Returns (notes, values). notes: [{table, column, kind, text}], values: {'table.col': [..]}"""
+    q = engine.dialect.identifier_preparer.quote
+    notes, values = [], {}
+    with engine.connect() as c:
+        for t, meta in tables.items():
+            qt = q(t)
+            n = c.exec_driver_sql(f"SELECT COUNT(*) FROM {qt}").scalar() or 0
+            meta["rows"] = int(n)
+            if n == 0:
+                continue
+            src = qt if n <= SAMPLE_ABOVE else f"(SELECT * FROM {qt} LIMIT {SAMPLE_SIZE}) AS s"
+            cols = meta["columns"]
+            is_text = {col: any(h in meta["types"][col].upper() for h in TEXT_HINTS) for col in cols}
+            parts = [f"SUM(CASE WHEN {q(col)} IS NULL THEN 1 ELSE 0 END)" for col in cols]
+            tcols = [col for col in cols if is_text[col]]
+            parts += [f"COUNT(DISTINCT {q(col)})" for col in tcols]
+            row = c.exec_driver_sql(f"SELECT COUNT(*), {', '.join(parts)} FROM {src}").fetchone()
+            total, nulls, distincts = row[0] or 1, row[1:1 + len(cols)], row[1 + len(cols):]
+            for col, nl in zip(cols, nulls):
+                pct = 100 * (nl or 0) / total
+                if 1 <= pct < 100 and col not in meta["pk"]:
+                    notes.append(dict(table=t, column=col, kind="null",
+                                      text=f"{t}.{col} can be NULL ({pct:.0f}% of rows); use COALESCE when aggregating."))
+            for col, d in zip(tcols, distincts):
+                if not d or d > max_distinct:
+                    continue
+                rows = c.exec_driver_sql(
+                    f"SELECT {q(col)}, COUNT(*) FROM {src} WHERE {q(col)} IS NOT NULL "
+                    f"GROUP BY {q(col)} ORDER BY COUNT(*) DESC LIMIT {max_distinct}").fetchall()
+                vals = [str(r[0]) for r in rows]
+                values[f"{t}.{col}"] = vals
+                bad = similar_groups(vals)
+                if bad:
+                    notes.append(dict(table=t, column=col, kind="variants", text=(
+                        f"{t}.{col} has inconsistent spelling/case ({', '.join(vals)}); "
+                        f"compare with LOWER() and treat similar spellings as the same.")))
+                elif d <= 12:
+                    notes.append(dict(table=t, column=col, kind="values",
+                                      text=f"{t}.{col} has values: {', '.join(vals)}."))
+    return notes, values
+
+
+def _norm(s):
+    return re.sub(r"[^a-z0-9]", "", s.lower())
+
+def _base(table):
+    parts = [p for p in re.split(r"[_\W]+", re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", table).lower()) if p]
+    if len(parts) > 1 and parts[0] in PREFIXES: parts = parts[1:]
+    return "".join(parts)
+
+def infer_joins(tables):
+    """Guess foreign keys from column names when the database does not declare them."""
+    notes = []
+    for t, meta in tables.items():
+        declared = {c for fk in meta["fks"] for c in fk["cols"]}
+        meta["fks_inferred"] = []
+        for col in meta["columns"]:
+            if col in declared or col in meta["pk"]:
+                continue
+            nc = _norm(col)
+            for p, pm in tables.items():
+                if p == t or len(pm["pk"]) != 1:
+                    continue
+                pk = pm["pk"][0]
+                cands = {_norm(singular(_base(p))) + "id", _norm(_base(p)) + "id"}
+                if _norm(pk) != "id":
+                    cands.add(_norm(pk))
+                if nc in cands:
+                    meta["fks_inferred"].append({"cols": [col], "ref_table": p, "ref_cols": [pk]})
+                    notes.append(dict(table=t, column=col, kind="join", ref_table=p,
+                                      text=f"{t}.{col} references {p}.{pk} (inferred from column names)."))
+                    break
+    return notes

contextsql/learn.py

@@ -0,0 +1,48 @@
+"""One-time scan: schema + data profile + glossary import -> ~/.contextsql/knowledge.json"""
+from datetime import datetime, timezone
+from sqlalchemy import inspect
+from . import glossary, autoglossary
+from .config import KNOW, save_json
+from .introspect import ddl_for, col_type, profile, infer_joins
+
+EXCLUDE_DEFAULT = {"query_log", "business_glossary"}
+
+def scan(engine, exclude=EXCLUDE_DEFAULT, say=print):
+    """Read schema + profile data. Returns (knowledge, all_table_names). Saves nothing."""
+    insp = inspect(engine)
+    all_tables = insp.get_table_names()
+    tables = {}
+    for t in all_tables:
+        if t in exclude:
+            continue
+        cols = insp.get_columns(t)
+        tables[t] = {
+            "ddl": ddl_for(engine, insp, t),
+            "columns": [c["name"] for c in cols],
+            "types": {c["name"]: col_type(engine, c) for c in cols},
+            "pk": insp.get_pk_constraint(t).get("constrained_columns") or [],
+            "fks": [{"cols": f["constrained_columns"], "ref_table": f["referred_table"],
+                     "ref_cols": f["referred_columns"]} for f in insp.get_foreign_keys(t)],
+        }
+        say(f"  schema   {t} ({len(cols)} columns)")
+    say("  profiling data ...")
+    notes, values = profile(engine, tables)
+    notes += infer_joins(tables)
+    know = {"dialect": engine.dialect.name, "tables": tables, "notes": notes, "values": values,
+            "learned_at": datetime.now(timezone.utc).isoformat(timespec="seconds")}
+    try:
+        know["autogloss"] = autoglossary.build_verified(know, engine)
+    except Exception:
+        know["autogloss"] = {}
+    return know, all_tables
+
+def learn(engine, exclude=EXCLUDE_DEFAULT, say=print):
+    know, all_tables = scan(engine, exclude, say)
+    imported = 0
+    if "business_glossary" in all_tables:
+        with engine.connect() as c:
+            rows = c.exec_driver_sql(
+                "SELECT term, definition, sql_snippet, tables_used FROM business_glossary").fetchall()
+        imported = glossary.merge_imported([tuple(r) for r in rows])
+    save_json(KNOW, know)
+    return know, imported

contextsql/names.py

@@ -0,0 +1,29 @@
+"""Name helpers: abbreviation expansion + simple singularization."""
+FULL2ABBR = {"customer": "cust", "quantity": "qty", "amount": "amt", "employee": "emp", "department": "dept",
+             "number": "no", "description": "descr", "address": "addr", "date": "dt", "price": "prc",
+             "order": "ord", "product": "prod", "invoice": "inv", "account": "acct", "transaction": "txn",
+             "patient": "pat", "doctor": "doc", "appointment": "appt", "student": "stud", "course": "crs",
+             "shipment": "shp", "warehouse": "wh", "subscription": "sub", "ticket": "tkt", "property": "prop",
+             "listing": "lst", "member": "mem", "book": "bk", "room": "rm", "booking": "bkg", "salary": "sal",
+             "category": "cat", "status": "sts", "created": "crt", "supplier": "supp", "manager": "mgr",
+             "policy": "pol", "claim": "clm", "passenger": "pax", "flight": "flt", "aircraft": "acft",
+             "subscriber": "subs", "meter": "mtr", "reading": "rdg", "farm": "frm", "crop": "crp",
+             "harvest": "hrv", "project": "proj", "worker": "wrk", "task": "tsk", "material": "mtl",
+             "event": "evt", "attendee": "atnd", "donor": "dnr", "donation": "dntn", "campaign": "cmpgn",
+             "volunteer": "vlnt", "vehicle": "veh", "owner": "ownr", "service": "svc", "drug": "drg",
+             "purchase": "purch", "stock": "stk", "candidate": "cand", "application": "appl",
+             "interview": "intv", "artist": "artst", "track": "trk", "stream": "strm", "listener": "lstnr",
+             "reservation": "resv", "discount": "disc", "percent": "pct"}
+ABBR2FULL = {v: k for k, v in FULL2ABBR.items()}
+ABBR2FULL.update({"cus": "customer", "acc": "account", "desc": "description", "trx": "transaction",
+                  "tx": "transaction", "nbr": "number", "amnt": "amount", "num": "number", "qnty": "quantity",
+                  "emp": "employee", "dpt": "department", "cust": "customer", "prd": "product", "ord": "order",
+                  "qty": "quantity", "amt": "amount", "prc": "price", "dt": "date", "sts": "status"})
+PREFIXES = {"tbl", "t", "tb", "mst", "dim"}
+
+def singular(w):
+    if w.endswith("statuses"): return w[:-2]
+    if w.endswith("ies") and len(w) > 4: return w[:-3] + "y"
+    if w.endswith("sses"): return w[:-2]
+    if w.endswith("s") and not w.endswith("ss") and len(w) > 3: return w[:-1]
+    return w

contextsql/phrasing.py

@@ -0,0 +1,61 @@
+"""Request-phrase helpers: requested-column extraction (inference + training) and paraphrasing (training)."""
+import re
+from .retrieve import toks, name_alts, hit, all_fks, GENERIC_COLS
+
+CONN = r"\b(?:with their|with its|along with|alongside|including|plus|together with|as well as|with)\b"
+BLOCK = {"at", "least", "more", "than", "over", "under", "less", "only", "between", "greater", "fewer"}
+CONN_CHOICES = ["with their", "along with", "alongside", "including", "plus", "together with", "as well as", "with"]
+
+def requested_columns(question, know, tables):
+    """Columns the user asked to see ('... with country and email'). A known data value means it is a filter instead."""
+    parts = re.split(CONN, question, flags=re.I)
+    if len(parts) < 2:
+        return []
+    seg = parts[-1]
+    if re.search(r"\d", seg) or BLOCK & set(seg.lower().split()):
+        return []
+    q = set(toks(seg))
+    if not q:
+        return []
+    for vals in know.get("values", {}).values():
+        for v in vals:
+            vt = set(toks(v))
+            if vt and vt <= q:
+                return []
+    out = []
+    for t in tables:
+        meta = know["tables"][t]
+        skip = {c for fk in all_fks(meta) for c in fk["cols"]} | set(meta["pk"])
+        alts_t = name_alts(t)
+        tname = set().union(*alts_t) if alts_t else set()
+        for col in meta["columns"]:
+            alts = [a for a in name_alts(col) if not (a & GENERIC_COLS)]
+            if not alts or col in skip or set().union(*alts) <= tname:
+                continue
+            if hit(alts, q) and col not in out:
+                out.append(col)
+    return out
+
+WORD_SWAPS = {"total": ["total", "overall"], "average": ["average", "mean"], "highest": ["highest", "largest", "biggest"],
+              "lowest": ["lowest", "smallest"], "per": ["per", "for each"], "Top": ["Top", "Best"],
+              "give": ["give", "show", "list"], "Number": ["Number", "Count"]}
+
+def paraphrase(q, rng, extra=False):
+    for k, opts in WORD_SWAPS.items():
+        if re.search(rf"\b{k}\b", q) and rng.random() < 0.4:
+            q = re.sub(rf"\b{k}\b", rng.choice(opts), q, count=1)
+    if extra and rng.random() < 0.8:
+        q = re.sub(CONN, rng.choice(CONN_CHOICES), q, count=1, flags=re.I)
+    return q
+
+def typo(q, rng, p=0.06):
+    """Light real-world noise: drop or swap one letter in a long word."""
+    if rng.random() > p:
+        return q
+    ws = q.split(" ")
+    idx = [i for i, w in enumerate(ws) if len(w) >= 6 and w.isalpha()]
+    if not idx:
+        return q
+    i = rng.choice(idx); w = ws[i]; j = rng.randrange(1, len(w) - 2)
+    ws[i] = w[:j] + w[j + 1:] if rng.random() < 0.5 else w[:j] + w[j + 1] + w[j] + w[j + 2:]
+    return " ".join(ws)

contextsql/prompt.py

@@ -0,0 +1,5 @@
+DATE_RULE = "Date rule: 'last month' = order_date >= DATE_FORMAT(CURDATE() - INTERVAL 1 MONTH,'%Y-%m-01') AND order_date < DATE_FORMAT(CURDATE(),'%Y-%m-01')"
+
+def build_prompt(schema, ctx_lines, question):
+    ctx = "\n".join(list(ctx_lines) + [DATE_RULE]) if ctx_lines else DATE_RULE
+    return f"### Schema\n{schema}\n\n### Business context\n{ctx}\n\n### Question\n{question}\n\n### SQL\n"

contextsql/retrieve.py

@@ -0,0 +1,110 @@
+"""Pick relevant tables, columns, glossary entries and data notes for a question (no extra model needed)."""
+import re
+from .names import ABBR2FULL, PREFIXES
+
+STOP = set("""the a an of in on by for per to from and or with without is are was were be our we us my me i you your
+all any show list give get find what which who whom how many much top best most least highest lowest total number count
+last past each every please can could tell about than then that this these those it its there their them do does did has have
+had day days week weeks month months year years""".split())
+GENERIC_COLS = {"id", "name", "date", "type", "value", "created", "updated", "at"}
+
+def stem(w):
+    if w.endswith("ies") and len(w) > 4: return w[:-3] + "y"
+    if w.endswith("s") and len(w) > 3 and not w.endswith("ss"): return w[:-1]
+    return w
+
+def toks(text):
+    return [stem(w) for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOP and not w.isdigit()]
+
+def name_alts(name):
+    """[{alternatives for token 1}, {alternatives for token 2}, ...]  e.g. tbl_cust -> [{cust, customer}]"""
+    name = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
+    parts = [p for p in re.split(r"[_\W]+", name.lower()) if p]
+    if len(parts) > 1 and parts[0] in PREFIXES: parts = parts[1:]
+    out = []
+    for p in parts:
+        alt = {stem(p)}
+        if p in ABBR2FULL: alt.add(stem(ABBR2FULL[p]))
+        out.append(alt)
+    return out
+
+def skel(w):
+    return w[0] + re.sub(r"[aeiou]", "", w[1:])
+
+def tok_match(t, w):
+    if t == w:
+        return True
+    if len(t) >= 4 and len(w) >= 4 and (w.startswith(t) or t.startswith(w)):
+        return True
+    return len(t) >= 4 and len(w) >= 5 and skel(w) == t      # cstmr == skeleton of "customer"
+
+def hit(alts, q):
+    return bool(alts) and all(any(tok_match(t, w) for t in alt for w in q) for alt in alts)
+
+def all_fks(meta):
+    return list(meta["fks"]) + list(meta.get("fks_inferred", []))
+
+def select(question, know, gloss, max_tables=5):
+    q = set(toks(question))
+    tables = know["tables"]
+    score = {t: 0.0 for t in tables}
+    hit_cols = set()
+
+    for t, meta in tables.items():
+        if hit(name_alts(t), q): score[t] += 4
+        fk_cols = {c for fk in all_fks(meta) for c in fk["cols"]}
+        for col in meta["columns"]:
+            if col in fk_cols: continue
+            alts = [a for a in name_alts(col) if not (a & GENERIC_COLS)]
+            if alts and hit(alts, q):
+                score[t] += 1; hit_cols.add(f"{t}.{col}")
+    for key, vals in know.get("values", {}).items():
+        t = key.split(".")[0]
+        for v in vals:
+            vt = set(toks(v))
+            if vt and vt <= q:
+                score[t] += 2; hit_cols.add(key)
+
+    matched = [k for k, g in gloss.items() if set(toks(k)) and set(toks(k)) <= q]
+    for k in list(matched):
+        d = gloss[k]["definition"].lower()
+        matched += [o for o in gloss if o not in matched and o != k and o in d]
+    for k in matched:
+        for t in gloss[k].get("tables", []):
+            if t in score: score[t] += 5
+
+    chosen = [t for t, s in sorted(score.items(), key=lambda x: -x[1]) if s > 0][:max_tables]
+    if not chosen:
+        chosen = sorted(tables, key=lambda t: -tables[t].get("rows", 0))[:min(4, len(tables))]
+
+    if len(chosen) <= 1:   # one table matched: include its parent tables so joins stay possible
+        for t in list(chosen):
+            for fk in all_fks(tables[t]):
+                if fk["ref_table"] in tables and fk["ref_table"] not in chosen and len(chosen) < 3:
+                    chosen.append(fk["ref_table"])
+
+    for t, meta in tables.items():   # bridge tables connecting two chosen tables
+        if t in chosen: continue
+        links = {fk["ref_table"] for fk in all_fks(meta)} | {o for o, m in tables.items()
+                                                             if any(fk["ref_table"] == t for fk in all_fks(m))}
+        if len(links & set(chosen)) >= 2 and len(chosen) < max_tables + 1:
+            chosen.append(t)
+
+    notes = []
+    for n in know.get("notes", []):
+        if n["table"] not in chosen: continue
+        if n["kind"] == "join":
+            if n["ref_table"] in chosen: notes.append(n)
+        elif n["kind"] in ("variants", "null") or f"{n['table']}.{n['column']}" in hit_cols:
+            notes.append(n)
+
+    # columns worth keeping when a table is too wide to show in full
+    gtext = " ".join(gloss[k]["sql"] + " " + gloss[k]["definition"] for k in matched).lower()
+    keep = {}
+    for t in chosen:
+        meta = tables[t]
+        ks = set(meta["pk"]) | {c for fk in all_fks(meta) for c in fk["cols"]} | set(meta["columns"][:6])
+        ks |= {c for c in meta["columns"] if f"{t}.{c}" in hit_cols or c.lower() in gtext}
+        ks |= {n["column"] for n in notes if n["table"] == t}
+        keep[t] = ks
+    return {"tables": chosen, "glossary": matched, "notes": notes, "scores": score, "keep": keep}