diff options
Diffstat (limited to 'assistant.py')
| -rw-r--r-- | assistant.py | 125 |
1 files changed, 122 insertions, 3 deletions
diff --git a/assistant.py b/assistant.py index 8cdf232..4107092 100644 --- a/assistant.py +++ b/assistant.py @@ -406,6 +406,7 @@ ctx = {"model": None, "max": 0, "used": 0} lock = threading.RLock() cur = {"id": None, "history": [], "turns": [], "pending": None} inflight = {"conn": None} # the HTTP connection the current reply streams over +writing = {"on": False} # a write in flight, which interrupt() also ends turn = 0 # bumped by every ask, stop, new and open: older replies are stale DB = None @@ -519,6 +520,9 @@ def interrupt(): # prompt pass would otherwise keep the GPU busy until that chunk. global turn turn += 1 + if writing["on"]: # a write has no pending turn to close + writing["on"] = False + emit({"type": "state", "state": "idle"}) conn = inflight["conn"] if conn is not None and conn.sock is not None: try: @@ -553,6 +557,13 @@ def handle(raw): 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 "bad write" if kind == "write" and not ( + isinstance(cmd.get("files"), list) and cmd["files"] + and all(isinstance(f, str) for f in cmd["files"]) + and isinstance(cmd.get("prompt"), str) + and isinstance(cmd.get("dir"), str) and os.path.isabs(cmd["dir"]) + and isinstance(cmd.get("name"), str) and cmd["name"].strip(". ") + and "/" not in cmd["name"]) else None) if bad: emit({"type": "error", "message": bad}) @@ -569,7 +580,11 @@ def handle(raw): def command(kind, cmd): # Caller holds lock; arguments are already checked. - if kind == "ask": + if kind == "write": + interrupt() + threading.Thread(target=write, daemon=True, + args=(turn, cmd["files"], cmd["prompt"], cmd["dir"], cmd["name"])).start() + elif kind == "ask": text, files = cmd["text"].strip(), cmd.get("files", []) if text or files: interrupt() @@ -678,7 +693,7 @@ def llm_model(): return model_inputs()[0] -def ask(messages, calls, tools=None, mine=None): +def ask(messages, calls, tools=None, mine=None, think=None): # Yields ("think", text) and ("token", text) as they stream; tool-call # fragments go into calls. mine is the reply's turn: a stop shuts the # connection from the other thread (interrupt), which ends the reading. @@ -686,7 +701,7 @@ def ask(messages, calls, tools=None, mine=None): context_size(model) body = {"model": model, "messages": messages, "stream": True, "stream_options": {"include_usage": True}, - "chat_template_kwargs": {"enable_thinking": CFG["thinking"]}} + "chat_template_kwargs": {"enable_thinking": CFG["thinking"] if think is None else think}} if tools: body["tools"] = tools if CFG["temperature"] is not None: @@ -773,6 +788,80 @@ def say(mine, kind, value, done=False): return True +def free_name(d, name): + # The first of name, "stem (1).ext", "stem (2).ext" not taken in d. + stem, ext = os.path.splitext(name) + path, n = os.path.join(d, name), 1 + while os.path.lexists(path): + path, n = os.path.join(d, f"{stem} ({n}){ext}"), n + 1 + return path + + +def unfence(text): + # Models often wrap a whole document in one code fence; the file wants + # what is inside it. + lines = text.strip().split("\n") + if len(lines) >= 2 and lines[0].startswith("```") and lines[-1].strip() == "```": + return "\n".join(lines[1:-1]) + "\n" + return text + + +def write(mine, files, prompt, d, name): + # A one-off request whose reply goes to a file, for the canvas's To .md: + # no history in, nothing kept in the conversation, no thinking, no tools. + tmp = None + try: + free = (ctx["max"] - ctx["used"]) // 2 * 3 if ctx["max"] else 24000 + blocks, refs, _ = attach(files, max(free, 2000)) + model, inputs = model_inputs() + images = inputs is None or "image" in inputs + if refs and not images: + raise Refused(f"{model} cannot see images") + with lock: + if turn != mine: + return + writing["on"] = True + emit({"type": "state", "state": "thinking"}) + body = "\n\n".join(blocks + [prompt]) + msgs = [{"role": "user", "content": [{"type": "text", "text": body}] + refs if refs else body}] + text, shown = "", 0.0 + for kind, piece in ask(wire(msgs, images), {}, None, mine, think=False): + if turn != mine: + return + if kind == "token": + text += piece + if time.time() - shown >= 1: + shown = time.time() + emit({"type": "progress", "chars": len(text)}) + if turn != mine: # stopped: the hang-up ended the stream + return + os.makedirs(d, exist_ok=True) + tmp = os.path.join(d, f".write-{os.getpid()}-{mine}.tmp") + with open(tmp, "w") as f: + f.write(unfence(text)) + with lock: + if turn != mine: + return + path = free_name(d, name) + os.rename(tmp, path) + tmp = None + writing["on"] = False + emit({"type": "written", "path": path}) + emit({"type": "state", "state": "idle"}) + log(f"wrote {path}") + except Exception as e: + log(f"refused: {e}" if isinstance(e, Refused) + else "write failed:\n" + traceback.format_exc().rstrip()) + with lock: + if turn == mine: + writing["on"] = False + emit({"type": "error", "message": str(e) or type(e).__name__, "turn": True}) + emit({"type": "state", "state": "idle"}) + finally: + if tmp is not None and os.path.exists(tmp): + os.remove(tmp) + + def reply(mine, text, files=()): try: answer(mine, text, files) @@ -1329,6 +1418,36 @@ def test(): assert "chat supports" in until("error")[-1]["message"] send(type="ask", text="x", files="not a list") assert until("error")[-1]["message"] == "bad files" + # write: the reply goes to a file, not the conversation. + before = len(list_convs(DB)) + with open(os.path.join(fd, "doc.md"), "w") as fh: # f is the socket's + fh.write("taken") + script["replies"] = [['data: {"choices": [{"delta": {"content": "```markdown\\n# Title\\n\\nbody\\n"}}]}\n\n', + {"content": "```"}]] + send(type="write", files=[pdf], prompt="As Markdown.", dir=fd, name="doc.md") + w = until("written")[-1] + assert w["path"] == os.path.join(fd, "doc (1).md"), w + with open(w["path"]) as fh: + assert fh.read() == "# Title\n\nbody\n" # the one fence unwrapped + assert script["body"]["chat_template_kwargs"]["enable_thinking"] is False + assert "tools" not in script["body"] and script["body"]["messages"][-1]["role"] == "user" + assert len(list_convs(DB)) == before # nothing in History + assert until("state")[-1]["state"] == "idle" + # A stopped write leaves nothing, and the state goes back to idle. + script.update(hold=True, hung_up=None) + script["replies"] = [[]] + send(type="write", files=[pdf], prompt="As Markdown.", dir=fd, name="held.md") + until("state") + send(type="stop") + assert until("state")[-1]["state"] == "idle" + time.sleep(0.3) + assert not os.path.exists(os.path.join(fd, "held.md")) + assert not [n for n in os.listdir(fd) if n.startswith(".write-")] + script["hold"] = False + for bad in ({"dir": "relative", "name": "x.md"}, {"dir": fd, "name": "a/b.md"}, + {"dir": fd, "name": ""}): + send(type="write", files=[pdf], prompt="p", **bad) + assert until("error")[-1]["message"] == "bad write", bad # llama-server down: an error, not a dead thread. CFG["llm"] = "http://127.0.0.1:9/v1" send(type="ask", text="anyone?") |