Initial Commit
This commit is contained in:
@@ -0,0 +1,287 @@
|
||||
"""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 ""
|
||||
Reference in New Issue
Block a user