diff options
Diffstat (limited to 'assistant.py')
| -rw-r--r-- | assistant.py | 28 |
1 files changed, 24 insertions, 4 deletions
diff --git a/assistant.py b/assistant.py index 937700a..db527a5 100644 --- a/assistant.py +++ b/assistant.py @@ -488,11 +488,20 @@ def ask(messages, calls, tools=None): body["temperature"] = CFG["temperature"] req = urllib.request.Request(CFG["llm"] + "/chat/completions", json.dumps(body).encode(), {"Content-Type": "application/json"}) - with urllib.request.urlopen(req, timeout=300) as r: + try: + r = urllib.request.urlopen(req, timeout=300) + except urllib.error.HTTPError as e: + # The body says why (a full context, a bad request); the status alone does not. + raise RuntimeError(f"llama-server {e.code}: " + f"{e.read().decode(errors='replace')[:300]}") from None + with r: for line in r: if not line.startswith(b"data: ") or line.strip() == b"data: [DONE]": continue chunk = json.loads(line[6:]) + if "error" in chunk: # a failure mid-stream arrives as a chunk + err = chunk["error"] + raise RuntimeError(err.get("message", str(err)) if isinstance(err, dict) else str(err)) if chunk.get("usage"): # How full the model's context is: history, tool results and # tool definitions all count. @@ -626,7 +635,7 @@ def main(): DB = db_open(DB_PATH) try: context_size(llm_model()) # so the first hello has it - except OSError: + except Exception: pass # llama-server down: asked again per request serve() served = True @@ -635,6 +644,8 @@ def main(): signal.pause() except (KeyboardInterrupt, SystemExit) as e: log(f"exit: {type(e).__name__} {e}".rstrip()) + if not served: + raise SystemExit(1) # refused or failed before serving: not a clean stop except BaseException: log("crashed:\n" + traceback.format_exc().rstrip()) raise @@ -744,12 +755,13 @@ def test(): self.wfile.write(json.dumps(body).encode()) def do_POST(self): - self.rfile.read(int(self.headers["Content-Length"])) + script["body"] = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) self.send_response(200) self.end_headers() try: for delta in script["replies"].pop(0): - self.wfile.write(b"data: " + json.dumps({"choices": [{"delta": delta}]}) + self.wfile.write(delta.encode() if isinstance(delta, str) else + b"data: " + json.dumps({"choices": [{"delta": delta}]}) .encode() + b"\n\n") self.wfile.flush() time.sleep(script["delay"]) @@ -801,12 +813,20 @@ def test(): ("state", "thinking"), ("token", "It is "), ("token", "noon."), ("done", None)], got assert until("state")[-1]["state"] == "idle" + 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"]}] # 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)"} + # An error chunk mid-stream reaches the client with its message. + script["replies"] = [['data: {"error": {"message": "context full"}}\n\n']] + CFG["thinking"] = False + send(type="ask", text="too much") + assert until("error")[-1]["message"] == "context full" + assert script["body"]["chat_template_kwargs"]["enable_thinking"] is False + CFG["thinking"] = True # Stop mid-stream: no token after done, nothing saved. script["replies"] = [[{"content": f"w{i} "} for i in range(50)]] script["delay"] = 0.02 |