lessismore/lessismore.py

"""lessismore — squeeze text before it hits an LLM.

Deterministic passes first (free, safe, cacheable). Optional ML pruning
(LLMLingua-2) only when installed and only worth it on big inputs.

    from lessismore import compress
    small = compress(big_log, level=2)

    $ python3 lessismore.py dump.log -l 2 > small.log
    $ tail -5000 app.log | python3 lessismore.py -l 2 | llm ...
"""

import json
import re
import sys
from collections import Counter
from functools import lru_cache

# ---------------------------------------------------------------- tokens

@lru_cache(maxsize=1)
def _encoder():
    try:
        import tiktoken
        return tiktoken.get_encoding("o200k_base")
    except ImportError:
        return None


def count_tokens(text: str) -> int:
    enc = _encoder()
    if enc:
        return len(enc.encode(text))
    return max(1, len(text) // 4)  # ponytail: chars/4 heuristic; pip install tiktoken for real counts


# ---------------------------------------------------------------- passes
# Each pass is (str) -> str. Order matters: whitespace before dedupe.

def collapse_whitespace(text: str) -> str:
    text = re.sub(r"[ \t]+$", "", text, flags=re.M)        # trailing whitespace
    text = re.sub(r"(?<=\S)[ \t]{2,}", " ", text)          # interior runs (leading indent kept: code-safe)
    text = re.sub(r"\n{3,}", "\n\n", text)                 # blank-line runs
    return text


def dedupe_lines(text: str) -> str:
    """Collapse runs of identical lines — the classic log killer."""
    lines = text.split("\n")
    out, i = [], 0
    while i < len(lines):
        j = i
        while j < len(lines) and lines[j] == lines[i]:
            j += 1
        run = j - i
        # marker must actually be shorter than the lines it replaces
        if run >= 4 and lines[i].strip() and (run - 1) * (len(lines[i]) + 1) > 45:
            out.append(lines[i])
            out.append(f"[previous line repeated {run - 1} more times]")
        else:
            out.extend(lines[i:j])
        i = j
    # ponytail: only consecutive repeats; add block-level dedupe (repeated stack traces) if logs demand it
    return "\n".join(out)


# no '/' in the class: URL paths are ≥64-char alnum+slash runs and they ARE the content.
# '-' and '_' admitted for base64url (real JWTs), guarded below so slugs survive.
_BLOB = re.compile(r"\b(?:[A-Za-z0-9+_-]{64,}={0,2}|[0-9a-fA-F]{48,})\b")


def _squash_blob(m: "re.Match") -> str:
    s = m.group()
    # a run with - or _ might be a kebab/snake identifier, not a blob: demand
    # blob-typical entropy (digits + both cases) before eating it
    if ("-" in s or "_" in s) and not (any(c.isdigit() for c in s)
                                       and s.lower() != s and s.upper() != s):
        return s
    return f"{s[:12]}[+{len(s) - 16} chars]{s[-4:]}"


def squash_blobs(text: str) -> str:
    """Base64/hex runs are token-dense and semantically opaque — keep head and tail.

    The tail suffix keeps two different blobs from squashing to the same stub
    and then being falsely merged as "repeated" by dedupe_lines.
    """
    return _BLOB.sub(_squash_blob, text)


_MON = r"(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)"
_TIMESTAMP = re.compile(
    r"\b\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:[.,]\d+)?(?:Z|[+-]\d{2}:?\d{2})?\b ?"  # ISO-8601
    r"|\b" + _MON + r" {1,2}\d{1,2} \d{2}:\d{2}:\d{2}\b ?"                              # syslog
    r"|\b\d{2}/" + _MON + r"/\d{4}:\d{2}:\d{2}:\d{2}(?: [+-]\d{4})?\b ?"                # nginx CLF
)


def strip_timestamps(text: str) -> str:
    """Line order already encodes sequence; per-line timestamps are ~8 tokens each."""
    # ponytail: epoch timestamps skipped on purpose — any 10-digit number matches
    return _TIMESTAMP.sub("", text)


def minify_json(text: str) -> str:
    """Pretty JSON → minified, lossless when it fires, untouched when not.

    Fires on the whole input AND on multi-line JSON embedded in other text
    ('response body:\\n{...}' — the common shape of captured tool output).
    Single-line JSON (JSONL) is left alone: ', ' and ',' are 1 token either way;
    only the newline+indent of pretty-printing costs real tokens.
    """
    try:
        return json.dumps(json.loads(text), separators=(",", ":"), ensure_ascii=False)
    except ValueError:
        pass
    dec = json.JSONDecoder()
    out, last = [], 0
    for m in re.finditer(r"[{\[](?=[ \t]*\n)", text):  # pretty JSON opens then breaks the line
        if m.start() < last:
            continue  # inside a block we already minified
        try:
            obj, end = dec.raw_decode(text, m.start())
        except ValueError:
            continue
        mini = json.dumps(obj, separators=(",", ":"), ensure_ascii=False)
        if len(mini) < end - m.start():
            out += [text[last:m.start()], mini]
            last = end
    out.append(text[last:])
    return "".join(out)


def collapse_cr(text: str) -> str:
    """Keep only the final state of \\r-overwritten progress lines (pip/tqdm/wget)."""
    text = text.replace("\r\n", "\n")
    return "\n".join(l.rsplit("\r", 1)[-1] for l in text.split("\n"))


_ANSI = re.compile(r"\x1b\[[0-9;?]*[A-Za-z]")  # ponytail: CSI only; add OSC if titles show up


def strip_ansi(text: str) -> str:
    """Color codes carry nothing for an LLM, and they make identical lines differ."""
    return _ANSI.sub("", text)


_UUID = re.compile(r"\b[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}\b")


def squash_uuids(text: str) -> str:
    """36 chars → 8-hex prefix, git-short-hash style; cross-references still resolve."""
    return _UUID.sub(lambda m: m.group()[:8] + "…", text)


_PKGPATH = re.compile(r'[^\s"\']+/(?:site-packages|dist-packages|lib/python3\.\d+)/')


def squash_pkgpaths(text: str) -> str:
    """Traceback path spam: …/httpx/_client.py:1054 is still unique without the venv prefix."""
    return _PKGPATH.sub("…/", text)


# the lookbehinds keep "not just X" from becoming "not X" — meaning inversion
_FILLER = re.compile(
    r"(?<!\bnot )(?<!n't )"
    r"\b(?:could you please|can you please|i would like you to|i want you to|"
    r"go ahead and|hey there,?|please|kindly|basically|actually|currently|"
    r"really|simply|just|very|quite)\b ?",
    re.I,
)


def strip_filler(text: str) -> str:
    """Aggressive: eats words like 'just' anywhere, including inside strings. Prose only."""
    return _FILLER.sub("", text)


# caveman-style word dropping: every function word is a whole token.
# NEVER add negations (not/no/never), modals (must/should), or order words
# (before/after) — dropping those changes meaning, not just style.
_STICKS = re.compile(
    r"\b(?:the|a|an|is|are|was|were|be|been|being|am|i|we|you|they|it|"
    r"that|which|who|have|has|had|do|does|did|there)\b ?",
    re.I,
)


def two_sticks(text: str) -> str:
    """Gist-only: strips text to caveman. Fine for articles/transcripts, never instructions."""
    # keep IT/US-style acronyms that case-insensitively collide with function words
    return _STICKS.sub(lambda m: m.group() if m.group().strip().isupper()
                       and len(m.group().strip()) > 1 else "", text)


_SKEL = re.compile(r"[^A-Za-z]+")
_DIGITS = re.compile(r"\d+")


def _variant_summary(lines: list) -> str:
    """Digits are often THE signal (status codes, ports, counts) — summarize the
    values a similar-line collapse would otherwise silently eat."""
    runs = [_DIGITS.findall(l) for l in lines]
    if not runs[0] or any(len(r) != len(runs[0]) for r in runs):
        return ""
    parts = []
    for col in zip(*runs):
        vals = sorted({int(v) for v in col})
        if len(vals) == 1:
            continue
        parts.append("/".join(map(str, vals)) if len(vals) <= 4
                      else f"{vals[0]}–{vals[-1]}")
    return "; values " + ", ".join(parts) if parts else ""


def dedupe_similar(text: str) -> str:
    """Collapse runs of lines identical after masking non-letters — 'same words,
    different numbers' (progress lines, per-item CI steps). Keeps first and last
    so progression endpoints survive, and summarizes the varying values in the
    marker so distinct facts (status codes, ports) aren't lost."""
    lines = text.split("\n")
    out, i = [], 0
    while i < len(lines):
        k, j = _SKEL.sub(" ", lines[i]).strip(), i
        while j < len(lines) and _SKEL.sub(" ", lines[j]).strip() == k:
            j += 1
        if j - i >= 4 and k:
            out += [lines[i],
                    f"[{j - i - 2} similar lines omitted{_variant_summary(lines[i:j])}]",
                    lines[j - 1]]
        else:
            out.extend(lines[i:j])
        i = j
    return "\n".join(out)


_REF = re.compile(r"@\d+")


def alias_repeats(text: str, min_count: int = 4, min_len: int = 30) -> str:
    """Dictionary-code scattered duplicate lines that consecutive dedupe can't reach.

    Lossless — the legend keeps every line verbatim. A code only pays when it
    replaces a repeated multi-token sequence; single words are already 1 BPE
    token each, so word-level codebooks lose (measured, see README).
    """
    lines = text.split("\n")
    if any(_REF.fullmatch(l) for l in lines):  # already aliased, or real @N content — bail
        return text
    counts = Counter(l for l in lines if len(l) >= min_len)
    # ponytail: c*len>200 chars is the payoff heuristic; tune if legends ever dominate
    worth = [l for l, c in counts.items() if c >= min_count and c * len(l) > 200]
    if not worth:
        return text
    ref = {l: f"@{i}" for i, l in enumerate(worth, 1)}
    legend = [f"@{i} = {l}" for i, l in enumerate(worth, 1)]
    return "\n".join(["[repeated lines aliased:]"] + legend + [""] +
                     [ref.get(l, l) for l in lines])


# order: strippers leave doubled spaces, so collapse_whitespace runs after them;
# normalized lines then match better in alias_repeats/dedupe_lines/dedupe_similar
_STRIP2 = [minify_json, collapse_cr, strip_ansi, strip_timestamps,
           squash_blobs, squash_uuids, squash_pkgpaths]
_MERGE = [collapse_whitespace, alias_repeats, dedupe_lines, dedupe_similar]
LEVELS = {
    1: [collapse_whitespace, dedupe_lines],                       # code-safe
    2: _STRIP2 + _MERGE,                                          # logs/dumps/tool output
    3: _STRIP2 + [strip_filler] + _MERGE,                         # prose
    4: _STRIP2 + [strip_filler, two_sticks] + _MERGE,             # gist-only
}


def compress(text: str, level: int = 2) -> str:
    for f in LEVELS[level]:
        text = f(text)
    return text


def budget(text: str, max_tokens: int) -> str:
    """Hard cap: keep head and tail lines, drop the middle with a marker.

    The only pass with GUARANTEED bounded output — run it after compression,
    as the backstop, never instead of it.
    """
    total = count_tokens(text)
    if total <= max_tokens:
        return text
    lines = text.split("\n")
    limit = max(0, max_tokens - 8)  # reserve for the marker line
    head, tail = [], []
    h_tok = t_tok = 0
    hi, ti = 0, len(lines) - 1
    while hi <= ti:
        take_head = h_tok <= t_tok
        line = lines[hi] if take_head else lines[ti]
        tok = count_tokens(line) + 1
        if h_tok + t_tok + tok > limit:
            break
        if take_head:
            head.append(line); h_tok += tok; hi += 1
        else:
            tail.append(line); t_tok += tok; ti -= 1
    if not head and not tail:  # one giant line (e.g. minified JSON): slice by chars
        keep = limit * 2  # ~4 chars/token, half per side
        head, tail = [text[:keep]], [text[-keep:]]
        h_tok = count_tokens(head[0]); t_tok = count_tokens(tail[0])
    omitted = max(0, total - h_tok - t_tok)
    return "\n".join(head + ["[~%d tokens omitted]" % omitted] + list(reversed(tail)))


def serve(port=7777):
    """Paste-in demo page on localhost. Zero deps, binds 127.0.0.1 only."""
    import html
    import http.server
    import urllib.parse

    page = ("<!doctype html><meta charset=utf-8><title>lessismore</title>"
            "<style>body{font-family:system-ui;max-width:900px;margin:2rem auto;padding:0 1rem;"
            "background:#10141a;color:#e6edf3}textarea{width:100%;height:35vh;font-family:monospace;"
            "background:#1a2028;color:#e6edf3;border:1px solid #2d3743;border-radius:8px;padding:8px}"
            "select,button{font-size:16px;padding:8px 14px;border-radius:8px;border:1px solid #2d3743;"
            "background:#1f6feb;color:#fff;cursor:pointer}select{background:#1a2028}</style>"
            "<h1>lessismore \U0001F4C9</h1>"
            "<form method=post><textarea name=text placeholder='paste your ugliest log here'></textarea>"
            "<p>level <select name=level><option>1<option selected>2<option>3<option>4</select> "
            "<button>squeeze</button></p></form>{result}")

    class Handler(http.server.BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass

        def _page(self, result=""):
            body = page.replace("{result}", result).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            self._page()

        def do_POST(self):
            n = int(self.headers.get("Content-Length", 0))
            q = urllib.parse.parse_qs(self.rfile.read(n).decode("utf-8", "replace"))
            text = q.get("text", [""])[0]
            level = min(4, max(1, int(q.get("level", ["2"])[0])))
            out = compress(text, level)
            b, a = count_tokens(text), count_tokens(out)
            self._page("<p><b>{:,} → {:,} tokens ({}% saved)</b></p>"
                       "<textarea readonly>{}</textarea>".format(
                           b, a, round(100 * (1 - a / max(b, 1))), html.escape(out)))

    print("lessismore UI on http://localhost:%d (Ctrl+C to stop)" % port)
    http.server.ThreadingHTTPServer(("127.0.0.1", port), Handler).serve_forever()


# ---------------------------------------------------------------- optional ML pass

@lru_cache(maxsize=1)
def _llmlingua():
    from llmlingua import PromptCompressor  # pip install llmlingua
    return PromptCompressor(
        model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
        use_llmlingua2=True,
    )


def compress_ml(text: str, rate: float = 0.5) -> str:
    """Perplexity-based token pruning (LLMLingua-2).

    Runs a local classifier model — only pays for itself on multi-KB inputs.
    Run the deterministic passes first; never feed it code you need verbatim.
    """
    return _llmlingua().compress_prompt(text, rate=rate)["compressed_prompt"]


# ---------------------------------------------------------------- CLI

def run_command(cmd: str) -> "tuple[str, int]":
    """Run CMD in a shell, return (merged raw output, exit code) — exit code
    survives, unlike piping through a filter."""
    import subprocess
    r = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    return r.stdout.decode("utf-8", "replace"), r.returncode


def hook(level: int = 2, min_chars: int = 2000) -> None:
    """Claude Code PostToolUse hook mode: read the hook JSON on stdin, emit
    updatedToolOutput JSON with the tool output compressed. Zero deps, no jq.

    Small outputs pass through untouched (marker lines aren't worth the churn).
    """
    try:
        d = json.load(sys.stdin)
    except ValueError:
        return  # not hook JSON — emit nothing, Claude Code keeps the original
    out = d.get("tool_response") or d.get("tool_output") or {}
    if isinstance(out, dict):
        raw = out.get("stdout") or out.get("output") or ""
    else:
        raw = str(out)
    if len(raw) < min_chars:
        return
    small = compress(raw, level)
    if len(small) >= len(raw):
        return
    json.dump({"hookSpecificOutput": {"hookEventName": "PostToolUse",
                                      "updatedToolOutput": {"stdout": small}}},
              sys.stdout)


def main():
    import argparse

    p = argparse.ArgumentParser(prog="lessismore", description=__doc__.splitlines()[0])
    p.add_argument("file", nargs="?", help="input file (default: stdin)")
    p.add_argument("-l", "--level", type=int, default=2, choices=sorted(LEVELS),
                   help="1=code-safe 2=logs/dumps 3=prose 4=gist-only caveman (default 2)")
    p.add_argument("--ml", type=float, metavar="RATE",
                   help="also run LLMLingua-2 keeping RATE of tokens (needs: pip install llmlingua)")
    p.add_argument("--budget", type=int, metavar="N",
                   help="hard cap output at ~N tokens: keep head+tail, drop the middle")
    p.add_argument("--run", metavar="CMD",
                   help="run CMD in a shell, print its output compressed, exit with its status")
    p.add_argument("--hook", action="store_true",
                   help="Claude Code PostToolUse hook mode (reads hook JSON on stdin)")
    p.add_argument("--serve", nargs="?", const=7777, type=int, metavar="PORT",
                   help="serve a paste-in demo page on localhost (default port 7777)")
    a = p.parse_args()

    if a.serve:
        return serve(a.serve)
    if a.hook:
        return hook(a.level)

    if a.run:
        raw, code = run_command(a.run)
    else:
        # newline="" / buffer.read(): keep \r intact for collapse_cr
        raw = (open(a.file, encoding="utf-8", errors="replace", newline="").read() if a.file
               else sys.stdin.buffer.read().decode("utf-8", "replace"))
        code = 0
    out = compress(raw, a.level)
    if a.ml:
        try:
            out = compress_ml(out, a.ml)
        except ImportError:
            sys.exit("--ml needs: pip install llmlingua")
    if a.budget:
        out = budget(out, a.budget)
    sys.stdout.write(out)
    before, after = count_tokens(raw), count_tokens(out)
    print(f"lessismore: {before} → {after} tokens ({1 - after / max(before, 1):.0%} saved)",
          file=sys.stderr)
    sys.exit(code)


if __name__ == "__main__":
    main()