contextsql-garvit 0.1.0__py3-none-any.whl

contextsql/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

contextsql/agent.py

@@ -0,0 +1,192 @@
+"""Question -> SQL (local LLM) -> validate -> run read-only -> self-correct on errors."""
+import os, re, time, json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from . import feedback, glossary as gl
+from .config import LOG, HOME
+from .hints import FAIL, EMPTY, MISS
+from .phrasing import requested_columns
+from .prompt import build_prompt
+from .retrieve import select
+
+MAX_PROMPT_TOKENS, MAX_ROWS, TIMEOUT_S = 1000, 200, 15
+FORBID = re.compile(r"\b(insert|update|delete|drop|alter|truncate|create|grant|revoke|replace|merge|call|exec|"
+                    r"outfile|dumpfile|load_file|sleep|benchmark|pg_sleep)\b", re.I)
+
+class LlamaGenerator:
+    def __init__(self, path, n_ctx=2048, threads=None):
+        try:
+            from llama_cpp import Llama
+        except ImportError:
+            raise SystemExit("llama-cpp-python is not installed. Run:\n  pip install llama-cpp-python "
+                             "--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu")
+        threads = threads or max(1, os.cpu_count() or 2)
+        self.llm = Llama(model_path=path, n_ctx=n_ctx, n_threads=threads, n_batch=256, verbose=False)
+        try:                                   # reuse the processed prompt prefix between questions
+            from llama_cpp.llama_cache import LlamaRAMCache
+            self.llm.set_cache(LlamaRAMCache(capacity_bytes=256 << 20))
+        except Exception:
+            pass
+
+    def count_tokens(self, text):
+        return len(self.llm.tokenize(text.encode("utf-8"), add_bos=False))
+
+    def generate(self, prompt, temperature=0.0):
+        out = self.llm(prompt, max_tokens=300, temperature=temperature, stop=[";"], echo=False)
+        return out["choices"][0]["text"]
+
+def prune_ddl(meta, keep, max_cols=25):
+    """Show only relevant columns of very wide tables (keeps keys, matched columns and the first few)."""
+    if not keep or len(meta["columns"]) <= max_cols:
+        return meta["ddl"]
+    lines = meta["ddl"].split("\n")
+    head, body, tail = lines[0], lines[1:-1], lines[-1]
+    out = [l for l in body if l.strip().split(" ")[0].rstrip(",") in keep or
+           l.strip().startswith(("PRIMARY", "CONSTRAINT", "FOREIGN"))]
+    out = [l.rstrip(",") for l in out]
+    return head + "\n" + ",\n".join(out) + "\n" + tail
+
+def clean_sql(text):
+    text = re.sub(r"```(?:sql)?", "", text).strip()
+    return " ".join(text.split()).rstrip(";").strip() + ";"
+
+def validate(sql):
+    s = re.sub(r"--[^\n]*|/\*.*?\*/", "", sql, flags=re.S).strip().rstrip(";")
+    s = re.sub(r"'(?:[^']|'')*'", "''", s)          # ignore words inside string literals
+    if ";" in s: return False, "only one statement is allowed"
+    if not re.match(r"^\s*(select|with)\b", s, re.I): return False, "only SELECT queries are allowed"
+    m = FORBID.search(s)
+    if m: return False, f"forbidden keyword: {m.group(0)}"
+    return True, ""
+
+@dataclass
+class Result:
+    question: str
+    sql: str = ""
+    columns: list = field(default_factory=list)
+    rows: list = field(default_factory=list)
+    ok: bool = False
+    error: str = ""
+    attempts: int = 0
+    latency_ms: int = 0
+    tables: list = field(default_factory=list)
+    truncated: bool = False
+    confidence: str = "high"      # "low": weak table match, retries needed, or still missing requested columns
+
+class Agent:
+    def __init__(self, engine, know, generator, max_attempts=3):
+        self.engine, self.know, self.gen, self.max_attempts = engine, know, generator, max_attempts
+        self._cache = {}
+
+    def _count(self, text):
+        return self.gen.count_tokens(text) if hasattr(self.gen, "count_tokens") else len(text) // 3
+
+    def glossary(self):
+        return {**self.know.get("autogloss", {}), **gl.load()}     # your own entries override automatic ones
+
+    def make_prompt(self, question, hint=None):
+        gloss = self.glossary()
+        sel = select(question, self.know, gloss)
+        tables, notes, terms = list(sel["tables"]), list(sel["notes"]), list(sel["glossary"])
+        req = requested_columns(question, self.know, tables)
+        low = max(sel["scores"].values(), default=0) <= 0
+        shots = feedback.similar(question, tables)
+        while True:
+            ctx = [f"- {k}: {gloss[k]['definition']} SQL: {gloss[k]['sql']}" for k in terms]
+            ctx += [f"- data note: {n['text']}" for n in notes if n["table"] in tables
+                    and (n["kind"] != "join" or n["ref_table"] in tables)]
+            if req: ctx.append(f"- requested columns: {', '.join(req)}")
+            ctx += [f"- example: Q: {q} SQL: {sq}" for q, sq in shots]
+            if hint: ctx.append(f"- note: {hint}")
+            schema = "\n\n".join(prune_ddl(self.know["tables"][t], sel["keep"].get(t)) for t in tables)
+            prompt = build_prompt(schema, ctx, question)
+            if self._count(prompt) <= MAX_PROMPT_TOKENS: break
+            if shots:
+                shots = []                                                  # drop examples first
+            elif any(n["kind"] == "values" for n in notes):
+                notes = [n for n in notes if n["kind"] != "values"]       # drop value lists first
+            elif len(tables) > 1:
+                tables.remove(min(tables, key=lambda t: sel["scores"].get(t, 0)))
+            else:
+                break
+        return prompt, tables, low, req
+
+    def run_sql(self, sql):
+        d = self.engine.dialect.name
+        body = sql.rstrip(";")
+        if d in ("mysql", "postgresql"): body = body.replace("%", "%%")
+        with self.engine.connect() as c:
+            try:
+                if d == "mysql": c.exec_driver_sql("SET SESSION TRANSACTION READ ONLY")
+                elif d == "postgresql": c.exec_driver_sql("SET TRANSACTION READ ONLY")
+            except Exception:
+                pass
+            for stmt in (f"SET SESSION MAX_EXECUTION_TIME={TIMEOUT_S * 1000}", f"SET SESSION max_statement_time={TIMEOUT_S}",
+                         f"SET statement_timeout = {TIMEOUT_S * 1000}"):
+                try:
+                    c.exec_driver_sql(stmt); break
+                except Exception:
+                    continue
+            c.exec_driver_sql(("EXPLAIN QUERY PLAN " if d == "sqlite" else "EXPLAIN ") + body)
+            res = c.exec_driver_sql(body)
+            cols = list(res.keys()); rows = res.fetchmany(MAX_ROWS + 1)
+        return cols, [tuple(r) for r in rows[:MAX_ROWS]], len(rows) > MAX_ROWS
+
+    @staticmethod
+    def missing(req, sql):
+        m = re.search(r"select\s+(.*?)\s+from\b", sql, re.I | re.S)
+        clause = (m.group(1) if m else sql).lower()
+        return [c for c in req if c.lower() not in clause]
+
+    def ask(self, question):
+        key = " ".join(question.lower().split())
+        if key in self._cache:
+            return self._cache[key]
+        t0, r, hint, best, low0 = time.time(), Result(question=question), None, None, False
+        for i in range(1, self.max_attempts + 1):
+            prompt, tables, low, req = self.make_prompt(question, hint)
+            low0 = low0 or (i == 1 and low)
+            raw = self.gen.generate(prompt, temperature=0.0 if i == 1 else 0.2 * (i - 1))
+            r.tables, r.attempts = tables, i
+            if raw.strip().lower().startswith("-- unanswerable"):          # the model says the data cannot answer this
+                r.ok, r.sql, r.confidence = False, "", "low"
+                r.error = raw.strip().split(":", 1)[-1].strip() or "this database has no matching information"
+                best = None
+                break
+            sql = clean_sql(raw)
+            r.sql = sql
+            ok, err = validate(sql)
+            if ok:
+                try:
+                    cols, rows, trunc = self.run_sql(sql)
+                    r.columns, r.rows, r.truncated, r.ok, r.error = cols, rows, trunc, True, ""
+                    miss = self.missing(req, sql)
+                    if best is None or rows or not best[2]:
+                        best = (sql, cols, rows, trunc, tables, i)
+                    if i < self.max_attempts and (not rows or miss):
+                        hint = EMPTY.format(sql=sql) if not rows else MISS.format(cols=", ".join(miss), sql=sql)
+                        r.ok = False
+                        continue
+                    r.confidence = "low" if (low0 or i > 1 or miss or not rows) else "high"
+                    break
+                except Exception as e:
+                    err = str(getattr(e, "orig", e))[:300]
+            r.error = err
+            hint = FAIL.format(err=err, sql=sql)
+        if not r.ok and best:      # retries failed: fall back to the earlier valid answer
+            r.sql, r.columns, r.rows, r.truncated, r.tables, r.attempts = best
+            r.ok, r.error, r.confidence = True, "", "low"
+        if not r.ok:
+            r.confidence = "low"
+        r.latency_ms = int((time.time() - t0) * 1000)
+        try:
+            HOME.mkdir(parents=True, exist_ok=True)
+            with open(LOG, "a", encoding="utf-8") as f:
+                f.write(json.dumps({"ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+                                    "question": question, "sql": r.sql, "ok": r.ok, "error": r.error,
+                                    "attempts": r.attempts, "latency_ms": r.latency_ms, "confidence": r.confidence, "tables": r.tables}) + "\n")
+        except Exception:
+            pass
+        if r.ok:
+            self._cache[key] = r
+        return r

contextsql/autoglossary.py

@@ -0,0 +1,279 @@
+"""Auto-built business glossary: detects common patterns in ANY schema (valid-status filters, revenue,
+'without child', 'active', outcome rates) and can generate verified training examples from them."""
+import re
+from .names import singular, ABBR2FULL, PREFIXES
+from .retrieve import name_alts, all_fks
+
+BAD = ("cancel", "void", "reject", "fail", "refund", "return", "declin", "abandon", "delet", "invalid", "expire",
+       "lost", "no_show", "drop", "default", "withdraw", "lapse", "churn")
+GOOD = ("complet", "deliver", "paid", "resolv", "clos", "won", "approv", "done", "success", "checked", "arriv", "receiv", "hired")
+STATUS = {"status", "statu", "state", "stage", "outcome", "sts"}
+QTY = {"quantity", "qty"}
+PRICE = {"price", "prc"}
+PRICE2 = {"cost", "fare"}
+DISC = {"discount", "disc"}
+TEXT = ("CHAR", "TEXT", "ENUM", "STRING")
+DATE = ("DATE", "TIME")
+NUM = ("INT", "DECIMAL", "NUMERIC", "FLOAT", "DOUBLE", "REAL")
+DISP = {"name", "title", "label", "company"}
+
+def _flat(col):
+    s = set()
+    for a in name_alts(col): s |= a
+    return s
+
+def _is(meta, col, kinds):
+    return any(k in meta["types"][col].upper() for k in kinds)
+
+def _flag(meta, col):
+    t = meta["types"][col].upper()
+    return "TINYINT(1)" in t or "BOOL" in t or t.startswith("BIT")
+
+def _keys(meta):
+    return {c for fk in all_fks(meta) for c in fk["cols"]} | set(meta["pk"])
+
+def cols_of(meta, tokens, kinds=None, exclude=()):
+    return [c for c in meta["columns"] if c not in exclude and _flat(c) & tokens
+            and (kinds is None or _is(meta, c, kinds)) and not _flag(meta, c)]
+
+def label(name, sing=False):
+    n = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
+    parts = [p for p in re.split(r"[_\W]+", n.lower()) if p]
+    if len(parts) > 1 and parts[0] in PREFIXES: parts = parts[1:]
+    words = [ABBR2FULL.get(p, p) for p in parts]
+    if sing and words: words[-1] = singular(words[-1])
+    return " ".join(words)
+
+class Labels:
+    """Plain-English names. With synthetic metadata they are exact; otherwise derived from the physical names."""
+    def __init__(self, meta=None):
+        self.m = {}
+        if meta:
+            for T in meta["tables"].values():
+                self.m[T["phys"]] = (T["label_sg"], T["label_pl"])
+                for c in T["cols"]: self.m[(T["phys"], c["phys"])] = c["label"]
+    def tbl(self, t):
+        if t in self.m: return self.m[t]
+        sg = label(t, True); return sg, sg + "s"
+    def col(self, t, c): return self.m.get((t, c)) or label(c)
+
+def disp_col(meta):
+    for c in meta["columns"]:
+        if _is(meta, c, TEXT) and _flat(c) & DISP: return c
+    for c in meta["columns"]:
+        if _is(meta, c, TEXT): return c
+    return meta["pk"][0] if meta["pk"] else meta["columns"][0]
+
+def cat_cols(know, t, exclude=()):
+    return [c for c in know["tables"][t]["columns"] if f"{t}.{c}" in know.get("values", {}) and c not in exclude]
+
+def build(know, labels=None):
+    """-> {term: entry}. entry: kind, term, definition, tpl (with {a} alias slot), sql, tables, params."""
+    L = labels or Labels()
+    T, E, valid = know["tables"], {}, {}
+    for t, m in T.items():                                    # 1) valid-status filter
+        for col in cols_of(m, STATUS, TEXT):
+            vals = know.get("values", {}).get(f"{t}.{col}", [])
+            bad = sorted({v.lower() for v in vals if any(b in v.lower() for b in BAD)})
+            if bad:
+                sg, _ = L.tbl(t); term = f"valid {sg}"
+                tpl = "LOWER({a}.%s) NOT IN (%s)" % (col, ", ".join(f"'{b}'" for b in bad))
+                valid[t] = dict(col=col, tpl=tpl, term=term)
+                E[term] = dict(kind="valid", term=term, table=t, col=col, tpl=tpl, sql=tpl.format(a="t"), tables=[t],
+                               definition=f"A {sg} whose {L.col(t, col)} is not {' or '.join(bad)} (compare in lowercase).")
+                break
+    for t, m in T.items():                                    # 2) revenue
+        if "revenue" in E: break
+        ex = _keys(m)
+        qty, pr = cols_of(m, QTY, NUM, ex), cols_of(m, PRICE, NUM, ex) or cols_of(m, PRICE2, NUM, ex)
+        if not (qty and pr): continue
+        disc = cols_of(m, DISC, NUM)
+        expr = ("SUM({i}.%s * {i}.%s * (1 - COALESCE({i}.%s,0)/100))" % (qty[0], pr[0], disc[0])) if disc \
+            else "SUM({i}.%s * {i}.%s)" % (qty[0], pr[0])
+        par = next((fk for fk in all_fks(m) if fk["ref_table"] in valid and len(fk["cols"]) == 1), None)
+        sg, pl = L.tbl(t)
+        defi = f"Sum of {L.col(t, qty[0])} times {L.col(t, pr[0])}{' after discount' if disc else ''}, "
+        if par:
+            psg, _ = L.tbl(par["ref_table"]); defi += f"for valid {psg}s only."
+        else:
+            defi += f"over all {pl}."
+        E["revenue"] = dict(kind="revenue_A", term="revenue", table=t, tpl=expr, sql=expr.format(i="i"), definition=defi,
+                            tables=[t] + ([par["ref_table"]] if par else []), par=par, qty=qty[0], price=pr[0])
+    if "revenue" not in E:
+        for t, m in T.items():
+            qty = cols_of(m, QTY, NUM, _keys(m))
+            if not qty: continue
+            for fk in all_fks(m):
+                M = fk["ref_table"]
+                if M not in T or len(fk["cols"]) != 1: continue
+                mp = cols_of(T[M], PRICE, NUM, _keys(T[M])) or cols_of(T[M], PRICE2, NUM, _keys(T[M]))
+                if mp:
+                    expr = "SUM({i}.%s * {m}.%s)" % (qty[0], mp[0])
+                    E["revenue"] = dict(kind="revenue_B", term="revenue", table=t, tpl=expr, sql=expr.format(i="i", m="m"),
+                                        definition=f"Sum of {L.col(t, qty[0])} times {L.col(M, mp[0])} of the linked {L.tbl(M)[0]}.",
+                                        tables=[t, M], par=fk, qty=qty[0], price=mp[0], pt=M)
+                    break
+            if "revenue" in E: break
+    n_wo = n_act = 0
+    for ct, m in T.items():                                   # 3) without-child, 4) active parent
+        for fk in all_fks(m):
+            P = fk["ref_table"]
+            if P not in T or len(fk["cols"]) != 1 or P == ct: continue
+            psg, ppl = L.tbl(P); csg, cpl = L.tbl(ct)
+            term = f"{psg} without {csg}"
+            if term not in E and n_wo < 6:
+                n_wo += 1
+                tpl = "NOT EXISTS (SELECT 1 FROM %s c WHERE c.%s = {a}.%s)" % (ct, fk["cols"][0], fk["ref_cols"][0])
+                E[term] = dict(kind="without", term=term, table=P, child=ct, fk=fk, tpl=tpl, sql=tpl.format(a="p"),
+                               tables=[P, ct], definition=f"A {psg} that has no {csg} records.")
+            dcols = [c for c in m["columns"] if _is(m, c, DATE) and c not in _keys(m)]
+            term = f"active {psg}"
+            if dcols and term not in E and n_act < 6:
+                n_act += 1
+                tpl = "{a}.%s >= DATE_SUB(CURDATE(), INTERVAL 90 DAY)" % dcols[0]
+                E[term] = dict(kind="active", term=term, table=P, child=ct, fk=fk, date=dcols[0], tpl=tpl, sql=tpl.format(a="c"),
+                               tables=[P, ct], definition=f"A {psg} with at least one {csg} in the last 90 days (by {L.col(ct, dcols[0])}).")
+    for t, m in T.items():                                    # 5) outcome share
+        for col in cols_of(m, STATUS, TEXT):
+            vals = know.get("values", {}).get(f"{t}.{col}", [])
+            good = [v for v in vals if any(g in v.lower() for g in GOOD)]
+            if good:
+                v, sg = good[0], L.tbl(t)[0]; term = f"{v.lower().replace('_', ' ')} {sg}"
+                tpl = "ROUND(100 * SUM({a}.%s = '%s') / COUNT(*), 2)" % (col, v)
+                E.setdefault(term, dict(kind="outcome", term=term, table=t, col=col, value=v, tpl=tpl, sql=tpl.format(a="t"),
+                                        tables=[t], definition=f"Percentage of {L.tbl(t)[1]} whose {L.col(t, col)} is {v}."))
+                break
+    for t, m in T.items():                                    # 6) soft delete -> "current" records
+        for col in m["columns"]:
+            if "deleted" in _flat(col) and col not in _keys(m):
+                sg, _ = L.tbl(t); term = f"current {sg}"
+                fl = _flag(m, col)
+                tpl = ("{a}.%s = 0" if fl else "{a}.%s IS NULL") % col
+                E[term] = dict(kind="current", term=term, table=t, col=col, tpl=tpl, sql=tpl.format(a="t"), tables=[t],
+                               definition=f"A {sg} that has not been deleted ({L.col(t, col)} {'is 0' if fl else 'is empty'}).")
+                break
+    return E
+
+def probe_sql(e):
+    k = e["kind"]
+    if k == "valid": return f"SELECT COUNT(*) FROM {e['table']} t WHERE {e['tpl'].format(a='t')}"
+    if k == "revenue_A": return f"SELECT {e['tpl'].format(i='i')} FROM {e['table']} i"
+    if k == "revenue_B":
+        fk = e["par"]; return f"SELECT {e['tpl'].format(i='i', m='m')} FROM {e['table']} i JOIN {e['pt']} m ON i.{fk['cols'][0]} = m.{fk['ref_cols'][0]}"
+    if k == "without": return f"SELECT COUNT(*) FROM {e['table']} p WHERE {e['tpl'].format(a='p')}"
+    if k == "active":
+        fk = e["fk"]; return (f"SELECT COUNT(*) FROM {e['table']} p JOIN {e['child']} c ON c.{fk['cols'][0]} = p.{fk['ref_cols'][0]} "
+                              f"WHERE {e['tpl'].format(a='c')}")
+    return f"SELECT {e['tpl'].format(a='t')} FROM {e['table']} t"
+
+def build_verified(know, engine):
+    """Entries whose probe query actually runs on this database, in the public glossary format."""
+    out = {}
+    with engine.connect() as c:
+        for term, e in build(know).items():
+            try:
+                c.exec_driver_sql(probe_sql(e).replace("%", "%%") if engine.dialect.name in ("mysql", "postgresql") else probe_sql(e)).fetchall()
+            except Exception:
+                continue
+            out[term] = dict(definition=e["definition"], sql=e["sql"], tables=e["tables"], auto=True)
+    return out
+
+def examples(know, labels, entries):
+    """Training examples that use the glossary entries. -> [{family, variants, sql, tables, used, terms, req, extra}]"""
+    L, T, out = labels, know["tables"], []
+    def add(fam, variants, sql, tables, terms, used=()):
+        out.append(dict(family=fam, variants=variants, sql=" ".join(sql.split()), tables=list(dict.fromkeys(tables)),
+                        used=list(used), terms=list(terms), req=[], extra=False))
+    for term, e in entries.items():
+        k, t = e["kind"], e["table"]
+        sg, pl = L.tbl(t)
+        if k == "valid":
+            f = e["tpl"].format(a="t"); m = T[t]
+            add("gl_valid_count", [f"How many valid {pl} are there?", f"Number of valid {pl}", f"Count the valid {pl}"],
+                f"SELECT COUNT(*) FROM {t} t WHERE {f}", [t], [term], [f"{t}.{e['col']}"])
+            for c in cat_cols(know, t, {e["col"]})[:2]:
+                add("gl_valid_group", [f"Number of valid {pl} per {L.col(t, c)}", f"How many valid {pl} for each {L.col(t, c)}?"],
+                    f"SELECT t.{c}, COUNT(*) AS n FROM {t} t WHERE {f} GROUP BY t.{c} ORDER BY n DESC", [t], [term], [f"{t}.{c}"])
+            for n in [c for c in m["columns"] if _is(m, c, NUM) and not _flag(m, c) and c not in _keys(m)][:2]:
+                add("gl_valid_sum", [f"Total {L.col(t, n)} of valid {pl}", f"What is the total {L.col(t, n)} across valid {pl}?"],
+                    f"SELECT ROUND(SUM(t.{n}),2) AS total FROM {t} t WHERE {f}", [t], [term])
+            for d in [c for c in m["columns"] if _is(m, c, DATE) and c not in _keys(m)][:1]:
+                for days in (90, 365):
+                    add("gl_valid_recent", [f"How many valid {pl} in the last {days} days?", f"Number of valid {pl} within the past {days} days"],
+                        f"SELECT COUNT(*) FROM {t} t WHERE {f} AND t.{d} >= DATE_SUB(CURDATE(), INTERVAL {days} DAY)", [t], [term])
+        elif k in ("revenue_A", "revenue_B"):
+            fk, par = e.get("par"), e.get("par")
+            ex = e["tpl"].format(i="i", m="m")
+            if k == "revenue_A":
+                join, where, ptab, terms = f"{t} i", "", None, [term]
+                if fk:
+                    ptab = fk["ref_table"]; vt = next(x for x in entries.values() if x["kind"] == "valid" and x["table"] == ptab)
+                    join = f"{t} i JOIN {ptab} p ON i.{fk['cols'][0]} = p.{fk['ref_cols'][0]}"
+                    where = " WHERE " + vt["tpl"].format(a="p"); terms = [term, vt["term"]]
+            else:
+                ptab = e["pt"]; join = f"{t} i JOIN {ptab} m ON i.{fk['cols'][0]} = m.{fk['ref_cols'][0]}"; where, terms = "", [term]
+            al = "p" if k == "revenue_A" else "m"
+            add("gl_revenue", ["What is our total revenue?", "Total revenue", "How much revenue have we made?"],
+                f"SELECT ROUND({ex},2) AS revenue FROM {join}{where}", [t] + ([ptab] if ptab else []), terms)
+            if ptab:
+                pm = T[ptab]
+                for c in cat_cols(know, ptab)[:2]:
+                    if k == "revenue_A" and c == next(x for x in entries.values() if x["kind"] == "valid" and x["table"] == ptab)["col"]: continue
+                    add("gl_revenue_group", [f"Revenue by {L.col(ptab, c)}", f"Show revenue per {L.col(ptab, c)}"],
+                        f"SELECT {al}.{c}, ROUND({ex},2) AS revenue FROM {join}{where} GROUP BY {al}.{c} ORDER BY revenue DESC", [t, ptab], terms, [f"{ptab}.{c}"])
+                if k == "revenue_A":
+                    d, pk = disp_col(pm), (pm["pk"] or [pm["columns"][0]])[0]
+                    psg, ppl = L.tbl(ptab)
+                    for kk in (3, 5, 10):
+                        add("gl_revenue_top", [f"Top {kk} {ppl} by revenue", f"Which {kk} {ppl} bring in the most revenue?"],
+                            f"SELECT p.{d}, ROUND({ex},2) AS revenue FROM {join}{where} GROUP BY p.{pk}, p.{d} ORDER BY revenue DESC LIMIT {kk}", [t, ptab], terms)
+                    for dc in [c for c in pm["columns"] if _is(pm, c, DATE) and c not in _keys(pm)][:1]:
+                        for days in (90, 365):
+                            w2 = (where + " AND " if where else " WHERE ") + f"p.{dc} >= DATE_SUB(CURDATE(), INTERVAL {days} DAY)"
+                            add("gl_revenue_recent", [f"Revenue in the last {days} days", f"How much revenue did we make in the past {days} days?"],
+                                f"SELECT ROUND({ex},2) AS revenue FROM {join}{w2}", [t, ptab], terms)
+        elif k == "without":
+            P, C, fk = t, e["child"], e["fk"]; psg, ppl = L.tbl(P); csg, cpl = L.tbl(C); d = disp_col(T[P])
+            f = e["tpl"].format(a="p")
+            add("gl_without_count", [f"How many {ppl} have no {cpl}?", f"Number of {ppl} without any {cpl}", f"Count {ppl} that have never had a {csg}"],
+                f"SELECT COUNT(*) FROM {P} p WHERE {f}", [P, C], [term])
+            for kk in (5, 10):
+                add("gl_without_list", [f"List {kk} {ppl} without {cpl}", f"Show {kk} {ppl} that have no {cpl}"],
+                    f"SELECT p.{d} FROM {P} p WHERE {f} LIMIT {kk}", [P, C], [term])
+        elif k == "active":
+            P, C, fk = t, e["child"], e["fk"]; psg, ppl = L.tbl(P); csg, cpl = L.tbl(C)
+            f = e["tpl"].format(a="c"); pk = fk["ref_cols"][0]
+            base = f"FROM {P} p JOIN {C} c ON c.{fk['cols'][0]} = p.{pk} WHERE {f}"
+            add("gl_active_count", [f"How many active {ppl} do we have?", f"Number of active {ppl}", f"Count active {ppl}"],
+                f"SELECT COUNT(DISTINCT p.{pk}) AS active {base}", [P, C], [term])
+            for c2 in cat_cols(know, P)[:2]:
+                add("gl_active_group", [f"Active {ppl} by {L.col(P, c2)}", f"How many active {ppl} per {L.col(P, c2)}?"],
+                    f"SELECT p.{c2}, COUNT(DISTINCT p.{pk}) AS active {base} GROUP BY p.{c2} ORDER BY active DESC", [P, C], [term], [f"{P}.{c2}"])
+        elif k == "current":
+            f = e["tpl"].format(a="t"); m = T[t]
+            add("gl_current_count", [f"How many current {pl} do we have?", f"Number of {pl} that are not deleted", f"Count the active {pl} records"],
+                f"SELECT COUNT(*) FROM {t} t WHERE {f}", [t], [term])
+            add("gl_deleted_count", [f"How many {pl} have been deleted?", f"Number of deleted {pl}"],
+                f"SELECT COUNT(*) FROM {t} t WHERE NOT ({f})", [t], [term])
+            for c in cat_cols(know, t)[:2]:
+                add("gl_current_group", [f"Number of current {pl} per {L.col(t, c)}", f"{L.col(t, c)} wise count of current {pl}"],
+                    f"SELECT t.{c}, COUNT(*) AS n FROM {t} t WHERE {f} GROUP BY t.{c} ORDER BY n DESC", [t], [term], [f"{t}.{c}"])
+            for n in [c for c in m["columns"] if _is(m, c, NUM) and not _flag(m, c) and c not in _keys(m)][:1]:
+                add("gl_current_sum", [f"Total {L.col(t, n)} of current {pl}", f"What is the total {L.col(t, n)} across current {pl}?"],
+                    f"SELECT ROUND(SUM(t.{n}),2) AS total FROM {t} t WHERE {f}", [t], [term])
+        elif k == "outcome":
+            v, col = e["value"], e["col"]; vtxt = v.replace("_", " ")
+            add("gl_outcome_pct", [f"What percentage of {pl} are {vtxt}?", f"Share of {pl} that are {vtxt}", f"What is the {vtxt} rate of {pl}?"],
+                f"SELECT ROUND(100 * SUM(t.{col} = '{v}') / COUNT(*), 2) AS pct FROM {t} t", [t], [term], [f"{t}.{col}"])
+            for c in cat_cols(know, t, {col})[:1]:
+                add("gl_outcome_group", [f"Percentage of {pl} that are {vtxt} by {L.col(t, c)}", f"{vtxt.capitalize()} rate of {pl} per {L.col(t, c)}"],
+                    f"SELECT t.{c}, ROUND(100 * SUM(t.{col} = '{v}') / COUNT(*), 2) AS pct FROM {t} t GROUP BY t.{c} ORDER BY pct DESC", [t], [term], [f"{t}.{c}"])
+    for t, m in T.items():                                    # flag columns (no glossary term needed)
+        sg, pl = L.tbl(t)
+        for c in [c for c in m["columns"] if _flag(m, c)][:2]:
+            fl = re.sub(r"^is ", "", L.col(t, c))
+            add("flag_count", [f"How many {pl} are {fl}?", f"Number of {pl} that are {fl}", f"Count the {fl} {pl}"],
+                f"SELECT COUNT(*) FROM {t} WHERE {c} = 1", [t], [])
+            add("flag_pct", [f"What percentage of {pl} are {fl}?", f"Share of {pl} that are {fl}"],
+                f"SELECT ROUND(100 * AVG({c}), 2) AS pct FROM {t}", [t], [])
+    return out

contextsql/cli.py

@@ -0,0 +1,165 @@
+"""contextsql connect | learn | ask | glossary | status"""
+import argparse, getpass, json, os, sys
+from pathlib import Path
+from . import glossary as gl
+from .config import CONFIG, KNOW, LOG, load_json, save_json
+from .db import make_url, get_engine
+
+HF_REPO = "Garvitarunav/contextsql-v4"
+HF_FILE = "contextsql-v4-q4_k_m.gguf"
+
+def default_model():
+    p = Path("models/contextsql-q4_k_m.gguf")
+    if p.exists():
+        return str(p.resolve())
+    from huggingface_hub import hf_hub_download
+    print("Downloading model (1.8 GB, first time only) ...")
+    return hf_hub_download(HF_REPO, HF_FILE)
+
+def need(path, hint):
+    obj = load_json(path)
+    if obj is None:
+        sys.exit(f"Not ready yet - run: {hint}")
+    return obj
+
+def cmd_connect(a):
+    d = a.dialect
+    if d == "sqlite":
+        db = a.db or input("SQLite file path: ")
+        url = make_url("sqlite", database=db)
+    else:
+        host = a.host or input("Host [127.0.0.1]: ") or "127.0.0.1"
+        port = a.port or int(input(f"Port [{3306 if d == 'mysql' else 5432}]: ") or (3306 if d == "mysql" else 5432))
+        user = a.user or input("User (a READ-ONLY user is strongly recommended): ")
+        pw = a.password if a.password is not None else getpass.getpass("Password: ")
+        db = a.db or input("Database name: ")
+        url = make_url(d, host, port, user, pw, db)
+    try:
+        with get_engine(url).connect() as c:
+            c.exec_driver_sql("SELECT 1")
+    except Exception as e:
+        sys.exit(f"Could not connect: {getattr(e, 'orig', e)}")
+    model = a.model or default_model()
+    save_json(CONFIG, {"url": url, "model": model})
+    print("Connected and saved.")
+    print("Next: contextsql learn")
+
+def cmd_learn(a):
+    from .learn import learn
+    cfg = need(CONFIG, "contextsql connect")
+    print("Learning your database (one-time) ...")
+    know, imported = learn(get_engine(cfg["url"]))
+    rows = sum(t.get("rows", 0) for t in know["tables"].values())
+    print(f"\nDone: {len(know['tables'])} tables, {rows:,} rows scanned, {len(know['notes'])} data notes, "
+          f"{len(know.get('autogloss', {}))} auto glossary terms, {imported} glossary terms imported.")
+    for n in know["notes"][:6]:
+        print("  -", n["text"])
+    if not gl.load():
+        print("\nTip: teach it your key metrics (2 minutes):\n"
+              '  contextsql glossary add "revenue" "Sum of quantity*price for completed orders" '
+              '"SUM(quantity*unit_price)" --tables orders,order_items')
+    print("\nNext: contextsql ask \"your question\"")
+
+def show(r):
+    print(f"\nSQL ({r.attempts} attempt{'s' if r.attempts != 1 else ''}, {r.latency_ms} ms; tables: {', '.join(r.tables)}):\n  {r.sql}")
+    if r.confidence == "low":
+        print("  (low confidence - please check the SQL above, or rephrase / add a glossary term)")
+    if not r.ok:
+        print(f"\nCould not answer: {r.error}"); return
+    if not r.rows:
+        print("\n(no rows)"); return
+    cells = [[("NULL" if v is None else str(v)) for v in row] for row in r.rows[:20]]
+    w = [max(len(str(c)), *(len(row[i]) for row in cells)) for i, c in enumerate(r.columns)]
+    print("\n" + "  ".join(str(c).ljust(w[i]) for i, c in enumerate(r.columns)))
+    print("  ".join("-" * x for x in w))
+    for row in cells: print("  ".join(v.ljust(w[i]) for i, v in enumerate(row)))
+    more = len(r.rows) - 20
+    print(f"\n{len(r.rows)}{'+' if r.truncated else ''} row(s)" + (f" (showing 20, {more} more)" if more > 0 else ""))
+
+def cmd_ask(a):
+    from .agent import Agent, LlamaGenerator
+    cfg = need(CONFIG, "contextsql connect")
+    know = need(KNOW, "contextsql learn")
+    if not cfg.get("model") or not Path(cfg["model"]).exists():
+        cfg["model"] = default_model()
+        save_json(CONFIG, cfg)
+    print("Loading model ...")
+    agent = Agent(get_engine(cfg["url"]), know, LlamaGenerator(cfg["model"], threads=a.threads))
+    q = " ".join(a.question).strip()
+    if q:
+        show(agent.ask(q)); return
+    print("Ask anything about your data (empty line to quit).")
+    while True:
+        q = input("\n> ").strip()
+        if not q: break
+        show(agent.ask(q))
+
+def cmd_glossary(a):
+    if a.action == "add":
+        if not (a.term and a.definition and a.sql): sys.exit('usage: contextsql glossary add "term" "definition" "sql" [--tables a,b]')
+        gl.add(a.term, a.definition, a.sql, [t for t in (a.tables or "").split(",") if t])
+        print(f"Added '{a.term}'. Re-run learn is not needed.")
+    elif a.action == "remove":
+        print("Removed." if gl.remove(a.term) else "Not found.")
+    else:
+        auto = (load_json(KNOW) or {}).get("autogloss", {})
+        for k, g in {**auto, **gl.load()}.items():
+            tag = " (auto)" if k in auto and k not in gl.load() else ""
+            print(f"- {k}{tag}: {g['definition']}\n    SQL: {g['sql']}  tables: {', '.join(g['tables'])}")
+
+def _last():
+    try:
+        return json.loads(open(LOG, encoding="utf-8").read().strip().splitlines()[-1])
+    except Exception:
+        sys.exit("No previous question found - ask something first.")
+
+def cmd_good(a):
+    from . import feedback
+    e = _last()
+    if not e.get("ok"):
+        sys.exit('The last answer did not succeed. Use: contextsql fix "SELECT ..."')
+    feedback.save(e["question"], e["sql"], e.get("tables", []))
+    print("Saved. It will be used as an example for similar questions.")
+
+def cmd_fix(a):
+    from . import feedback
+    from .agent import Agent, validate, clean_sql
+    e, sql = _last(), clean_sql(" ".join(a.sql))
+    ok, err = validate(sql)
+    if not ok: sys.exit(f"Not saved: {err}")
+    cfg = need(CONFIG, "contextsql connect")
+    try:
+        Agent(get_engine(cfg["url"]), None, None).run_sql(sql)
+    except Exception as ex:
+        sys.exit(f"Not saved - that SQL does not run: {getattr(ex, 'orig', ex)}")
+    feedback.save(e["question"], sql, e.get("tables", []))
+    print("Saved your corrected SQL for that question.")
+
+def cmd_status(a):
+    cfg, know = load_json(CONFIG), load_json(KNOW)
+    print("connected:", bool(cfg), "| model:", (cfg or {}).get("model") or "-")
+    if know:
+        print(f"learned: {len(know['tables'])} tables, {len(know['notes'])} notes, at {know['learned_at']}")
+    print("glossary terms:", len(gl.load()))
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="contextsql", description="Ask your database questions in plain language - local and free.")
+    s = p.add_subparsers(dest="cmd", required=True)
+    c = s.add_parser("connect"); c.set_defaults(fn=cmd_connect)
+    c.add_argument("--dialect", default="mysql", choices=["mysql", "postgresql", "sqlite"])
+    c.add_argument("--host"); c.add_argument("--port", type=int); c.add_argument("--user")
+    c.add_argument("--password"); c.add_argument("--db"); c.add_argument("--model")
+    s.add_parser("learn").set_defaults(fn=cmd_learn)
+    k = s.add_parser("ask"); k.add_argument("question", nargs="*"); k.add_argument("--threads", type=int); k.set_defaults(fn=cmd_ask)
+    g = s.add_parser("glossary"); g.set_defaults(fn=cmd_glossary)
+    g.add_argument("action", choices=["add", "list", "remove"], nargs="?", default="list")
+    g.add_argument("term", nargs="?"); g.add_argument("definition", nargs="?"); g.add_argument("sql", nargs="?")
+    g.add_argument("--tables")
+    s.add_parser("good").set_defaults(fn=cmd_good)
+    f = s.add_parser("fix"); f.add_argument("sql", nargs="+"); f.set_defaults(fn=cmd_fix)
+    s.add_parser("status").set_defaults(fn=cmd_status)
+    a = p.parse_args(argv)
+    a.fn(a)
+
+if __name__ == "__main__":
+    main()