288 lines
11 KiB
Python
288 lines
11 KiB
Python
"""Live progress, written to a file for the Open WebUI filter to collect.
|
|
|
|
WHY A FILE AND NOT A POST TO OPEN WEBUI
|
|
|
|
The first version of this posted straight to
|
|
`/api/v1/chats/{id}/messages/{id}/event`, which works and needs an Open
|
|
WebUI **admin** API key - because anyone may invoke vm_agent, the chat
|
|
belongs to whoever is talking, and that endpoint refuses a key whose user
|
|
neither owns the chat nor is an admin. An admin key living in a container
|
|
is an agent that could read every conversation on the instance. It was
|
|
kept out of the agent itself (nginx held it), but it still existed.
|
|
|
|
So it does not exist any more. The agent writes progress to the job
|
|
volume it already shares with the api container, and a filter INSIDE Open
|
|
WebUI reads it and emits the events itself - where `__event_emitter__` is
|
|
already bound to the requesting user's own session and no credential is
|
|
involved at any point.
|
|
|
|
The trade is that progress is polled rather than pushed, and the cost of
|
|
that is a 2-second poll inside Open WebUI's event loop for the length of
|
|
a session. That is the whole of it.
|
|
|
|
WHAT REACHES THE MODEL: NOTHING
|
|
|
|
Open WebUI stores a message's status history and its `embeds` as SIBLING
|
|
fields, and rebuilds the payload sent to the model from an allowlist -
|
|
role, thinking, content, tool_calls, tool_call_id, images. Neither is on
|
|
it. So a screenshot posted this way is seen by the person watching and
|
|
never by the model that called the tool. That is why this exists at all
|
|
rather than returning images from the tool: a tool result that IS a data:
|
|
URI goes to the MODEL as a vision input, which is the opposite of what is
|
|
wanted here.
|
|
|
|
FAILURE IS ALWAYS SILENT
|
|
|
|
A progress update that cannot be written must never break a session. The
|
|
user loses the live view of a job that is otherwise running perfectly,
|
|
which is a bad trade to make the other way round.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
from collections import deque
|
|
|
|
log = logging.getLogger("vmagent.events")
|
|
|
|
JOBS = os.environ.get("OWUITOOLS_VM_JOBS_DIR", "/vmjobs")
|
|
PROGRESS_DIR = os.path.join(JOBS, "progress")
|
|
|
|
# Chat ids are Open WebUI UUIDs and become a filename here, so they are
|
|
# checked rather than trusted - the same discipline as everywhere else
|
|
# in this suite that turns a name into a path.
|
|
CHAT_ID_RE = re.compile(r"^[0-9a-fA-F-]{36}$")
|
|
|
|
# Twenty lines, fixed. Enough to see a command and what it said, short
|
|
# enough that the view does not become the conversation - and a constant,
|
|
# so the box does not jump about as the session runs.
|
|
TERMINAL_LINES = 20
|
|
# The live line above the task list, and how often it may be written.
|
|
# Short enough that the tail of it is what a one-line box actually shows:
|
|
# the text is already only the newest characters, and anything past the
|
|
# width of the box is clipped, so a longer cap only lags further behind.
|
|
SPEECH_CHARS = 160
|
|
SPEECH_SECONDS = 1.0
|
|
|
|
# Every screenshot write is ~80 KB on a 16 MB tmpfs and every read is a
|
|
# poll from Open WebUI. Once every few seconds is plenty for watching a
|
|
# window open; the agent takes them far more often than that.
|
|
MIN_SCREENSHOT_INTERVAL = 3.0
|
|
|
|
JPEG_QUALITY = 55
|
|
MAX_SHOT_BYTES = 400_000
|
|
|
|
|
|
class Events:
|
|
def __init__(self, chat_id: str | None, message_id: str | None,
|
|
interface: str = "cli") -> None:
|
|
self.chat_id = (chat_id or "").strip()
|
|
self.enabled = bool(CHAT_ID_RE.fullmatch(self.chat_id))
|
|
self.path = os.path.join(PROGRESS_DIR, f"{self.chat_id}.json") if self.enabled else ""
|
|
self._seq = 0
|
|
self._shot_seq = 0
|
|
self._shot_b64 = ""
|
|
self._shot_caption = ""
|
|
self._last_shot = 0.0
|
|
self._status = ""
|
|
self._done = False
|
|
# THE TWO VIEWS, AND WHY A SESSION ONLY EVER GETS ONE.
|
|
#
|
|
# A CLI session has nothing to photograph, and a GUI session's
|
|
# screen already shows whatever terminal it opened. Sending both
|
|
# would put a picture of a desktop next to a transcript of
|
|
# commands that are not on it.
|
|
# Kept only so the progress file still says what kind of session
|
|
# this was; there is one kind now, and nothing branches on it.
|
|
self.interface = interface if interface in ("cli", "gui") else "gui"
|
|
# ONE HEIGHT. There used to be two, because a desktop session
|
|
# showed a screen above the feed and the feed was halved to make
|
|
# room. The screen now appears only once the agent has actually
|
|
# taken a picture, so the feed is not competing with something
|
|
# that may never arrive.
|
|
self._terminal: deque[str] = deque(maxlen=TERMINAL_LINES)
|
|
self._terminal_seq = 0
|
|
self._tasks = ""
|
|
# What the model is saying, right now, in one line. See speech().
|
|
self._speech = ""
|
|
self._speech_at = 0.0
|
|
if not self.enabled:
|
|
log.info("live progress disabled: no usable chat id was forwarded")
|
|
return
|
|
try:
|
|
os.makedirs(PROGRESS_DIR, mode=0o700, exist_ok=True)
|
|
except OSError as exc:
|
|
log.warning("cannot create %s: %s", PROGRESS_DIR, exc)
|
|
self.enabled = False
|
|
|
|
# -- writing -------------------------------------------------------
|
|
def _flush(self) -> None:
|
|
if not self.enabled:
|
|
return
|
|
self._seq += 1
|
|
payload = {
|
|
"seq": self._seq,
|
|
"status": self._status,
|
|
"tasks": self._tasks,
|
|
"speech": self._speech,
|
|
"interface": self.interface,
|
|
"terminal": list(self._terminal),
|
|
"terminal_seq": self._terminal_seq,
|
|
"done": self._done,
|
|
"shot_seq": self._shot_seq,
|
|
"shot": self._shot_b64,
|
|
"caption": self._shot_caption,
|
|
"updated": round(time.time(), 2),
|
|
}
|
|
try:
|
|
# Written whole and renamed into place: the reader is a
|
|
# separate process polling this path, and a half-written JSON
|
|
# file is the one thing that would make it throw rather than
|
|
# simply see nothing new.
|
|
tmp = self.path + ".tmp"
|
|
with open(tmp, "w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle)
|
|
os.replace(tmp, self.path)
|
|
except OSError as exc:
|
|
log.warning("could not write progress: %s", exc)
|
|
|
|
# -- text ----------------------------------------------------------
|
|
def status(self, description: str, done: bool = False, force: bool = False) -> None:
|
|
if not force and description == self._status and done == self._done:
|
|
return
|
|
self._status = description
|
|
self._done = done
|
|
self._flush()
|
|
|
|
def tasks(self, rendered: str) -> None:
|
|
"""The list the user watches the session against.
|
|
|
|
Above everything else in the view on purpose: a status line says
|
|
what is happening this second, and the list says whether that
|
|
amounts to progress.
|
|
"""
|
|
rendered = (rendered or "").strip()
|
|
if rendered and rendered != self._tasks:
|
|
self._tasks = rendered[:2000]
|
|
self._flush()
|
|
|
|
def speech(self, text: str) -> None:
|
|
"""The one line above the task list: what the model is saying now.
|
|
|
|
Rewritten rather than appended, and flushed at most once a second.
|
|
The tokens arrive far faster than anybody reads and faster than
|
|
the watcher polls, so writing the file per token would be work
|
|
nobody sees. The tail is kept rather than the head - the end of
|
|
the sentence is the part that is still moving.
|
|
"""
|
|
if not self.enabled:
|
|
return
|
|
line = " ".join(str(text or "").split())[-SPEECH_CHARS:]
|
|
now = time.time()
|
|
if line == self._speech or now - self._speech_at < SPEECH_SECONDS:
|
|
return
|
|
self._speech = line
|
|
self._speech_at = now
|
|
self._flush()
|
|
|
|
def terminal(self, *lines: str) -> None:
|
|
"""Append to the rolling view of what the agent is doing.
|
|
|
|
Capped at the deque's own maxlen, which is a `deque` rather
|
|
than a trim-on-read so the cost is bounded at the point where the
|
|
lines arrive - a command that prints a megabyte does not get to
|
|
put a megabyte through this file on its way to being cut down.
|
|
"""
|
|
if not self.enabled:
|
|
return
|
|
added = False
|
|
for line in lines:
|
|
for part in str(line).replace("\r", "").split("\n"):
|
|
self._terminal.append(part[:200])
|
|
added = True
|
|
if added:
|
|
self._terminal_seq += 1
|
|
self._flush()
|
|
|
|
def heartbeat(self) -> None:
|
|
"""Re-stamp the file so the watcher can tell alive from wedged.
|
|
|
|
A thinking model on a long context is quiet for minutes at a
|
|
time, and "no new status" and "the agent has died" look identical
|
|
without this.
|
|
"""
|
|
self._flush()
|
|
|
|
# -- pictures ------------------------------------------------------
|
|
def screenshot(self, png_bytes: bytes, caption: str = "") -> None:
|
|
if not self.enabled:
|
|
return
|
|
now = time.monotonic()
|
|
if now - self._last_shot < MIN_SCREENSHOT_INTERVAL:
|
|
return
|
|
self._last_shot = now
|
|
encoded = _to_jpeg_b64(png_bytes)
|
|
if not encoded:
|
|
return
|
|
self._shot_b64 = encoded
|
|
self._shot_caption = (caption or "")[:200]
|
|
self._shot_seq += 1
|
|
self._flush()
|
|
|
|
def finish(self) -> None:
|
|
"""Mark the session over and drop the picture.
|
|
|
|
The status line stays - the watcher shows it as the final state -
|
|
but the screenshot goes, so a stale frame of a machine that no
|
|
longer exists is not left on screen.
|
|
"""
|
|
self._shot_b64 = ""
|
|
self._shot_seq += 1
|
|
# The terminal stays. It is a record of what was done, and the
|
|
# last thing on it is usually the reason the session ended.
|
|
self._done = True
|
|
self._flush()
|
|
|
|
def clear_embeds(self) -> None:
|
|
self.finish()
|
|
|
|
|
|
def _to_jpeg_b64(png_bytes: bytes) -> str:
|
|
"""PNG in, JPEG out, because size is the constraint here.
|
|
|
|
Every write lands on a 16 MB tmpfs and every read crosses into Open
|
|
WebUI. A desktop screenshot as PNG is a few hundred KB before base64
|
|
adds a third; JPEG at modest quality is an order of magnitude smaller
|
|
and entirely good enough for a human watching a window open.
|
|
|
|
Pillow is optional: without it a small enough PNG is passed through
|
|
and a large one is skipped. A missing dependency should cost the live
|
|
view, not the session.
|
|
"""
|
|
try:
|
|
from PIL import Image # noqa: PLC0415
|
|
except ImportError:
|
|
if len(png_bytes) <= MAX_SHOT_BYTES:
|
|
return "data:image/png;base64," + base64.b64encode(png_bytes).decode()
|
|
return ""
|
|
try:
|
|
image = Image.open(io.BytesIO(png_bytes)).convert("RGB")
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True)
|
|
data = buffer.getvalue()
|
|
if len(data) > MAX_SHOT_BYTES:
|
|
image.thumbnail((960, 960))
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format="JPEG", quality=45, optimize=True)
|
|
data = buffer.getvalue()
|
|
return "data:image/jpeg;base64," + base64.b64encode(data).decode()
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("could not encode a screenshot: %s", exc)
|
|
return ""
|