Take files with a question

ask carries paths; the daemon reads them as the user. Text files and a
PDF's text go inline within half the free context, shared equally, and
are stored that way, so a reopened conversation still has them. Images,
and PDF pages with no text, go as image parts to a model the router
reports as seeing images, and are stored as a path and hash, resolved
when sent: a file changed or gone since is reported as no longer
available. turns gains a files column for the chips, added to an older
database on open. A refused file is one log line, not a traceback.

Co-Authored-By: Claude Haiku 5.5 <noreply@anthropic.com>

1 files changed, 307 insertions, 31 deletions

diff --git a/assistant.py b/assistant.py
index 777ceba..7b95255 100644
--- a/assistant.py
+++ b/assistant.py
@@ -21,8 +21,11 @@ finished turn is kept in a SQLite history with full-text search.
 Run: python3 assistant.py        (--test runs the self-check)
 """
 
+import base64
 import glob
+import hashlib
 import json
+import mimetypes
 import os
 import re
 import signal
@@ -147,25 +150,30 @@ def db_open(path):
     conn.execute("PRAGMA journal_mode=WAL")
     conn.execute("PRAGMA foreign_keys=ON")
     conn.executescript(SCHEMA)
+    # Added with attachments; a database from before them gains the column.
+    if "files" not in {r[1] for r in conn.execute("PRAGMA table_info(turns)")}:
+        conn.execute("ALTER TABLE turns ADD COLUMN files TEXT")
     return conn
 
 
-def save_turn(conn, conv, user, reply, tools, messages):
+def save_turn(conn, conv, user, reply, tools, messages, files=()):
     # Returns the conversation id: the row is made at its first finished turn,
-    # so a New that is never used leaves nothing behind.
+    # so a New that is never used leaves nothing behind. A question that is
+    # only files is titled by their names.
     now = int(time.time())
+    title = " ".join(user.split()[:6]) or ", ".join(f["name"] for f in files)
     with conn:
         if conv is None:
             conv = conn.execute("INSERT INTO conversations(title, created, updated) "
-                                "VALUES (?, ?, ?)",
-                                (" ".join(user.split()[:6]), now, now)).lastrowid
+                                "VALUES (?, ?, ?)", (title, now, now)).lastrowid
         else:
             conn.execute("UPDATE conversations SET updated = ? WHERE id = ?", (now, conv))
         seq = conn.execute("SELECT COALESCE(MAX(seq), 0) + 1 FROM turns WHERE conv = ?",
                            (conv,)).fetchone()[0]
-        conn.execute("INSERT INTO turns(conv, seq, user, reply, tools, messages) "
-                     "VALUES (?, ?, ?, ?, ?, ?)",
-                     (conv, seq, user, reply, json.dumps(tools), json.dumps(messages)))
+        conn.execute("INSERT INTO turns(conv, seq, user, reply, tools, messages, files) "
+                     "VALUES (?, ?, ?, ?, ?, ?, ?)",
+                     (conv, seq, user, reply, json.dumps(tools), json.dumps(messages),
+                      json.dumps(list(files))))
     return conv
 
 
@@ -176,12 +184,13 @@ def list_convs(conn):
 
 def load_conv(conn, conv):
     # (model history, shown turns), or None for an id that is not there.
-    rows = conn.execute("SELECT user, reply, tools, messages FROM turns WHERE conv = ? "
-                        "ORDER BY seq", (conv,)).fetchall()
+    rows = conn.execute("SELECT user, reply, tools, messages, files FROM turns "
+                        "WHERE conv = ? ORDER BY seq", (conv,)).fetchall()
     if not rows:
         return None
-    return ([json.loads(m) for *_, m in rows],
-            [{"user": u, "reply": r, "tools": json.loads(t)} for u, r, t, _ in rows])
+    return ([json.loads(m) for _, _, _, m, _ in rows],
+            [{"user": u, "reply": r, "tools": json.loads(t), "files": json.loads(f or "[]")}
+             for u, r, t, _, f in rows])
 
 
 def search(conn, q):
@@ -209,6 +218,134 @@ def delete_conv(conn, conv):
         conn.execute("DELETE FROM conversations WHERE id = ?", (conv,))
 
 
+# Attachments. Paths come from the canvas; the daemon runs as the user and
+# reads them itself. Text, and a PDF's text, go inline and are stored that
+# way, so a reopened conversation still has them. Images are stored as a
+# path and hash and read again when sent, so the history does not grow by
+# every picture, and a file changed since is reported gone rather than
+# shown as something it no longer is.
+CODE_SUFFIXES = {
+    ".py", ".c", ".h", ".cpp", ".hpp", ".rs", ".go", ".js", ".ts", ".jsx",
+    ".tsx", ".sh", ".bash", ".zsh", ".lua", ".rb", ".pl", ".php", ".java",
+    ".kt", ".swift", ".sql", ".toml", ".yaml", ".yml", ".json", ".ini",
+    ".cfg", ".conf", ".md", ".rst", ".txt", ".csv", ".xml", ".html", ".css",
+    ".diff", ".patch", ".slackbuild", ".info", ".desktop", ".service", ".qml",
+}
+IMAGE_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp"}
+MAX_FILES = 8
+MAX_IMAGE = 20 << 20
+MAX_PAGE_IMAGES = 4                       # scanned PDF pages sent as images
+
+
+class Refused(Exception):
+    """A file or request the user has to change; logged without a trace."""
+
+
+def kind_of(path):
+    mime = mimetypes.guess_type(path)[0] or ""
+    if mime in IMAGE_TYPES:
+        return "image"
+    if mime == "application/pdf":
+        return "pdf"
+    if mime.startswith("text/") or os.path.splitext(path)[1].lower() in CODE_SUFFIXES:
+        return "text"
+    return "other"
+
+
+def sha256(path):
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def pdf_pages(path):
+    out = subprocess.run(["pdftotext", "-layout", path, "-"], capture_output=True,
+                         text=True, timeout=60)
+    if out.returncode:
+        last = out.stderr.strip().splitlines()[-1:] or ["pdftotext failed"]
+        raise Refused(f"cannot read {os.path.basename(path)}: {last[0]}")
+    pages = out.stdout.split("\f")
+    return pages[:-1] if out.stdout.endswith("\f") else pages   # each page ends in \f
+
+
+def attach(paths, budget):
+    # (text blocks, image references, chips) for the files; Refused says
+    # what is wrong with the first bad one. Text files and PDFs share the
+    # budget of characters equally.
+    if len(paths) > MAX_FILES:
+        raise Refused(f"at most {MAX_FILES} files")
+    loaded = []
+    for p in paths:
+        name = os.path.basename(p) or p
+        if not (os.path.isabs(p) and os.path.isfile(p) and os.access(p, os.R_OK)):
+            raise Refused(f"cannot read {name}")
+        kind = kind_of(p)
+        if kind == "other":
+            raise Refused(f"{name}: chat supports text, images and PDF")
+        if kind == "image" and os.path.getsize(p) > MAX_IMAGE:
+            raise Refused(f"{name}: images over 20 MB are refused")
+        loaded.append((p, name, kind))
+    share = budget // max(1, sum(k != "image" for _, _, k in loaded))
+    blocks, refs, chips = [], [], []
+    for p, name, kind in loaded:
+        chip = {"name": name, "kind": kind, "truncated": False}
+        chips.append(chip)
+        if kind == "image":
+            refs.append({"type": "image_ref", "name": name, "path": p, "sha256": sha256(p)})
+            continue
+        if kind == "pdf":
+            pages = pdf_pages(p)
+            chip["pages"] = len(pages)
+            body = "\n\n".join(f"[page {i}]\n{t.rstrip()}" for i, t in enumerate(pages, 1))
+            # A page with no text is a scan: the model reads it as an image.
+            blank = [i for i, t in enumerate(pages, 1) if not t.strip()][:MAX_PAGE_IMAGES]
+            digest = sha256(p) if blank else ""
+            refs += [{"type": "image_ref", "name": f"{name} page {n}", "path": p,
+                      "sha256": digest, "page": n} for n in blank]
+        else:
+            with open(p, "rb") as f:          # no more than the share could use
+                body = f.read(share * 4 + 4).decode("utf-8", "replace")
+        if len(body) > share:
+            body, chip["truncated"] = body[:share], True
+        note = " (truncated)" if chip["truncated"] else ""
+        blocks.append(f"--- file: {name}{note} ---\n{body}\n--- end file ---")
+    return blocks, refs, chips
+
+
+def resolve(part):
+    # An image reference back into an image the model can see, or a note
+    # that it is gone: moved, deleted, or changed since it was sent.
+    if part.get("type") != "image_ref":
+        return part
+    gone = {"type": "text", "text": f"[image {part['name']} is no longer available]"}
+    try:
+        if sha256(part["path"]) != part["sha256"]:
+            return gone
+        if "page" in part:
+            n = str(part["page"])
+            out = subprocess.run(["pdftoppm", "-f", n, "-l", n, "-r", "150", "-png", part["path"]],
+                                 capture_output=True, timeout=60)
+            if out.returncode or not out.stdout:
+                return gone
+            data, mime = out.stdout, "image/png"
+        else:
+            with open(part["path"], "rb") as f:
+                data = f.read()
+            mime = mimetypes.guess_type(part["path"])[0]
+    except (OSError, subprocess.TimeoutExpired):
+        return gone
+    return {"type": "image_url",
+            "image_url": {"url": f"data:{mime};base64,{base64.b64encode(data).decode()}"}}
+
+
+def wire(messages):
+    # What is sent: the stored messages with their image references resolved.
+    return [{**m, "content": [resolve(p) for p in m["content"]]}
+            if isinstance(m.get("content"), list) else m for m in messages]
+
+
 # The canvas's feed: JSON lines on a Unix socket, both ways. Every client
 # gets every event. A client that stalls for 2 s is dropped rather than
 # holding up the others.
@@ -354,6 +491,9 @@ def handle(raw):
     # equal an unsaved conversation's None.
     bad = ("bad id" if kind in ("open", "delete") and type(cmd.get("id")) is not int
            else "bad text" if kind == "ask" and not isinstance(cmd.get("text"), str)
+           else "bad files" if kind == "ask" and not (
+               isinstance(cmd.get("files", []), list)
+               and all(isinstance(f, str) for f in cmd.get("files", [])))
            else "bad query" if kind == "search" and not isinstance(cmd.get("q"), str)
            else None)
     if bad:
@@ -372,10 +512,10 @@ def handle(raw):
 def command(kind, cmd):
     # Caller holds lock; arguments are already checked.
     if kind == "ask":
-        text = cmd["text"].strip()
-        if text:
+        text, files = cmd["text"].strip(), cmd.get("files", [])
+        if text or files:
             interrupt()
-            threading.Thread(target=reply, args=(turn, text), daemon=True).start()
+            threading.Thread(target=reply, args=(turn, text, files), daemon=True).start()
     elif kind == "stop":
         interrupt()
     elif kind == "new":
@@ -474,6 +614,16 @@ def llm_model():
                 models[0]["id"])
 
 
+def sees_images(model):
+    # The router reports each model's inputs; a server that does not is
+    # left to decide for itself.
+    with urllib.request.urlopen(CFG["llm"] + "/models", timeout=10) as r:
+        models = json.load(r)["data"]
+    m = next((m for m in models if m["id"] == model), {})
+    inputs = (m.get("architecture") or {}).get("input_modalities")
+    return inputs is None or "image" in inputs
+
+
 def ask(messages, calls, tools=None):
     # Yields ("think", text) and ("token", text) as they stream; tool-call
     # fragments go into calls.
@@ -540,11 +690,13 @@ def say(mine, kind, value, done=False):
         return True
 
 
-def reply(mine, text):
+def reply(mine, text, files=()):
     try:
-        answer(mine, text)
+        answer(mine, text, files)
     except Exception as e:
-        log("reply failed:\n" + traceback.format_exc().rstrip())
+        # A refused file is the user's to fix, not a bug: one line, no trace.
+        log(f"refused: {e}" if isinstance(e, Refused)
+            else "reply failed:\n" + traceback.format_exc().rstrip())
         with lock:
             if turn == mine:
                 cur["pending"] = None
@@ -554,16 +706,25 @@ def reply(mine, text):
                 emit({"type": "state", "state": "idle"})
 
 
-def answer(mine, text):
+def answer(mine, text, files=()):
+    blocks, refs, chips = [], [], []
+    if files:
+        # Half of what the context has free, at about 3 characters a token.
+        free = (ctx["max"] - ctx["used"]) // 2 * 3 if ctx["max"] else 24000
+        blocks, refs, chips = attach(files, max(free, 2000))
+        model = llm_model()
+        if refs and not sees_images(model):
+            raise Refused(f"{model} cannot see images")
     with lock:
         if turn != mine:
             return
         history = list(cur["history"])
-        cur["pending"] = {"user": text, "reply": "", "think": "", "tools": []}
-        emit({"type": "user", "text": text})
+        cur["pending"] = {"user": text, "reply": "", "think": "", "tools": [], "files": chips}
+        emit({"type": "user", "text": text, "files": chips})
         emit({"type": "state", "state": "thinking"})
-    log(f"you: {text}")
-    question = {"role": "user", "content": text}
+    log(f"you: {text}" + "".join(f" [{c['name']}]" for c in chips))
+    body = "\n\n".join(blocks + [text] if text else blocks)
+    question = {"role": "user", "content": [{"type": "text", "text": body}] + refs if refs else body}
     # The model has no clock: without the date it makes one up for tools.
     # Also stamped on the question itself, which history does not keep: with
     # the time only in the system prompt, "what time is it?" asked again an
@@ -573,7 +734,8 @@ def answer(mine, text):
     # ponytail: last 10 turns, summarise if long sessions need memory
     msgs = ([{"role": "system", "content": f"{CFG['system']} It is {now}."}]
             + [m for past in history[-10:] for m in past]
-            + [{"role": "user", "content": f"[{now}] {text}"}])
+            + [{"role": "user", "content": [{"type": "text", "text": f"[{now}] {body}"}] + refs
+                if refs else f"[{now}] {body}"}])
     first = len(msgs)                            # this turn's tool messages start here
     said, used = "", []
     # ponytail: a stopped reply notices at its next chunk, so llama-server's
@@ -581,7 +743,8 @@ def answer(mine, text):
     for rnd in range(CFG["tool_rounds"] + 1):
         calls, content = {}, ""
         # The last round offers no tools, which forces an answer.
-        for kind, piece in ask(msgs, calls, TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
+        for kind, piece in ask(wire(msgs), calls,
+                               TOOL_SCHEMA if rnd < CFG["tool_rounds"] else None):
             if not say(mine, kind, piece):
                 log("interrupted")
                 return
@@ -619,9 +782,9 @@ def answer(mine, text):
         if turn != mine:
             log("interrupted")
             return
-        cur["id"] = save_turn(DB, cur["id"], text, said, used, entry)
+        cur["id"] = save_turn(DB, cur["id"], text, said, used, entry, chips)
         cur["history"].append(entry)
-        cur["turns"].append({"user": text, "reply": said, "tools": used})
+        cur["turns"].append({"user": text, "reply": said, "tools": used, "files": chips})
         cur["pending"] = None
         emit({"type": "done", "id": cur["id"]})
         emit({"type": "state", "state": "idle"})
@@ -716,8 +879,8 @@ def test():
     hist, shown = load_conv(db, a)
     assert hist == [[{"role": "user", "content": "q"}], [{"role": "user", "content": "q2"}]]
     assert shown == [{"user": "che tempo fa in città oggi pomeriggio?", "reply": "Piove.",
-                      "tools": ["weather"]},
-                     {"user": "e domani?", "reply": "Sole.", "tools": []}]
+                      "tools": ["weather"], "files": []},
+                     {"user": "e domani?", "reply": "Sole.", "tools": [], "files": []}]
     assert load_conv(db, 999) is None
     hits = search(db, "citta")                       # no accent typed, accented word found
     assert [h["id"] for h in hits] == [a] and "**città**" in hits[0]["snippet"]
@@ -740,16 +903,103 @@ def test():
         assert os.stat(os.path.dirname(path)).st_mode & 0o777 == 0o700
         assert os.stat(path).st_mode & 0o777 == 0o600
 
+    # Attachments.
+    import struct
+    import zlib
+
+    def tiny_pdf(path, pages):
+        # One Helvetica line per page; "" is a page with no text, a scan.
+        objs = ["<< /Type /Catalog /Pages 2 0 R >>",
+                "<< /Type /Pages /Kids [%s] /Count %d >>"
+                % (" ".join(f"{4 + 2 * i} 0 R" for i in range(len(pages))), len(pages)),
+                "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"]
+        for i, t in enumerate(pages):
+            stream = f"BT /F1 24 Tf 72 700 Td ({t}) Tj ET" if t else ""
+            objs += [f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources "
+                     f"<< /Font << /F1 3 0 R >> >> /Contents {5 + 2 * i} 0 R >>",
+                     f"<< /Length {len(stream)} >>\nstream\n{stream}\nendstream"]
+        out, offsets = b"%PDF-1.4\n", []
+        for k, o in enumerate(objs, 1):
+            offsets.append(len(out))
+            out += f"{k} 0 obj\n{o}\nendobj\n".encode()
+        xref = len(out)
+        out += (f"xref\n0 {len(objs) + 1}\n0000000000 65535 f \n".encode()
+                + b"".join(f"{o:010d} 00000 n \n".encode() for o in offsets)
+                + f"trailer\n<< /Size {len(objs) + 1} /Root 1 0 R >>\n"
+                  f"startxref\n{xref}\n%%EOF\n".encode())
+        with open(path, "wb") as f:
+            f.write(out)
+
+    def tiny_png(path, rgb=b"\xff\x00\x00"):
+        def chunk(t, d):
+            return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d))
+        with open(path, "wb") as f:
+            f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0))
+                    + chunk(b"IDAT", zlib.compress(b"\x00" + rgb)) + chunk(b"IEND", b""))
+
+    files_dir = tempfile.TemporaryDirectory()
+    fd = files_dir.name
+    code, png, pdf, other = (os.path.join(fd, n) for n in ("a.py", "red.png", "doc.pdf", "x.bin"))
+    with open(code, "w") as f:
+        f.write("print('hi')\n" * 50)
+    tiny_png(png)
+    tiny_pdf(pdf, ["hello pdf", ""])
+    with open(other, "wb") as f:
+        f.write(b"\x00\x01")
+    assert [kind_of(p) for p in (code, png, pdf, other, "/x/README.md", "/x/b.SlackBuild")] == \
+        ["text", "image", "pdf", "other", "text", "text"]
+    blocks, refs, chips = attach([code, pdf, png], 100000)
+    assert blocks[0].startswith("--- file: a.py ---\nprint('hi')") and "(truncated)" not in blocks[0]
+    assert "[page 1]\nhello pdf" in blocks[1] and "[page 2]" in blocks[1]
+    assert chips == [{"name": "a.py", "kind": "text", "truncated": False},
+                     {"name": "doc.pdf", "kind": "pdf", "truncated": False, "pages": 2},
+                     {"name": "red.png", "kind": "image", "truncated": False}]
+    # The blank page is a scan: it goes as an image, rendered when sent.
+    assert [r["name"] for r in refs] == ["doc.pdf page 2", "red.png"] and refs[0]["page"] == 2
+    sent = [resolve(r) for r in refs]
+    assert all(x["type"] == "image_url" for x in sent)
+    assert sent[1]["image_url"]["url"].startswith("data:image/png;base64,")
+    assert base64.b64decode(sent[0]["image_url"]["url"].split(",", 1)[1])[:4] == b"\x89PNG"
+    small = attach([code], 24)[0][0]                 # the share is the cut: two lines
+    assert "(truncated)" in small and small.count("print") == 2
+    for bad, why in ((["a.py"], "cannot read a.py"), ([fd + "/nope.txt"], "cannot read nope.txt"),
+                     ([fd], "cannot read"), ([other], "x.bin: chat supports text, images and PDF"),
+                     ([code] * 9, "at most 8 files")):
+        try:
+            attach(bad, 1000)
+            assert False, bad
+        except Refused as e:
+            assert why in str(e), (bad, e)
+    # Stored as a reference: survives a save and reload, gone once changed.
+    a = save_turn(db, None, "", "A red pixel.", [], [{"role": "user", "content": refs[1:]}],
+                  chips[2:])
+    hist, shown = load_conv(db, a)
+    assert shown[0]["files"] == chips[2:] and list_convs(db)[0]["title"] == "red.png"
+    assert wire(hist[0])[0]["content"][0]["type"] == "image_url"
+    tiny_png(png, b"\x00\xff\x00")
+    assert wire(hist[0])[0]["content"][0] == {
+        "type": "text", "text": "[image red.png is no longer available]"}
+    os.remove(png)
+    assert resolve(refs[1])["type"] == "text"        # deleted reads the same
+    # An old database gains the files column.
+    with tempfile.TemporaryDirectory() as d:
+        old = sqlite3.connect(os.path.join(d, "h.db"))
+        old.executescript(SCHEMA)
+        old.close()
+        assert "files" in {r[1] for r in db_open(os.path.join(d, "h.db"))
+                           .execute("PRAGMA table_info(turns)")}
+
     # Protocol, against a fake llama-server.
     global DB, TOOLS, TOOL_SCHEMA
-    script = {"replies": [], "delay": 0}
+    script = {"replies": [], "delay": 0, "inputs": ["text", "image"]}
 
     class Fake(http.server.BaseHTTPRequestHandler):
         def log_message(self, *a):
             pass
 
         def do_GET(self):
-            body = ({"data": [{"id": "m", "status": {"value": "loaded"}}]}
+            body = ({"data": [{"id": "m", "status": {"value": "loaded"},
+                               "architecture": {"input_modalities": script["inputs"]}}]}
                     if self.path.endswith("/models")
                     else {"default_generation_settings": {"n_ctx": 4096}})
             self.send_response(200)
@@ -818,7 +1068,7 @@ def test():
         assert script["body"]["chat_template_kwargs"]["enable_thinking"] is True
         conv = cur["id"]
         assert load_conv(DB, conv)[1] == [{"user": "what time is it", "reply": "It is noon.",
-                                           "tools": ["now"]}]
+                                           "tools": ["now"], "files": []}]
         # The tool result went to the model but history keeps a placeholder.
         assert load_conv(DB, conv)[0][0][2] == {"role": "tool", "tool_call_id": "c1",
                                                 "content": "(shown to the user at the time)"}
@@ -888,6 +1138,32 @@ def test():
         seen = until("conversations")
         assert {"type": "conversation", "id": None, "turns": []} in seen   # was current
         assert seen[-1]["items"] == []
+        # Files: the text inline, the image resolved, chips on the turn.
+        tiny_png(png)
+        script["replies"] = [[{"content": "Red."}]]
+        send(type="ask", text="what is it?", files=[pdf, png])
+        user = until("user")[-1]
+        assert user["files"][1] == {"name": "red.png", "kind": "image", "truncated": False}
+        until("done")
+        q = script["body"]["messages"][-1]["content"]
+        assert "hello pdf" in q[0]["text"] and q[0]["text"].endswith("what is it?")
+        assert [x["type"] for x in q] == ["text", "image_url", "image_url"]
+        assert cur["turns"][-1]["files"] == user["files"]
+        # Only files, no text.
+        script["replies"] = [[{"content": "Code."}]]
+        send(type="ask", text="", files=[code])
+        assert until("user")[-1]["text"] == ""
+        until("done")
+        # A model without vision is refused before anything is sent.
+        script["inputs"] = ["text"]
+        send(type="ask", text="see?", files=[png])
+        e = until("error")[-1]
+        assert e["message"] == "m cannot see images" and e["turn"] is True
+        script["inputs"] = ["text", "image"]
+        send(type="ask", text="x", files=[other])
+        assert "chat supports" in until("error")[-1]["message"]
+        send(type="ask", text="x", files="not a list")
+        assert until("error")[-1]["message"] == "bad files"
         # llama-server down: an error, not a dead thread.
         CFG["llm"] = "http://127.0.0.1:9/v1"
         send(type="ask", text="anyone?")