134 lines
5.3 KiB
Python
134 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Streaming exists for the person watching, and may not change the turn.
|
|
|
|
python3 tests/test_vm_agent_stream.py
|
|
|
|
WHY
|
|
|
|
The agent's calls were non-streaming on purpose: it acts on a finished
|
|
message, so nothing in the loop needed tokens as they arrived. That was
|
|
true of the agent and false of the operator, who asked to see what the
|
|
model is saying while it says it - a turn on this model runs tens of
|
|
seconds, and a motionless box makes "still thinking" look like "wedged".
|
|
|
|
So the reply is now read as it arrives. The rule that makes that safe is
|
|
that it is a DISPLAY change and nothing else: the assembled message handed
|
|
back must be identical in shape and content to the one the non-streaming
|
|
call returned, because the agent is not allowed to care which path
|
|
produced it. These cases are all about that equivalence, plus the one
|
|
about the callback - which belongs to the watcher, so it must never be
|
|
able to fail a turn.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, os.path.join(HERE, "..", "vmagent", "app"))
|
|
|
|
FAILS: list[str] = []
|
|
|
|
|
|
def check(label: str, condition: bool, detail: str = "") -> None:
|
|
if condition:
|
|
print(f"ok {label}")
|
|
else:
|
|
FAILS.append(label)
|
|
print(f"FAIL {label}\n {detail}")
|
|
|
|
|
|
import ollama # noqa: E402
|
|
|
|
CALL = {"function": {"name": "run", "arguments": {"command": "uname -a"}}}
|
|
|
|
|
|
def chunks(*entries) -> list:
|
|
return [json.dumps(e).encode() for e in entries]
|
|
|
|
|
|
# --- the assembled message ----------------------------------------------
|
|
seen: list[str] = []
|
|
body = ollama._read_stream(chunks(
|
|
{"message": {"content": "Check"}},
|
|
{"message": {"content": "ing the"}},
|
|
{"message": {"content": " kernel."}},
|
|
{"message": {"tool_calls": [CALL]}},
|
|
{"done": True, "done_reason": "stop", "prompt_eval_count": 120, "eval_count": 44,
|
|
"message": {}},
|
|
), seen.append)
|
|
|
|
check("the text is joined in order", body["message"]["content"] == "Checking the kernel.",
|
|
body["message"]["content"])
|
|
check("tool calls survive the stream", body["message"]["tool_calls"] == [CALL])
|
|
check("the counts come off the final chunk",
|
|
body.get("prompt_eval_count") == 120 and body.get("eval_count") == 44)
|
|
check("and so does why it stopped", body.get("done_reason") == "stop")
|
|
check("the role is set, as the non-streaming reply sets it",
|
|
body["message"]["role"] == "assistant")
|
|
|
|
# The watcher is handed the text SO FAR, not the delta - it paints a line,
|
|
# it does not append to one.
|
|
check("the callback gets the running text, not the fragment",
|
|
seen == ["Check", "Checking the", "Checking the kernel."], seen)
|
|
|
|
# --- thinking ------------------------------------------------------------
|
|
thought: list[str] = []
|
|
body = ollama._read_stream(chunks(
|
|
{"message": {"thinking": "weighing it up"}},
|
|
{"message": {"content": "Right, "}},
|
|
{"message": {"thinking": " and more"}},
|
|
{"message": {"content": "doing it."}},
|
|
{"done": True, "message": {}},
|
|
), thought.append)
|
|
check("thinking is kept on the message", body["message"]["thinking"] == "weighing it up and more")
|
|
check("it is shown while there is nothing else to show", thought[0] == "weighing it up",
|
|
"a model that is thinking has not said anything yet, and a blank minute reads as hung")
|
|
check("but once it speaks, its words are what is shown",
|
|
thought[-1] == "Right, doing it." and " and more" not in thought[-1],
|
|
thought)
|
|
|
|
# --- nothing in a frame may cost the turn --------------------------------
|
|
body = ollama._read_stream(
|
|
[b'{"message": {"content": "a"}}', b"not json at all", b"",
|
|
b'{"message": {"content": "b"}, "done": true}'],
|
|
lambda _t: None)
|
|
check("a malformed frame is skipped, not raised on",
|
|
body["message"]["content"] == "ab",
|
|
"losing a whole turn over one bad line would be a poor trade")
|
|
|
|
|
|
def explode(_text):
|
|
raise RuntimeError("the watcher's callback failed")
|
|
|
|
|
|
try:
|
|
body = ollama._read_stream(chunks(
|
|
{"message": {"content": "still works"}}, {"done": True, "message": {}}), explode)
|
|
survived = body["message"]["content"] == "still works"
|
|
except Exception as exc: # noqa: BLE001
|
|
survived = False
|
|
why = f"{type(exc).__name__}: {exc}"
|
|
check("a callback that throws cannot fail the turn", survived, locals().get("why", ""))
|
|
|
|
# --- an empty reply is still a reply --------------------------------------
|
|
body = ollama._read_stream(chunks({"done": True, "message": {}}), seen.append)
|
|
check("a reply with nothing in it assembles rather than crashing",
|
|
body["message"]["content"] == "" and body["message"]["tool_calls"] == [])
|
|
|
|
# --- and the old path is untouched ----------------------------------------
|
|
src = open(os.path.join(HERE, "..", "vmagent", "app", "ollama.py"), encoding="utf-8").read()
|
|
check("streaming is off unless somebody is watching",
|
|
'"stream": on_text is not None' in src,
|
|
"no caller without a callback may be made to pay for this")
|
|
check("only the agent's own turn streams; the summariser does not",
|
|
open(os.path.join(HERE, "..", "vmagent", "app", "loop.py"),
|
|
encoding="utf-8").read().count("on_text=") == 1,
|
|
"nobody watches a compaction summary")
|
|
|
|
print()
|
|
if FAILS:
|
|
print(f"FAILED {len(FAILS)}")
|
|
raise SystemExit(1)
|
|
print("the watcher sees the model talking, and the turn cannot tell the difference")
|